Repository: allenai/tango (language: python)

Issue:
TorchTrainStep should save GradScaler state with checkpoints
Should save and restore `scale`, `growth_factor`, `backoff_factor`, and `growth_interval`.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b189de960..7a25a26f6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - If you try to import something from a tango integration that is not fully installed due to missing dependencies, an `IntegrationMissingError` will be raised
 instead of `ModuleNotFound`.
+- Fixed bug which caused GradScaler state to not be saved and loaded with checkpoints. 
 
 ### Added
 
diff --git a/tango/integrations/torch/training_engine.py b/tango/integrations/torch/training_engine.py
index 4dc447a9f..3f9c3f79f 100644
--- a/tango/integrations/torch/training_engine.py
+++ b/tango/integrations/torch/training_engine.py
@@ -281,6 +281,8 @@ def save_state(state: Dict[str, Any], name: str):
         save_state(self.optimizer.state_dict(), "optimizer"),
         if self.lr_scheduler is not None:
             save_state(self.lr_scheduler.state_dict(), "lr_scheduler")
+        if self.grad_scaler is not None:
+            save_state(self.grad_scaler.state_dict(), "grad_scaler")
         save_state(client_state, "trainer")
 
     def load_checkpoint(self, checkpoint_dir: Path) -> Dict[str, Any]:
@@ -294,6 +296,10 @@ def load_checkpoint(self, checkpoint_dir: Path) -> Dict[str, Any]:
             self.lr_scheduler.load_state_dict(
                 torch.load(checkpoint_dir / f"worker{self.train_config.worker_id}_lr_scheduler.pt")
             )
+        if self.grad_scaler is not None:
+            self.grad_scaler.load_state_dict(
+                torch.load(checkpoint_dir / f"worker{self.train_config.worker_id}_grad_scaler.pt")
+            )
         return torch.load(checkpoint_dir / f"worker{self.train_config.worker_id}_trainer.pt")
 
     def save_complete_weights_from_checkpoint(


Test patch:
diff --git a/tests/integrations/torch/training_engine_test.py b/tests/integrations/torch/training_engine_test.py
new file mode 100644
index 000000000..c6ff4ec67
--- /dev/null
+++ b/tests/integrations/torch/training_engine_test.py
@@ -0,0 +1,41 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from tango.common.testing import TangoTestCase
+from tango.integrations.torch.model import Model
+from tango.integrations.torch.training_engine import TorchTrainingEngine
+
+
+@Model.register("dummy_model")
+class DummyModel(Model):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 1)
+
+    def forward(self, x, y=None):
+        return self.linear(x)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Requires CUDA devices")
+class TestTorchTrainingEngine(TangoTestCase):
+    def test_grad_scaler(self):
+        training_engine = TorchTrainingEngine.from_params(
+            {
+                "train_config": {"step_id": "001", "work_dir": self.TEST_DIR},
+                "model": {
+                    "type": "dummy_model",
+                },
+                "optimizer": {"type": "torch::Adam"},
+                "amp": True,
+            }
+        )
+
+        state_dict = {"training_steps": None}
+        training_engine.save_checkpoint(self.TEST_DIR, state_dict)
+        saved_grad_scaler = training_engine.grad_scaler
+        training_engine.load_checkpoint(self.TEST_DIR)
+
+        assert (self.TEST_DIR / "worker0_grad_scaler.pt").is_file()
+        assert training_engine.grad_scaler == saved_grad_scaler