Add amp_scaling_state (apex) migration (#16161)

Co-authored-by: Adrian Wälchli <[email protected]>
Fixes #16149 (comment)

Author: carmocca

src/pytorch_lightning/CHANGELOG.md

@@ -7,6 +7,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [fabricLaunchPLVersion] - 202Y-MM-DD
 
+### Added
+
+- Added migration logic to warn about checkpoints with apex AMP state ([#16161](https://github.com/Lightning-AI/lightning/pull/16161))
+
 ### Removed
 
 - Removed the `pytorch_lightning.lite` module in favor of `lightning_fabric` ([#15953](https://github.com/Lightning-AI/lightning/pull/15953))

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -293,8 +293,6 @@ def restore_precision_plugin_state(self) -> None:
             prec_plugin.load_state_dict(self._loaded_checkpoint[prec_plugin.__class__.__qualname__])
 
         # old checkpoints compatibility
-        if "amp_scaling_state" in self._loaded_checkpoint:
-            rank_zero_warn("This checkpoint contains apex AMP data, but apex support has been removed.")
         if "native_amp_scaling_state" in self._loaded_checkpoint and isinstance(prec_plugin, MixedPrecisionPlugin):
             prec_plugin.load_state_dict(self._loaded_checkpoint["native_amp_scaling_state"])
 

src/pytorch_lightning/utilities/migration/migration.py

@@ -46,6 +46,7 @@ def _migration_index() -> Dict[str, List[Callable[[_CHECKPOINT], _CHECKPOINT]]]:
         "1.6.0": [_migrate_loop_global_step_to_progress_tracking, _migrate_loop_current_epoch_to_progress_tracking],
         "1.6.5": [_migrate_loop_batches_that_stepped],
         "1.9.0": [_migrate_model_checkpoint_save_on_train_epoch_end_default],
+        "2.0.0": [_drop_apex_amp_state],
     }
 
 
@@ -203,3 +204,18 @@ def new_key(old_key: str) -> str:
 
     checkpoint["callbacks"] = new_callback_states
     return checkpoint
+
+
+def _drop_apex_amp_state(checkpoint: _CHECKPOINT) -> _CHECKPOINT:
+    """Apex support was removed in v2.0.0, and this migration drops it from the state-keys saved in the checkpoint
+    dict.
+
+    Version: 2.0.0
+    Commit: e544676ff434ed96c6dd3b4e73a708bcb27ebcf1
+    PR: #16149
+    """
+    key = "amp_scaling_state"
+    if key in checkpoint:
+        rank_zero_warn("This checkpoint contains apex AMP data, but apex support has been removed in v2.0.0.")
+        del checkpoint[key]
+    return checkpoint

tests/tests_pytorch/utilities/migration/test_migration.py

@@ -146,3 +146,13 @@ def test_migrate_model_checkpoint_save_on_train_epoch_end_default_collision():
     with pytest.warns(PossibleUserWarning, match="callback states in this checkpoint.* colliding with each other"):
         updated_checkpoint, _ = migrate_checkpoint(old_checkpoint.copy(), target_version="1.9.0")
     assert updated_checkpoint["callbacks"] == old_checkpoint["callbacks"]  # no migration was performed
+
+
+def test_migrate_dropped_apex_amp_state(monkeypatch):
+    """Test that the migration warns about collisions that would occur if the keys were modified."""
+    monkeypatch.setattr(pl, "__version__", "2.0.0")  # pretend this version of Lightning is >= 2.0.0
+    old_checkpoint = {"amp_scaling_state": {"scale": 1.23}}
+    _set_version(old_checkpoint, "1.9.0")  # pretend a checkpoint prior to 2.0.0
+    with pytest.warns(UserWarning, match="checkpoint contains apex AMP data"):
+        updated_checkpoint, _ = migrate_checkpoint(old_checkpoint.copy())
+    assert "amp_scaling_state" not in updated_checkpoint