Just to see the diff

[Muennighoff]

No description provided.

[loubnabnl]

Looks great! I just left a comment on the order of the losses and the batch size we can fit

[loubnabnl]

Comment about the batch size: we're assuming that we can fit a batch size of 320 with our workers, but I think we can only fit 12 sequences on A100 40GB (so on 16 workers: batch of 16*12=192).

So we should probably either incorporate gradient accumulation and store the losses for 2 iterations (2 * 10 (small bz) * 16 gpus=320) or we can change the batch sizes from 320/32 to something that suits us with a 10% ratio like 160/16. In the paper they just talk about the 10% ratio but I'm not sure if using large batches si also important?

[Muennighoff]

There are no gradients here which means that a) we can likely fit a bigger batch size than 12 b) instead of grad acc. we can just run multiple times right after another & store the losses if it doesnt fit

[loubnabnl]

yes right! by grad acc. I also meant doing similar iterations over the losses

[loubnabnl]

Summary of the changes I added:

• fix device & size mismathes, (e.g len(dataloader) doesn't work on iterable dataset + take into account the number of processes)
• in transformers the forward pass of GPT2 returns the average loss over the entire batch, and the loss.repeat(batch_size) before calling accelerate.gather
  was just repeating that average so we would've ended up selecting the same value. I changed this line in
  gpt2_modeling of transformers to add reduction="none" (see requirements). This will return the loss of each token so average to get loss per sequence.
• added sanity checks on the order of irreducible losses (this requires using the same batch size.. during their computation and when loading them)
• after selecting a small batch size (32) the same batch would be in all workers so no need to gather values, one thing we will need to add is split the batch (32) to 2 or 3 equal chunks then do grad acc because 32 won't fit in one worker.

[Muennighoff]

• after selecting a small batch size (32) the same batch would be in all workers so no need to gather values, one thing we will need to add is split the batch (32) to 2 or 3 equal chunks then do grad acc because 32 won't fit in one worker.

Amazing work - do you want me to add the last point you mentioned?

[loubnabnl]

You can add it you have time, otherwise I will add it later 🤗

[Muennighoff]

You can add it you have time, otherwise I will add it later 🤗

Done, but not tested. May have a bug 👻