Fix grad accum when using loss_mask

[Aphoh]

Implement the fix described here. I don't think this is the cleanest way to do it, but my idea was to make a mixin for batch elements that allow them to specify when they have a specific num_elements that the microbatch losses should be divided by.

[dlwh]

thanks for tackling this! I think we can push more of the logic into grad_accum instead and that will be a bit cleaner. WDYT?

[dlwh]

imho i think the thing to do is to use reduction=None inside grad accum and do the reduction separately. Also, can't you infer batch_num_elements from loss_mask.sum?

[Aphoh]

@dlwh We can't get that from loss_mask.sum since the loss_mask we get is for the microbatch and not the whole batch. Instead of computing a mean loss per micro batch an averaging the mean losses, we need to sum all the microbatches dividing by the true number of losses in the batch.

[dlwh]

let's have the grad_fn return both the total (masked) loss and the number of elements (and the gradient?). WDYT?

[Aphoh]

So we'd change the signature of ComputeLossFunction to return something like
(model, input) -> (loss, extras) where extras has the number of elements in the microbatch in it?

I also find myself wanting to log things this way (by returning elements in an extras struct) so this would be a nice way to add that in.

The only concern is that if the number of microbatch elements or the losses are particularly large we could lose precision in the value of the loss right?
When computing (loss1 + loss2 + ... ) / num_elems v.s. loss1/num_elems + loss2/num_elems + ..., I would kinda expect the second to be more accurate. Is this a reasonable concern or am I prematurely optimizing?

[Aphoh]

I think this might be a valid concern given we'd be summing up fp16 gradients? This is what happens if you sum up a bunch of values near ln(vocab_size) which is around the maximum value you'd get from the loss.

[dlwh]

there's actually a standard way of dealing with this which I've encapsulated as RunningMean. The basic idea is to maintain the mean and the count rather than the sum and the count. This solves most precision problems usually. (You can actually do variance too but we don't need it)

[Aphoh]

@dlwh I tried to encapsulate this in this last commit. I'm pretty sure it doesn't work right since it doesn't pass the grad_accum tests (trees don't match). Could you give it a quick peek and let me know if you see anything?

[Aphoh]

passing tests, I think this is ready for a peek @dlwh