Repository: pytorch
Issue: Add `OptState.UPDATED` to `torch.cuda.amp.GradScaler`
Labels: feature,triaged,module: amp (automated mixed precision)

## 🚀 Feature

Add `OptState.UPDATED` to `torch.cuda.amp.GradScaler` to signal that the scaler has been stepped and `optimizer.step()` has been actually called.

## Motivation

`GradScaler` skips the `optimizer.step()` call when Inf/NaNs are found. This means that users can accidentally assume that `optimizer.step()` has been called if they don't know their model produced these types of values.

An example of a bug induced by this is https://github.com/pytorch/pytorch/issues/44511 where we want to run `scheduler.step()` only after `optimizer.step()`, but there's no easy way to be sure that happened when the `GradScaler` is used.

Having this would also make testing the scaler integration simpler for ecosystem libraries.

## Pitch

```diff
def step(self, optimizer, *args, **kwargs):
    if (not self._enabled):
        return optimizer.step(*args, **kwargs)

    if "closure" in kwargs:
        raise RuntimeError("Closure use is not currently supported if GradScaler is enabled.")

    self._check_scale_growth_tracker("step")

    optimizer_state = self._per_optimizer_states[id(optimizer)]

    if optimizer_state["stage"] is OptState.STEPPED:
        raise RuntimeError("step() has already been called since the last update().")

    retval = None

    if (hasattr(optimizer, "_step_supports_amp_scaling") and optimizer._step_supports_amp_scaling):
        # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
        # The contract with custom optimizers is that their step() should accept an additional,
        # optional grad_scaler kwarg.  We append self to the kwargs so the custom optimizer has full information:
        # it can query its own state, invoke unscale_ on itself, etc
        retval = optimizer.step(*args, **dict(kwargs, grad_scaler=self))       
-       optimizer_state["stage"] = OptState.STEPPED
+       optimizer_state["stage"] = OptState.UPDATED
        return retval

    if optimizer_state["stage"] is OptState.READY:
        self.unscale_(optimizer)

    assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."

    retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)

-   optimizer_state["stage"] = OptState.STEPPED
-
    return retval


def _maybe_opt_step(self, optimizer, optimizer_state, *args, **kwargs):
-   retval = None
    if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
        retval = optimizer.step(*args, **kwargs)
+       optimizer_state["stage"] = OptState.UPDATED
+       return retval
+   optimizer_state["stage"] = OptState.STEPPED
+   return
```

This feature request is particularly valuable with the implementation of #67589. With the addition of both, one could do:

```python
scaler.scale(loss).backward()
scaler.step(optimizer)
if scaler.state(optimizer) is OptState.UPDATED:
    lr_scheduler.step()
# note: this needs to be after since it resets `scaler._per_optimizer_states`
scaler.update()
```

## Alternatives

- Not do it.
- Have users subclass `GradScaler`.
- Keep suggesting https://discuss.pytorch.org/t/optimizer-step-before-lr-scheduler-step-error-using-gradscaler/92930/8. IMO this is an ugly solution that relies on the details of the scaler values.

## Additional context

Also fixes https://github.com/pytorch/pytorch/issues/55585

The warning `UserWarning: Detected call of lr_scheduler.step() before optimizer.step()` is impacting Lightning (https://github.com/PyTorchLightning/pytorch-lightning/issues/5558) because of this silent skipping behaviour and we haven't found a clean way to fix it. cc @mcarilli @ptrblck @akihironitta