Lightning-AI / Lightning-AI/pytorch-lightning

Mixed precision: scheduler and optimizer are called in the wrong order

Open
#5,558 36 comments 16 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug lr scheduler precision: amp priority: 2
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🐛 Bug

When using mixed-precision training, scheduler and optimizer are called in the wrong order. Warning is generated:

```
UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.
```

## Please reproduce using the BoringModel

https://colab.research.google.com/drive/1G7pk6E9XUYq-pS41DXKhqM9Srx8sikiP?usp=sharing

There are four tests. Three of them doesn't raise the warning:

1. test_amp_scheduler(precision=16, configure_optimizers=configure_optimizers_1)
2. test_amp_scheduler(precision=32, configure_optimizers=configure_optimizers_1)
3. test_amp_scheduler(precision=32, configure_optimizers=configure_optimizers_2)

This testcase raises the warning:

1. test_amp_scheduler(precision=16, configure_optimizers=configure_optimizers_2)

### To Reproduce

1. Create model with `configure_optimizers` in a following dictionary style:

```python
def configure_optimizers_2(model):
optimizer = torch.optim.SGD(model.layer.parameters(), lr=0.1)
scheduler = {'scheduler': torch.optim.lr_scheduler.StepLR(optimizer, step_size=1),
'name': 'learning_rate',
'interval':'step',
'frequency': 1}

return {"optimizer": optimizer, "lr_scheduler": scheduler}
```

2. Enable mixed-precision training by setting `precision=16` in a `Trainer`
3. Start training

#### Note

When scheduler is defined in another way, the issue seems to not occur:

```python
def configure_optimizers_1(model):
optimizer = torch.optim.SGD(model.layer.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)

return {"optimizer": optimizer, "lr_scheduler": scheduler}
```

### Expected behavior

No warning

### Environment

* CUDA:
- GPU:
- Tesla P100-PCIE-16GB
- available: True
- version: 10.1
* Packages:
- numpy: 1.19.5
- pyTorch_debug: True
- pyTorch_version: 1.7.0+cu101
- pytorch-lightning: 1.1.4
- tqdm: 4.41.1
* System:
- OS: Linux
- architecture:
- 64bit
-
- processor: x86_64
- python: 3.6.9
- version: #1 SMP Thu Jul 23 08:00:38 PDT 2020

cc @tchaton @rohitgr7 @carmocca @justusschock @awaelchli @akihironitta

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the four BoringModel test cases described in the issue, especially test_amp_scheduler with precision=16 and configure_optimizers_2. Compare the dictionary-style scheduler configuration with configure_optimizers_1 and trace the mixed-precision training flow. Done means the failing case no longer emits the optimizer-order warning while the other cases continue to pass.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.