returnn 1.20240718.93955__tar.gz → 1.20240719.111324__tar.gz

{returnn-1.20240718.93955 → returnn-1.20240719.111324}/PKG-INFO

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

returnn-1.20240719.111324/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240719.111324'
+long_version = '1.20240719.111324+git.12f2abc'

{returnn-1.20240718.93955 → returnn-1.20240719.111324}/returnn/frontend/attention.py

@@ -542,7 +542,7 @@ class RelPosCausalSelfAttention(CausalSelfAttention):
         output, _ = rf.merge_dims(att, dims=(self.num_heads, self.value_dim_per_head), out_dim=self.value_dim_total)
         if self.proj:
             output = self.proj(output)
-        return output
+        return output, new_state
 
 
 class CrossAttention(rf.Module):

{returnn-1.20240718.93955 → returnn-1.20240719.111324}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240718.93955 → returnn-1.20240719.111324}/returnn.egg-info/SOURCES.txt

@@ -385,6 +385,7 @@ tests/test_torch_engine.py
 tests/test_torch_frontend.py
 tests/test_torch_internal_frontend.py
 tests/test_torch_util.py
+tests/torch_utils.py
 tests/PyCharm.idea/.gitignore
 tests/PyCharm.idea/.name
 tests/PyCharm.idea/codeStyleSettings.xml

{returnn-1.20240718.93955 → returnn-1.20240719.111324}/tests/test_torch_util.py

@@ -11,6 +11,8 @@ import sys
 import unittest
 import torch
 
+from torch_utils import report_profile
+
 from returnn.util import better_exchook
 
 
@@ -62,7 +64,7 @@ def test_gradient_checkpoint_scope():
     b = 4  # size single f32
     t = shape[0] * shape[1] * b  # size tensor
     r = rng_state.numel() * rng_state.element_size()
-    _report_profile(
+    report_profile(
         prof,
         [
             # ignore private calls
@@ -95,7 +97,7 @@ def test_gradient_checkpoint_scope():
     with profile(activities=[ProfilerActivity.CPU], profile_memory=True, with_stack=True, record_shapes=True) as prof:
         with record_function("train_step_grad_ckpt"):
             model.demo_run()
-    _report_profile(
+    report_profile(
         prof,
         [
             ("pycall", {"callsite_name": "demo_run"}),
@@ -245,135 +247,6 @@ def test_saved_tensors_hooks_gc_segfault():
             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.
-    # Either the Python stack is missing, or the memory information is incomplete,
-    # or the Python/TorchOp events are missing.
-    # The only complete information source seems to be prof.profiler.kineto_results.experimental_event_tree().
-
-    import fnmatch
-
-    # noinspection PyProtectedMember
-    from torch.profiler._utils import traverse_dfs
-    from torch._C._profiler import _EventType  # noqa
-
-    _allocs = {}  # id -> dict with "size", "name"
-    check_events = list(check_events)
-
-    def _ev_visit(ev):
-        # ev: torch._C._profiler._ProfilerEvent
-        if ev.typed[0] == _EventType.Allocation:
-            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_Allocation
-            # ex.id/ex.allocation_id/ex.ptr redundant?
-            if ex.allocation_id in _allocs:
-                ev_name = "dealloc"  # deallocation
-                assert _allocs[ex.allocation_id]["size"] == -ex.alloc_size
-                name = _allocs[ex.allocation_id]["name"]
-                del _allocs[ex.allocation_id]
-            else:
-                ev_name = "alloc"
-                assert ex.alloc_size > 0
-                assert ev.parent
-                name = _ctx(ev.parent)
-                _allocs[ex.allocation_id] = {"size": ex.alloc_size, "name": name}
-            opts = {"id": ex.allocation_id, "name": name, "size": ex.alloc_size, "total_alloc": ex.total_allocated}
-        elif ev.typed[0] == _EventType.TorchOp:
-            ev_name = "torchop"
-            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_TorchOp
-            opts = {"name": ex.name}
-        elif ev.typed[0] == _EventType.PyCall:
-            ev_name = "pycall"
-            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_PyCall
-            ex0 = ex.caller  # torch._C._profiler._PyFrameState
-            ex1 = ex.callsite  # torch._C._profiler._PyFrameState
-            if _pycall_filter_fn(ex0.file_name) or _pycall_filter_fn(ex1.file_name):
-                opts = {
-                    "caller_loc": f"{ex0.file_name}:{ex0.line_number}",
-                    "caller_name": ex0.function_name,
-                    "callsite_name": ex1.function_name,
-                }
-            else:
-                return
-        else:
-            return
-
-        next_check = check_events[0] if check_events else None
-        if next_check:
-            next_check_name, next_check_opts = next_check
-            if ev_name == next_check_name:
-                for k, v in next_check_opts.items():
-                    if isinstance(v, str) and "*" in v:
-                        if not fnmatch.fnmatch(opts[k], v):
-                            mismatch = f"Pattern mismatch: {opts[k]} vs {v}"
-                            break
-                    elif k == "total_alloc":
-                        if abs(opts[k] - v) >= _size_threshold:
-                            mismatch = f"Size mismatch: {opts[k]} vs {v}"
-                            break
-                    elif opts[k] != v:
-                        mismatch = f"Value mismatch: {opts[k]} vs {v}"
-                        break
-                else:
-                    mismatch = None
-            else:
-                mismatch = f"Different event: {ev_name} vs {next_check_name}"
-        else:
-            mismatch = "No check event"
-
-        if ev_name in {"alloc", "dealloc"} and abs(opts["size"]) >= _size_threshold:
-            assert not mismatch, f"Event not matched: {ev_name} {opts} to {next_check}: {mismatch}"
-
-        if not mismatch:
-            print(f"{ev_name} {opts} ✓")
-            check_events.pop(0)
-        else:
-            print(f"({ev_name} {opts})")
-
-    def _ctx(ev) -> str:
-        stack = [None]
-        parent = ev
-        while parent and parent.typed[0] == _EventType.TorchOp:  # go to top torch op
-            stack[-1] = parent.typed[1].name
-            parent = parent.parent
-        if not stack[-1] and parent.typed[0] == _EventType.PyCCall:
-            stack[-1] = parent.typed[1].caller.function_name
-            parent = parent.parent
-        if not stack[-1]:
-            stack.pop(-1)
-        while parent:
-            if parent.typed[0] == _EventType.PyCall:
-                ex0 = parent.typed[1].caller  # torch._C._profiler._PyFrameState
-                ex1 = parent.typed[1].callsite  # torch._C._profiler._PyFrameState
-                if (
-                    _pycall_filter_fn(ex1.file_name)
-                    or (_pycall_filter_fn(ex0.file_name) and ex1.function_name == "backward")
-                ) and ex1.function_name not in {"__torch_dispatch__"}:
-                    stack.append(ex1.function_name)
-            parent = parent.parent
-        stack.reverse()
-        return "/".join(stack) or "unknown"
-
-    for ev_ in sorted(
-        traverse_dfs(prof.profiler.kineto_results.experimental_event_tree()), key=lambda ev: ev.start_time_ns
-    ):
-        # ev: torch._C._profiler._ProfilerEvent
-        _ev_visit(ev_)
-
-    assert not _allocs, f"Remaining allocs: {_allocs}"
-    assert not check_events, f"Remaining check events: {check_events}"
-
-
-def _pycall_filter_fn(filename: str) -> bool:
-    assert not filename.startswith("/")  # currently the case...
-    if os.path.basename(filename) == os.path.basename(__file__):
-        assert "/" not in filename  # currently the case...
-        return True
-    if filename.startswith("returnn/"):
-        return True
-    return False
-
-
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20240719.111324/tests/torch_utils.py

@@ -0,0 +1,230 @@
+"""
+Utilities for PyTorch tests
+"""
+
+from __future__ import annotations
+from typing import Optional, Any, Sequence, Tuple, Dict
+import torch
+
+
+def report_profile(
+    prof: torch.profiler.profiler,
+    check_events: Optional[Sequence[Tuple[str, Dict[str, Any]]]] = None,
+    *,
+    allow_remaining_allocs: bool = False,
+    _size_threshold: int = 100,
+) -> None:
+    """
+    Report profile.
+
+    Example usage::
+
+        from torch.profiler import profile, ProfilerActivity
+
+        with profile(
+            activities=[ProfilerActivity.CPU], profile_memory=True, with_stack=True, record_shapes=True
+        ) as prof:
+            # code to profile
+            ...
+
+        report_profile(prof)
+
+    :param prof: via torch.profiler.profile.
+    :param check_events: if given, will check that the report matches the given events.
+        Each entry is a tuple (event_name, event_opts).
+        The event_name can be "alloc", "dealloc", "torchop", "pycall".
+        The event_opts is a dict with the expected values.
+        You can use "*" as a wildcard in the event_opts.
+    :param allow_remaining_allocs: if True, will not raise an error if there are remaining allocations.
+    :param _size_threshold: internal
+    """
+    # 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.
+    # Either the Python stack is missing, or the memory information is incomplete,
+    # or the Python/TorchOp events are missing.
+    # The only complete information source seems to be prof.profiler.kineto_results.experimental_event_tree().
+
+    import fnmatch
+
+    # noinspection PyProtectedMember
+    from torch.profiler._utils import traverse_dfs
+    from torch._C._profiler import _EventType  # noqa
+
+    _printed_events = set()
+    _events_stack_depth = {}
+    _allocs = {}  # id -> dict with "size", "name"
+    check_events = list(check_events) if check_events is not None else None
+
+    def _ev_visit(ev):
+        # ev: torch._C._profiler._ProfilerEvent
+        if ev.typed[0] == _EventType.Allocation:
+            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_Allocation
+            # ex.id/ex.allocation_id/ex.ptr redundant?
+            if ex.allocation_id in _allocs:
+                ev_name = "dealloc"  # deallocation
+                assert _allocs[ex.allocation_id]["size"] == -ex.alloc_size
+                name = _allocs[ex.allocation_id]["name"]
+                del _allocs[ex.allocation_id]
+            else:
+                ev_name = "alloc"
+                assert ex.alloc_size > 0
+                assert ev.parent
+                name = _ctx(ev.parent)
+                _allocs[ex.allocation_id] = {"size": ex.alloc_size, "name": name}
+            opts = {"id": ex.allocation_id, "name": name, "size": ex.alloc_size, "total_alloc": ex.total_allocated}
+        elif ev.typed[0] == _EventType.TorchOp:
+            ev_name = "torchop"
+            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_TorchOp
+            opts = {"name": ex.name, "inputs": [_repr_tensor_metadata(i) for i in ex.inputs]}
+        elif ev.typed[0] == _EventType.PyCall:
+            ev_name = "pycall"
+            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_PyCall
+            ex0 = ex.caller  # torch._C._profiler._PyFrameState
+            ex1 = ex.callsite  # torch._C._profiler._PyFrameState
+            if _pycall_filter_fn(ex0.file_name) or _pycall_filter_fn(ex1.file_name):
+                opts = {
+                    "caller_loc": f"{ex0.file_name}:{ex0.line_number}",
+                    "caller_name": ex0.function_name,
+                    "callsite_name": ex1.function_name,
+                }
+            else:
+                return
+        elif ev.typed[0] == _EventType.PyCCall:
+            ev_name = "pyccall"
+            ex = ev.typed[1]  # torch._C._profiler._ExtraFields_PyCCall
+            ex0 = ex.caller  # torch._C._profiler._PyFrameState
+            if _pycall_filter_fn(ex0.file_name):
+                opts = {
+                    "caller_loc": f"{ex0.file_name}:{ex0.line_number}",
+                    "caller_name": ex0.function_name,
+                }
+            else:
+                return
+        else:
+            return
+
+        depth = 0
+        if ev.parent:
+            ev__ = ev.parent
+            while ev__:
+                if ev__ in _printed_events:
+                    depth = _events_stack_depth[ev__] + 1
+                    break
+                if ev__ in _events_stack_depth:
+                    depth = _events_stack_depth[ev__]
+                    break
+                ev__ = ev__.parent
+            ev__ = ev
+            while ev__:
+                if ev__ in _events_stack_depth:
+                    break
+                _events_stack_depth[ev__] = depth
+                ev__ = ev__.parent
+        else:
+            _events_stack_depth[ev] = 0
+        _printed_events.add(ev)
+        prefix = "  " * depth
+
+        if check_events is None:
+            print(f"{prefix}{ev_name} {opts}")
+            return
+
+        next_check = check_events[0] if check_events else None
+        if next_check:
+            next_check_name, next_check_opts = next_check
+            if ev_name == next_check_name:
+                for k, v in next_check_opts.items():
+                    if isinstance(v, str) and "*" in v:
+                        if not fnmatch.fnmatch(opts[k], v):
+                            mismatch = f"Pattern mismatch: {opts[k]} vs {v}"
+                            break
+                    elif k == "total_alloc":
+                        if abs(opts[k] - v) >= _size_threshold:
+                            mismatch = f"Size mismatch: {opts[k]} vs {v}"
+                            break
+                    elif opts[k] != v:
+                        mismatch = f"Value mismatch: {opts[k]} vs {v}"
+                        break
+                else:
+                    mismatch = None
+            else:
+                mismatch = f"Different event: {ev_name} vs {next_check_name}"
+        else:
+            mismatch = "No check event"
+
+        if ev_name in {"alloc", "dealloc"} and abs(opts["size"]) >= _size_threshold:
+            assert not mismatch, f"Event not matched: {ev_name} {opts} to {next_check}: {mismatch}"
+
+        if not mismatch:
+            print(f"{prefix}{ev_name} {opts} ✓")
+            check_events.pop(0)
+        else:
+            print(f"{prefix}({ev_name} {opts})")
+
+    def _ctx(ev) -> str:
+        stack = [None]
+        parent = ev
+        while parent and parent.typed[0] == _EventType.TorchOp:  # go to top torch op
+            stack[-1] = parent.typed[1].name
+            parent = parent.parent
+        if not stack[-1] and parent.typed[0] == _EventType.PyCCall:
+            stack[-1] = parent.typed[1].caller.function_name
+            parent = parent.parent
+        if not stack[-1]:
+            stack.pop(-1)
+        while parent:
+            if parent.typed[0] == _EventType.PyCall:
+                ex0 = parent.typed[1].caller  # torch._C._profiler._PyFrameState
+                ex1 = parent.typed[1].callsite  # torch._C._profiler._PyFrameState
+                if (
+                    _pycall_filter_fn(ex1.file_name)
+                    or (_pycall_filter_fn(ex0.file_name) and ex1.function_name == "backward")
+                ) and ex1.function_name not in {"__torch_dispatch__"}:
+                    stack.append(ex1.function_name)
+            parent = parent.parent
+        stack.reverse()
+        return "/".join(stack) or "unknown"
+
+    for ev_ in sorted(
+        traverse_dfs(prof.profiler.kineto_results.experimental_event_tree()), key=lambda ev: ev.start_time_ns
+    ):
+        # ev: torch._C._profiler._ProfilerEvent
+        _ev_visit(ev_)
+
+    if allow_remaining_allocs:
+        for alloc_id, alloc in _allocs.items():
+            print(f"Remaining alloc: {alloc_id} {alloc}")
+    else:
+        assert not _allocs, f"Remaining allocs: {_allocs}"
+    assert not check_events, f"Remaining check events: {check_events}"
+
+
+def _pycall_filter_fn(filename: str) -> bool:
+    assert not filename.startswith("/")  # currently the case...
+    if filename.startswith("test_"):
+        assert "/" not in filename  # currently the case...
+        return True
+    if filename.startswith("returnn/"):
+        return True
+    return False
+
+
+def _repr_tensor_metadata(x) -> Any:
+    """
+    :param x: torch._C._profiler._TensorMetadata or int
+    :return: obj which has a nice __repr__/__str__
+    """
+    if isinstance(x, (list, tuple)):
+        return type(x)(_repr_tensor_metadata(i) for i in x)
+    if not type(x).__module__.startswith("torch"):
+        return x
+    assert x is not None
+    return _ReprObj(f"<_T id={x.id} shape={x.sizes} dtype={x.dtype}>")
+
+
+class _ReprObj:
+    def __init__(self, obj: str):
+        self.obj = obj
+
+    def __repr__(self):
+        return self.obj

returnn-1.20240718.93955/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240718.093955'
-long_version = '1.20240718.093955+git.cf394a2'