Enable gradient accumulation for HuBERT recipe
@nateanl is already working on this.
Since Dec 13, 2022.
- Dominant language
- Python
- Stars
- 2.9k
- Forks
- 799
- Avg merge
- 58m
- Merged PRs (30d)
- 3
Description
🚀 The feature
It would be nice if gradient accumulation functionality could be added to the HuBERT recipe.
Motivation, pitch
Using gradient accumulation can simulate a larger cluster / larger effective batch sizes, for the purpose of replicating others' results or maintaining consistency across experiments that run on different numbers of GPUs.
For example, FAIR's HuBERT experiments ran on 32+ GPUs, but I generally have access to fewer devices than that.
Alternatives
In an earlier version of the recipes, before HuBERTPreTrainModule.automatic_optimization was set to False in #2744, it was as simple as passing accumulate_grad_batches=... as an argument to PyTorch Lightning's Trainer. However, that now fails with this message:
MisconfigurationException: Automatic gradient accumulation is not supported for manual optimization.
Remove `Trainer(accumulate_grad_batches=...)` or switch to automatic optimization.
So we could perhaps revert to the automatic optimization -- but then we'd lose the training step's customized NaN-handling, loss normalization, gradient clipping, and AMP training.
Alternatively, one can try to tweak the learning rate parameters (max and schedule) for each experiment, but that's rather tricky.
Additional context
I think some gradient accumulation functionality could be added to the custom training_step by tracking the batch_idx and only performing an update on accumulated gradients every N batches. However, this it might not have the same semantics as Trainer(accumulate_grad_batches=...) because the user would need to accordingly scale the number of max updates.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.