Lightning-AI / Lightning-AI/pytorch-lightning
Error when disabling an optimizer with native AMP turned on
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'm using 2 optimizers and trying to train with AMP (FP16). I can take steps with my first optimizer. When I take my first step with the second optimizer I get the following error:
```
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 450, in step
len(optimizer_state["found_inf_per_device"]) > 0
AssertionError: No inf checks were recorded for this optimizer.
```
I can train this correctly in FP32 -- so it seems to be an issue with AMP.
### What version are you seeing the problem on?
version 2.3.3
### How to reproduce the bug
```python
def training_step(self, batch: Dict, batch_idx: int):
"""
We have 2 sets of optimizers.
Every N batches (self.n_batches_per_optimizer), we make an optimizer update and
switch the optimizer to update.
If self.n_batches_per_optimizer = 1, then we make updates every batch and alternate optimizers
every batch.
If self.n_batches_per_optimizer > 1, then we're doing gradient accumulartion, where we are making
updates evern n_batches_per_optimizer batches and alternating optimizers every n_batches_per_optimizer
batches.
"""
opts = self.optimizers()
current_cycle = (batch_idx // self.n_batches_per_optimizer) % len(opts)
opt = opts[current_cycle]
opt.zero_grad()
if current_cycle == 0:
compute_model_1_loss = True
elif current_cycle == 1:
compute_model_1_loss = False
else:
raise NotImplementedError(f"Unknown optimizer {current_cycle}")
with opt.toggle_model():
loss = self.inner_training_step(batch=batch, compute_model_1_loss=compute_model_1_loss)
self.manual_backward(loss=loss)
# Perform the optimization step every accumulate_grad_batches steps
if (batch_idx + 1) % self.n_batches_per_optimizer == 0:
if not compute_model_1_loss:
print("About to take compute model 2 loss ...")
opt.step()
opt.zero_grad()
```
### Error messages and logs
```
Traceback (most recent call last):
File "/home/sahil/train.py", line 82, in
main(config)
File "/home/sahil/train.py", line 62, in main
trainer.fit(model, datamodule=data_module, ckpt_path=ckpt)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit
call._call_and_handle_interrupt(
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 44, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 986, in _run
results = self._run_stage()
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1030, in _run_stage
self.fit_loop.run()
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
self.advance()
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
self.epoch_loop.run(self._data_fetcher)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run
self.advance(data_fetcher)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 252, in advance
batch_output = self.manual_optimization.run(kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/manual.py", line 94, in run
self.advance(kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/manual.py", line 114, in advance
training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 311, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 390, in training_step
return self.lightning_module.training_step(*args, **kwargs)
File "/home/sahil/model/model.py", line 169, in training_step
opt.step()
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 153, in step
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 238, in optimizer_step
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 93, in optimizer_step
step_output = self.scaler.step(optimizer, **kwargs) # type: ignore[arg-type]
File "/home/sahil/.cache/pypoetry/virtualenvs/env-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 450, in step
len(optimizer_state["found_inf_per_device"]) > 0
AssertionError: No inf checks were recorded for this optimizer.
```
### Environment
Current environment
```
* CUDA:
- GPU:
- NVIDIA A100-SXM4-80GB
- available: True
- version: 12.1
* Lightning:
- lightning-utilities: 0.11.5
- pytorch-lightning: 2.3.3
- torch: 2.3.1
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1
* Packages:
- aiohttp: 3.9.5
- aiosignal: 1.3.1
- annotated-types: 0.7.0
- antlr4-python3-runtime: 4.9.3
- anyio: 4.4.0
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.3.0
- asttokens: 2.4.1
- async-lru: 2.0.4
- async-timeout: 4.0.3
- attrs: 23.2.0
- autocommand: 2.2.2
- babel: 2.15.0
- backports.tarfile: 1.2.0
- beautifulsoup4: 4.12.3
- bitsandbytes: 0.43.1
- bleach: 6.1.0
- boto3: 1.34.144
- botocore: 1.34.144
- braceexpand: 0.1.7
- certifi: 2024.7.4
- nvidia-curand-cu12: 10.3.2.106
- nvidia-cusolver-cu12: 11.4.5.107
- nvidia-cusparse-cu12: 12.1.0.106
- nvidia-nccl-cu12: 2.20.5
- nvidia-nvjitlink-cu12: 12.5.82
- nvidia-nvtx-cu12: 12.1.105
- omegaconf: 2.3.0
- opencv-python: 4.10.0.84
- ordered-set: 4.1.0
- overrides: 7.7.0
- packaging: 24.1
- pandocfilters: 1.5.1
- parso: 0.8.4
- pexpect: 4.9.0
- pillow: 10.4.0
- pip: 24.1
- platformdirs: 4.2.2
- pre-commit: 3.7.1
- proglog: 0.1.10
- prometheus-client: 0.20.0
- prompt-toolkit: 3.0.47
- protobuf: 5.27.2
- psutil: 6.0.0
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pycparser: 2.22
- pydantic: 2.8.2
- pydantic-core: 2.20.1
- pydantic-settings: 2.3.4
- pygments: 2.18.0
- python-dateutil: 2.9.0.post0
- python-dotenv: 1.0.1
- python-json-logger: 2.0.7
- pytorch-lightning: 2.3.3
- pyyaml: 6.0.1
- pyzmq: 26.0.3
- referencing: 0.35.1
- requests: 2.32.3
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rpds-py: 0.19.0
- s3transfer: 0.10.2
- send2trash: 1.8.3
- sentry-sdk: 2.10.0
- setproctitle: 1.3.3
- setuptools: 71.0.2
- six: 1.16.0
- smmap: 5.0.1
- sniffio: 1.3.1
- soupsieve: 2.5
- stack-data: 0.6.3
- sympy: 1.13.0
- terminado: 0.18.1
- tinycss2: 1.3.0
- tomli: 2.0.1
- torch: 2.3.1
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1
- tornado: 6.4.1
- tqdm: 4.66.4
- traitlets: 5.14.3
- triton: 2.3.1
- typeguard: 4.3.0
- types-python-dateutil: 2.9.0.20240316
- typing-extensions: 4.12.2
- uri-template: 1.3.0
- urllib3: 2.2.2
- virtualenv: 20.26.3
- wandb: 0.17.4
- wcwidth: 0.2.13
- webcolors: 24.6.0
- webdataset: 0.2.86
- webencodings: 0.5.1
- websocket-client: 1.8.0
- wheel: 0.43.0
- yarl: 1.9.4
- zipp: 3.19.2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor:
- python: 3.10.14
- release: 5.10.0-31-cloud-amd64
- version: #1 SMP Debian 5.10.221-1 (2024-07-14)
```
### More info
_No response_
cc @borda @justusschock @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 with the reproduced training_step flow and inspect pytorch_lightning/plugins/precision/amp.py, especially optimizer_step and its call into the scaler. Compare the first and second optimizer paths with native AMP enabled. Done means alternating optimizers no longer raises “No inf checks were recorded for this optimizer” while FP32 behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100