returnn 1.20240712.22135__tar.gz → 1.20240712.164051__tar.gz

{returnn-1.20240712.22135 → returnn-1.20240712.164051}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240712.22135
+Version: 1.20240712.164051
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240712.164051/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240712.164051'
+long_version = '1.20240712.164051+git.d5b954b'

{returnn-1.20240712.22135 → returnn-1.20240712.164051}/returnn/torch/util/gradient_checkpoint.py

@@ -530,8 +530,14 @@ def _unregister_custom_saved_tensors_hooks():
         assert thread in _custom_saved_tensors_hooks_registered_threads
         assert (
             not _custom_saved_tensors_hooks_tls_ctx.stack
-            and not _custom_saved_tensors_hooks_tls_ctx.callbacks
+            and (_custom_saved_tensors_hooks_tls_ctx.in_callback or not _custom_saved_tensors_hooks_tls_ctx.callbacks)
             and not _custom_saved_tensors_hooks_tls_ctx.queued_exits
+        ), (
+            f"_unregister_custom_saved_tensors_hooks:"
+            f" stack {_custom_saved_tensors_hooks_tls_ctx.stack},"
+            f" in_callback {_custom_saved_tensors_hooks_tls_ctx.in_callback},"
+            f" callbacks {_custom_saved_tensors_hooks_tls_ctx.callbacks},"
+            f" queued_exits {_custom_saved_tensors_hooks_tls_ctx.queued_exits}"
         )
         _custom_saved_tensors_hooks_tls_ctx.active = False
         _custom_saved_tensors_hooks_registered_threads.remove(thread)
@@ -579,6 +585,7 @@ def _register_custom_saved_tensors_hooks_thread_local_callback(cb: Callable[[],
     Register some thread-local callback function which is called on saved_tensors_hooks __enter__ and __exit__.
     If it returns True, it is kept alive, otherwise removed.
     """
+    assert not _custom_saved_tensors_hooks_tls_ctx.in_callback
     _custom_saved_tensors_hooks_tls_ctx.callbacks.append(cb)
 
 
@@ -590,5 +597,7 @@ def _custom_saved_tensors_hooks_call_callbacks():
         _custom_saved_tensors_hooks_tls_ctx.callbacks = [
             cb for cb in _custom_saved_tensors_hooks_tls_ctx.callbacks if cb()
         ]
+        if not _custom_saved_tensors_hooks_tls_ctx.active:  # was cleaned up by some of the callbacks
+            assert not _custom_saved_tensors_hooks_tls_ctx.callbacks and not _custom_saved_tensors_hooks_tls_ctx.stack
     finally:
         _custom_saved_tensors_hooks_tls_ctx.in_callback = False

{returnn-1.20240712.22135 → returnn-1.20240712.164051}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240712.22135
+Version: 1.20240712.164051
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20240712.22135 → returnn-1.20240712.164051}/tests/test_torch_util.py

@@ -151,6 +151,51 @@ def test_gradient_checkpoint_scope():
         torch.testing.assert_allclose(param_post_state[k], param_post_state_[k])
 
 
+def test_gradient_checkpoint_scope_twice():
+    # https://github.com/rwth-i6/returnn/issues/1579
+    shape = (101, 103)
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.var = torch.nn.Parameter(torch.randn(shape))
+            self.input_var = torch.nn.Parameter(torch.randn(shape))
+            self.opt = torch.optim.SGD(self.parameters(), lr=0.1)  # not common to have this here but ok for the test
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.get_var() * x
+
+        def get_var(self) -> torch.Tensor:
+            with gradient_checkpoint_scope():
+                return self.var + torch.randn(shape)
+
+        def get_input(self) -> torch.Tensor:
+            x = self.input_var
+            with gradient_checkpoint_scope():
+                return x + torch.randn(shape)
+
+        def demo_run(self):
+            self.opt.zero_grad()
+            y = self(self.get_input())
+            loss = y.sum()  # dummy loss
+            del y  # not needed anymore
+            loss.backward()
+            del loss  # not needed anymore
+            self.opt.step()
+
+    orig_gradient_checkpoint_scope_tensor_del_hook = gradient_checkpoint_scope._tensor_del_hook
+    try:
+        # Overwrite this here to trigger the case where the tensor del hook will not do the cleanup.
+        gradient_checkpoint_scope._tensor_del_hook = lambda self: None
+
+        model = _Model()
+        model.demo_run()
+        model.demo_run()
+
+    finally:
+        gradient_checkpoint_scope._tensor_del_hook = orig_gradient_checkpoint_scope_tensor_del_hook
+
+
 def _report_profile(prof: torch.profiler.profiler, check_events=(), *, _size_threshold=100):
     # Note: I tried prof.events(), prof.profiler.kineto_results.events(), prof._memory_profile().timeline,
     # but they all are not really giving me the information I want.

returnn-1.20240712.22135/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240712.022135'
-long_version = '1.20240712.022135+git.3da9c4b'