Batch size schedulers
Author: BramVanroy
Pull request: no
Labels: name: trainer

name: Feature request

### Feature request

Implement in the trainer the possibility to increment the batch size according to some schedule, similar to how lr can follow a schedule.

### Motivation

It's been shown (e.g. DeepSeek) that slowly increasing the batch size can have a positive effect on learning.

### Your contribution

Unfortunately no time to work on this. It's not trivial to add this, so it's really just an idea to throw out there.

Comments:
cc @younesbelkada @muellerzr 

I'll look at trying to do this today, should be very straightforward. Per my understanding reading DeepSeek v2 essentially:
- Increased bs from 2304 to 9216, or translated as a per-device bs of 288 -> 1152, starting pct of 25% -> 100%.
- This ramp-up occurred in the first 225B, and eventually training to their total 8.1T tokens, or ~2.78% of the total iterations

BTW @BramVanroy, didn't they say that you could do one or the other, and they're effectively the same? 

Changing the batch size dynamically can be done, but as you said it's a bit of a workaround to do so (basically we need to keep recreating the dataloader/sampler each step)

> BTW @BramVanroy, didn't they say that you could do one or the other, and they're effectively the same?
> 
> Changing the batch size dynamically can be done, but as you said it's a bit of a workaround to do so (basically we need to keep recreating the dataloader/sampler each step)

If "one" is to incrementally increase the batch size, then what's the "other"? In their DeepSeek-Lite model they did not use the batch size strategy, but I have not read any ablations around this so I am not sure about the performance difference.

I was going off of Soumith's comment here: https://github.com/pytorch/pytorch/issues/10470#issuecomment-412598294

As you are right, it's a bit different as they decayed LR and increased BS

@BramVanroy just to update you, we're working with the PyTorch team to see what the best way forward here would be, because legally PyTorch doesn't allow this 😅 

> @BramVanroy just to update you, we're working with the PyTorch team to see what the best way forward here would be, because legally PyTorch doesn't allow this 😅

What about illegally? 😏

Thanks for following up!

Mini-up here, would be a cool feature to have in the trainer !

@muellerzr @BramVanroy 

it would be even more useful, if we can determine the batch size dynamically and not through a pre-defined batch size scheduler. this would be helpful for curriculum learning approaches were the number of novel datapoints to draw from the training set is better to be adaptively decided

Hello everyone! +1 on the usefullness of the feature. Any progress on it? Is there any dirty workaround to consider? Like, where exactly should I be re-creating the dataloader/how do I ensure the recreation won't harm the items order shuffling?


+1

@BramVanroy @JackCai1206  @oserikov @marthos1 @ManuelFay @karynaur @mohblnk @amyeroberts @NicolasMejiaPetit
Hi all,

here's a repo supporting this. let me know what you think :)

https://github.com/ifiaposto/hf_trainer_with_online_batch_size