returnn 1.20240714.161251__tar.gz → 1.20240715.155107__tar.gz

{returnn-1.20240714.161251 → returnn-1.20240715.155107}/PKG-INFO

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

returnn-1.20240715.155107/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240715.155107'
+long_version = '1.20240715.155107+git.23d666c'

{returnn-1.20240714.161251 → returnn-1.20240715.155107}/returnn/torch/frontend/_backend.py

@@ -529,10 +529,10 @@ class TorchBackend(Backend[torch.Tensor]):
         f = getattr(torch, func, None)
         if not f:
             f = getattr(torch.nn.functional, func, None)
-        if not f:
-            f = TorchBackend._ActivationFuncMapping.get(func)
-        if not f:
-            raise ValueError(f"unknown activation function {func!r}")
+            if not f:
+                f = TorchBackend._ActivationFuncMapping.get(func)
+                if not f:
+                    raise ValueError(f"unknown activation function {func!r}")
         return f(raw_tensor)
 
     @staticmethod

{returnn-1.20240714.161251 → returnn-1.20240715.155107}/returnn/torch/util/gradient_checkpoint.py

@@ -206,7 +206,8 @@ class gradient_checkpoint_scope:
             # so we can exit saved_tensors_hooks_scope now.
             # (We might not always catch this properly in the Tensor _DelHook,
             #  e.g. when Tensor.__del__ runs in a different thread.)
-            self.exit_saved_tensors_hooks_scope()
+            if _can_exit_saved_tensors_hooks_inside_hooks():
+                self.exit_saved_tensors_hooks_scope()
             return x
         # _RecordGraph.__torch_dispatch__ should have recorded all newly created tensors.
         x_ = self.record_graph_scope.graph.graph_tensor_from_weak_raw_tensor.get(x, x)
@@ -217,7 +218,8 @@ class gradient_checkpoint_scope:
     @staticmethod
     def _unpack_hook(x: Union[torch.Tensor, _GraphTensor]) -> torch.Tensor:
         if isinstance(x, _GraphTensor):
-            x.op.graph.gradient_checkpoint_scope_backref._maybe_exit_saved_tensors_hooks_scope()
+            if _can_exit_saved_tensors_hooks_inside_hooks():
+                x.op.graph.gradient_checkpoint_scope_backref._maybe_exit_saved_tensors_hooks_scope()
             x.op.graph.maybe_recompute()
             return x.get_recomputed()
         return x
@@ -495,6 +497,18 @@ class _WeakMethod:
         return self.func(obj, *args, **kwargs)
 
 
+def _can_exit_saved_tensors_hooks_inside_hooks() -> bool:
+    """
+    Check whether we can call torch.autograd.graph.saved_tensors_hooks.__exit__
+    inside the pack or unpack hook itself.
+
+    https://github.com/rwth-i6/returnn/issues/1581
+    https://github.com/pytorch/pytorch/issues/130734
+    """
+    # TODO Return True in some later PyTorch version? Check https://github.com/pytorch/pytorch/issues/130734.
+    return False
+
+
 _orig_saved_tensors_hooks_enter = torch.autograd.graph.saved_tensors_hooks.__enter__
 _orig_saved_tensors_hooks_exit = torch.autograd.graph.saved_tensors_hooks.__exit__
 _custom_saved_tensors_hooks_tls_ctx = threading.local()
@@ -559,7 +573,9 @@ def _custom_saved_tensors_hooks_enter(self: torch.autograd.graph.saved_tensors_h
     return _orig_saved_tensors_hooks_enter(self)
 
 
-def _custom_saved_tensors_hooks_exit(self: torch.autograd.graph.saved_tensors_hooks, exc_type, exc_val, exc_tb):
+def _custom_saved_tensors_hooks_exit(
+    self: torch.autograd.graph.saved_tensors_hooks, exc_type=None, exc_val=None, exc_tb=None
+):
     if self not in _custom_saved_tensors_hooks_tls_ctx.stack:
         raise Exception(
             f"saved_tensors_hooks __exit__ mismatch."

{returnn-1.20240714.161251 → returnn-1.20240715.155107}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240714.161251 → returnn-1.20240715.155107}/tests/test_torch_util.py

@@ -12,7 +12,9 @@ import unittest
 import torch
 
 from returnn.util import better_exchook
-from returnn.torch.util.gradient_checkpoint import gradient_checkpoint_scope
+
+# noinspection PyProtectedMember
+from returnn.torch.util.gradient_checkpoint import gradient_checkpoint_scope, _can_exit_saved_tensors_hooks_inside_hooks
 
 
 def test_gradient_checkpoint_scope():
@@ -196,6 +198,45 @@ def test_gradient_checkpoint_scope_twice():
         gradient_checkpoint_scope._tensor_del_hook = orig_gradient_checkpoint_scope_tensor_del_hook
 
 
+def test_saved_tensors_hooks_gc_segfault():
+    # https://github.com/rwth-i6/returnn/issues/1581
+    # https://github.com/pytorch/pytorch/issues/130734
+    if not _can_exit_saved_tensors_hooks_inside_hooks():
+        raise unittest.SkipTest("Not yet fixed.")
+
+    shape = (101, 103)
+    for i in range(10):
+        print("**** iter", i)
+        v = torch.nn.Parameter(torch.randn(shape))
+
+        class _Handler:
+            def __init__(self):
+                self.scope = torch.autograd.graph.saved_tensors_hooks(self._pack_hook, self._unpack_hook)
+                self.scope.__enter__()
+                self.exited = False
+
+            def _pack_hook(self, x):
+                print(f"*** _pack_hook {self}")
+                return x
+
+            def _unpack_hook(self, x):
+                print(f"*** _unpack_hook {self}")
+                if not self.exited:
+                    self.exited = True
+                    print(
+                        f"*** exit {self.scope},"
+                        f" pack_hook {hex(id(self.scope.pack_hook))},"
+                        f" unpack_hook {hex(id(self.scope.unpack_hook))}"
+                    )
+                    self.scope.__exit__()
+                return x
+
+        with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
+            handler = _Handler()  # keep ref...  # noqa
+            x = v * torch.randn(shape)
+            x.sum().backward()
+
+
 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.20240714.161251/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240714.161251'
-long_version = '1.20240714.161251+git.2e54d6a'