split_batches instead of accumulate_grad_batches
🚀 Feature
Add a split_batches argument to the trainer that splits the batches of your dataloaders into smaller mini-batches (inverse of accumulate_grad_batches), without changing any of the other parts of training.
Optimizers would then only run every n mini-batches and for all other intents and purposes, it would appear the training/validation/testing happens with the original batch size.
Motivation
Imagine the following situation:
Researcher A develops a state-of-the-art technique for a certain tasks and develops an entire training pipeline with PyTorch Lightning. This researcher is part of a big fancy institute and has access to great computing infrastructure with the best GPUs with lots of VRAM. This allows him to train his model with a big batch_size and Researcher A thus figures it is necessary to set the batch_size of his training to 128 in order to reach good results. His entire pipeline (optimizer steps, training length, etc.) is thus built around this batch_size of 128.
He finally beats the SotA in his field and publishes a paper about it! 🎊
Researcher B, from a much smaller institute, reads about this paper and immediately sees a few possibilities for improving this new model! He quickly goes to github, searches for the new technique and is super happy to find that the original authors published some code in PyTorch. Even better, they used his favorite framework, PyTorch Lightning! 🎉
He pulls the code and data, and decides to first try and train a model to check whether he can reach the same accuracy. But alas, there comes the first issue, he does not have enough memory to train a model with that big batch_size. He tries to use smaller batch sizes, but this yields worse results, as this dataset is best trained with big batch sizes. Researcher B knows PyTorch Lightning well and thus uses a smaller batch_size of 32 and sets the accumulate_grad_batches to 4, in order to fake bigger batches of 128.
However, and here comes the frustrating part in all of this, in order to make these changes, researcher B had to go into the LightningDataModule and change the batch_size. He also needed to go into the LightningModules and adapt the optimizers and schedulers, as they were setup to run with an interval of 'step' (and now need to update 4 times slower). Finally, he cannot forget to run the training with the argument accumulate_grad_batches set to 4, otherwise his entire training will behave differently.
The moral of this made-up story is that the accumulate_grad_batches setting of the Trainer is an engineering trick, but in the current implementation it has quite a substantial influence on some of the hyperparameters of the model, which are RnD settings. IMO, the various Trainer parameters should not influence the results of training (or at least have minimal effect on it).
You should be able to define your training hyperparameters from a research/scientific point of view and only have to worry about the engineering part when starting your training, without both choices influencing each other.
Pitch
I would like to add a split_batches=n argument to the Trainer class, which will reduce the batch_size / n of all dataloaders, but run then run the forward pass of the n times before performing running the optimizer/scheduler/...
For every part of the training, it would act as if the training is using batch_size, but only the actual batches would be smaller.
In contrast to accumulate_grad_batches, this flag allows you to choose a different effective (mini-)batch size depending on your system and resources, without affecting anything else in the code (or needing to modify different parts of the model/data to match the original training).
The only minor issue I see with this flag, is that your choice of batch_size might not be divisible by your settings of split_batches. In reality, batches are usually factors of 2 and finding a proper split value should be doable, but if the values do not match, we can do the following:
batch_size = 130
split_batches = 4

mini_batch_size = round(130 / 4) = 32
new_batch_size = 4 * 32 = 128

This would mean the training uses an effective batch_size of 128, which is not the same as the original 130, but is  the closest value that matches and hopefully should only have a minor influence on the results.
We would also log a warning to the user that we are unable to match the original batch_size.
An alternative might be to ceil the mini-batch size, as a bigger batch_size is usually better than a small one, though this is open for discussion.
Alternatives
Obviously, this argument is quite similar to accumulate_grad_batches, but I think I highlighted the reasons for this argument enough in the text above.
I think both arguments can exist next to each other, as long as you dont use both at the same time.
In the long run, you could decide to deprecate accumulate_grad_batches, but that decision is not mine to make.
Additional context
I would love to contribute to this project and what better way than to implement something I am requesting myself.
However, I looked through the Trainer class and noticed it is quite a complicated setup you got there.
I think we would need to adapt the training loop logic in order to run through n batches, and set self.trainer.batch_idx = batch_idx / n.
Alternatively, we could provide something like the tbptt_split_batch functionality, in order to split an existing batch further during runtime. This involves adding some logic to should_accumulate(), in order to disable the backward pass.
If you agree that this feature would be nice to add, it would be nice to get some pointers on where to best start implementing this feature! 😇