Lightning-AI / Lightning-AI/pytorch-lightning
MisconfigurationException: Do not set `gradient_accumulation_steps` in the DeepSpeed config
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
I want to use gradient accumulation in my training process which is using a manually configured `DeepSpeedStrategy` (via a config file) to accomplish distributed training. My first go at this is to define the variable `gradient_accumulation_steps` in `deepspeed_config.json` whilst simultaneously passing in the same value in the `Trainer` as an argument. In this case, `lightning` raises the following exception:
```
lightning.fabric.utilities.exceptions.MisconfigurationException: Do not set `gradient_accumulation_steps` in the DeepSpeed config as this will be set with the `accumulate_grad_batches` argument passed via the Lightning Trainer.
```
That's fine, but when I follow this advice and unset `gradient_accumulation_steps` in the `DeepSpeed` config, the `deepspeed` library throws an exception:
```
AssertionError: Check batch related parameters. train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size 256 != 16 * 1 * 8
```
As a consequence I'm unable to use gradient accumulation with `DeepSpeedStrategy`. Am I doing something wrong here or is this actually a conflict of interest between the `deepspeed` and `lightning`. In my eyes, I would imagine it would be sufficient if `lightning` threw a warning here, or followed the `deepspeed` configuration file entirely.
I've tested this using `deepspeed==0.12.6` and `deepspeed==0.14.2` (latest).
### What version are you seeing the problem on?
v2.2
### How to reproduce the bug
```python
// this is the deepspeed config i'm using
{
"train_batch_size": 256,
// "gradient_accumulation_steps": 2,
"train_micro_batch_size_per_gpu": 16,
"gradient_clipping": 1.0,
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4,
"betas": [
0.9,
0.95
],
"eps": 1e-8,
"weight_decay": 1e-2
}
},
"scheduler": {
"type": "WarmupCosineLR",
"params": {
"total_num_steps": 250000,
"warmup_min_ratio": 1e-3,
"warmup_num_steps": 10000,
"cos_min_ratio": 1e-3,
"warmup_type": "linear"
}
},
"zero_optimization": {
"stage": 3,
"allgather_partitions": true,
"allgather_bucket_size": 2e6,
"reduce_scatter": true,
"reduce_bucket_size": 2e6,
"overlap_comm": true,
"contiguous_gradients": true
}
}
```
### Error messages and logs
Here's the full traceback of the first exception thrown by `lightning`.
```
Error executing job with overrides: ['ops.num_workers=48', 'wandb=offline']
Traceback (most recent call last):
File "/home/ubuntu/repos/klaymm/bin/train_lightning.py", line 138, in main
trainer.fit(lit_klaymm, datamodule=datamodule, ckpt_path=cfg.ops.model_ckpt_path)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 544, in fit
call._call_and_handle_interrupt(
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 43, in _call_and_handle_interrupt
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
return function(*args, **kwargs)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 580, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 963, in _run
self.strategy.setup(self)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 335, in setup
self._init_config_if_needed()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 802, in _init_config_if_needed
self._format_config()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 811, in _format_config
self._format_batch_size_and_grad_accum_config()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 898, in _format_batch_size_and_grad_accum_config
raise MisconfigurationException(
lightning.fabric.utilities.exceptions.MisconfigurationException: Do not set `gradient_accumulation_steps` in the DeepSpeed config as this will be set with the `accumulate_grad_batches` argument passed via the Lightni
ng Trainer.
```
And here's the full traceback of the second exception thrown by `deepspeed`.
```
Traceback (most recent call last):
File "/home/ubuntu/repos/klaymm/bin/train_lightning.py", line 138, in main
# fit
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 544, in fit
call._call_and_handle_interrupt(
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 43, in _call_and_handle_interrupt
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
return function(*args, **kwargs)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 580, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 963, in _run
self.strategy.setup(self)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 353, in setup
self.init_deepspeed()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 454, in init_deepspeed
self._initialize_deepspeed_train(self.model)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 490, in _initialize_deepspeed_train
model, deepspeed_optimizer = self._setup_model_and_optimizer(model, optimizer, scheduler)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/lightning/pytorch/strategies/deepspeed.py", line 426, in _setup_model_and_optimizer
deepspeed_engine, deepspeed_optimizer, _, _ = deepspeed.initialize(
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/deepspeed/__init__.py", line 157, in initialize
config_class = DeepSpeedConfig(config, mpu)
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/deepspeed/runtime/config.py", line 786, in __init__
self._configure_train_batch_size()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/deepspeed/runtime/config.py", line 966, in _configure_train_batch_size
self._batch_assertion()
File "/home/ubuntu/miniconda3/envs/klaymm/lib/python3.10/site-packages/deepspeed/runtime/config.py", line 914, in _batch_assertion
assert train_batch == micro_batch * grad_acc * self.world_size, (
AssertionError: Check batch related parameters. train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size 256 != 16 * 1 * 8
```
### Environment
Current environment
```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```
### More info
_No response_
cc @lantiga
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.
Research direction
Start in lightning/pytorch/strategies/deepspeed.py, especially _format_config and _format_batch_size_and_grad_accum_config, and reproduce the two failures with the supplied DeepSpeed JSON and Trainer settings. Done means a manually configured DeepSpeedStrategy can use gradient accumulation without either the Lightning MisconfigurationException or DeepSpeed batch-size assertion.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100