fix(train): check trainable layers separately during checkpointing

swift/trainers/utils.py

@@ -233,17 +233,12 @@ def _kwargs_to_args(func, args, kwargs) -> Optional[List[Any]]:
 
 def _add_gradient_checkpointing(module_list):
 
-    requires_grad = None
-
     def _new_forward(self, *args, **kwargs):
-        nonlocal requires_grad
-        if requires_grad is None:
-            requires_grad = any(p.requires_grad for p in self.parameters())
-
         new_args = _kwargs_to_args(self.__old_forward, args, kwargs)
         if new_args is not None and self.gradient_checkpointing and self.training:
-            if new_args and isinstance(new_args[0], torch.Tensor) and requires_grad and not new_args[0].requires_grad:
-                new_args[0].requires_grad_(True)
+            if new_args and isinstance(new_args[0], torch.Tensor) and not new_args[0].requires_grad:
+                if any(p.requires_grad for p in self.parameters()):
+                    new_args[0].requires_grad_(True)
             layer_ret = self._gradient_checkpointing_func(self.__old_forward, *new_args)
             logger.info_once('Successfully using dynamic gradient checkpointing.')
         else:

tests/general/test_dynamic_checkpointing.py

@@ -0,0 +1,84 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import copy
+import torch
+import unittest
+from functools import partial
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from swift.trainers.utils import dynamic_gradient_checkpointing
+
+
+class Block(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(4, 4)
+
+    def forward(self, hidden_states):
+        return hidden_states + self.proj(hidden_states).tanh()
+
+
+class TinyTower(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([Block() for _ in range(10)])
+        self.head = nn.Linear(4, 1)
+
+    def forward(self, hidden_states):
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return self.head(hidden_states)
+
+
+class TestDynamicCheckpointing(unittest.TestCase):
+
+    def check_gradients(self, model, reference, inputs):
+        model.zero_grad(set_to_none=True)
+        reference.zero_grad(set_to_none=True)
+        actual = model(inputs.clone())
+        expected = reference(inputs.clone())
+        torch.testing.assert_close(actual, expected)
+        actual.square().mean().backward()
+        expected.square().mean().backward()
+        for (name, parameter), (_, reference_parameter) in zip(model.named_parameters(), reference.named_parameters()):
+            with self.subTest(parameter=name):
+                if reference_parameter.grad is None:
+                    self.assertIsNone(parameter.grad)
+                else:
+                    self.assertIsNotNone(parameter.grad)
+                    torch.testing.assert_close(parameter.grad, reference_parameter.grad)
+
+    def prepare(self, model, use_reentrant):
+        dynamic_gradient_checkpointing(model)
+        for layer in model.layers:
+            layer.gradient_checkpointing = True
+            layer._gradient_checkpointing_func = partial(checkpoint, use_reentrant=use_reentrant)
+
+    def test_partially_frozen_tower_matches_eager_gradients(self):
+        for frozen_prefix, use_reentrant in ((5, True), (0, True), (5, False)):
+            with self.subTest(frozen_prefix=frozen_prefix, use_reentrant=use_reentrant):
+                torch.manual_seed(42)
+                model = TinyTower()
+                for layer in model.layers[:frozen_prefix]:
+                    layer.requires_grad_(False)
+                reference = copy.deepcopy(model)
+                self.prepare(model, use_reentrant)
+                self.check_gradients(model, reference, torch.randn(2, 4))
+
+    def test_unfreezing_layers_after_first_forward(self):
+        torch.manual_seed(42)
+        model = TinyTower()
+        model.layers.requires_grad_(False)
+        reference = copy.deepcopy(model)
+        self.prepare(model, True)
+        inputs = torch.randn(2, 4)
+        self.check_gradients(model, reference, inputs)
+        model.layers[4].requires_grad_(True)
+        reference.layers[4].requires_grad_(True)
+        self.check_gradients(model, reference, inputs)
+
+
+if __name__ == '__main__':
+    unittest.main()