torch-remat 0.2.0__tar.gz

torch_remat-0.2.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+(c) Meta Platforms, Inc. and affiliates.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list
+of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this
+list of conditions and the following disclaimer in the documentation
+and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may
+be used to endorse or promote products derived from this software without specific
+prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
+EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
+SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
+TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+DAMAGE.

torch_remat-0.2.0/PKG-INFO

@@ -0,0 +1,196 @@
+Metadata-Version: 2.4
+Name: torch_remat
+Version: 0.2.0
+Summary: Explicit activation rematerialization helpers for PyTorch
+Author: Meta Platforms, Inc.
+License: BSD-3-Clause
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.10.0
+Dynamic: license-file
+
+# torch_remat
+
+`torch_remat` is a modern activation checkpointing implementation for PyTorch.
+It takes the classic recipe of non-reentrant `torch.utils.checkpoint` (recompute
+everything before backwards) and enhances it with the ability to mark regions
+inside the recompute region to save their activations for backward, so they
+don't have to be recomputed.
+
+Compared with selective activation checkpointing (SAC), `torch_remat` allows you
+to specify save/recompute decisions on regions of source code, instead of writing
+a generic policy function that operates on a per-ATen operator basis.  This
+allows for finer-grained recompute policies (e.g., you can easily express that
+one matmul should be saved while another should be recomputed).  It's also
+easier to use `torch_remat` with code that has many custom kernels (you don't
+have to custom op'ify them), and the fact that `torch_remat` doesn't use a
+`TorchDispatchMode` means CPU overhead is lower.
+
+## Quick start
+
+Put one checkpoint around the unit you want to replay, usually a transformer
+block:
+
+```python
+import torch_remat as remat
+
+# torch.utils.checkpoint.checkpoint compatible API:
+# remat.checkpoint(**remat_kwargs)(func)(*args, **kwargs)
+output = remat.checkpoint(region_name="layers.0")(block)(hidden_states)
+```
+
+(`remat.checkpoint(block)(hidden_states)` is intentionally not supported,
+because this phrasing of the API is ambiguous with `torch.utils.checkpoint`.)
+
+Inside the block's `forward` method, annotate calls whose backward activations
+should be saved with `remat.region(..., recompute=False)`, and annotate
+outputs of save regions which will be needed for recompute with
+`remat.recompute_needs_tensor` (you can also omit these and `torch_remat` will
+tell you which ones you need to mark):
+
+```python
+def forward(self, hidden_states):
+    x = self.attention_norm(hidden_states)
+    attn = remat.region(
+        self.attention,
+        "attention",
+        recompute=False,
+    )(x)
+    remat.recompute_needs_tensor(attn)
+    hidden_states = hidden_states + attn
+
+    x = self.ffn_norm(hidden_states)
+    moe = remat.region(
+        self.moe,
+        "moe",
+        recompute=False,
+    )(x)
+    remat.recompute_needs_tensor(moe)
+    return hidden_states + moe
+```
+
+The interaction between recompute and save regions is somewhat subtle;
+check [Mental model](docs/mental_model.md) for more details.
+
+## State and side effects
+
+Code that is replayed must behave consistently with the original forward, even
+if saved region bodies are skipped during replay.  In particular, if you rely
+on mutable state in forwards (e.g., for RNG), you need to ensure you can
+snapshot and restore this state.  Code `remat.RecomputeStateHook` and pass it
+with `recompute_state_hooks=`. The hook restores state at checkpoint entry and
+at every non-recomputed function.  Here is an example that takes care of
+setting both a custom user RNG counter as well as standard PyTorch RNG state.
+
+```python
+from contextvars import ContextVar
+
+rng_counter = ContextVar("rng_counter", default=0)
+
+
+class CudaRNGStateHook:
+    def __init__(self, device):
+        self.device = device
+
+    def snapshot(self):
+        return torch.cuda.get_rng_state(self.device), rng_counter.get()
+
+    def restore(self, state):
+        cuda_rng_state, counter = state
+        torch.cuda.set_rng_state(cuda_rng_state, self.device)
+        rng_counter.set(counter)
+
+
+rng_hook = CudaRNGStateHook(hidden_states.device)
+output = remat.checkpoint(
+    region_name="layers.0",
+    recompute_state_hooks=(rng_hook,),
+)(block)(hidden_states)
+```
+
+We don't provide a "stock" save/restore hook; in particular, `torch_remat`
+doesn't support the `preserve_rng_state=True` kwarg that
+`torch.utils.checkpoint` supports.  The primary reason for this is that
+`preserve_rng_state` is documented to also save CPU RNG state, but in modern
+PyTorch code this is unnecessary (RNG should be sampled on-device) and
+expensive (a 5KB allocation is needed to snapshot the CPU MT19937 state)--and
+unlike `torch.utils.checkpoint`, we will repeatedly save/load RNG state many
+times per a `remat.checkpoint`.
+
+Separately, you can check whether other code is being replayed with
+`remat.is_recomputing()`, for example to suppress forward-only logging or
+metrics:
+
+```python
+if not remat.is_recomputing():
+    record_metric(value)
+```
+
+TODO: We should offer a simple way of checking that the replay is bitwise
+equivalent to the original.
+
+## Diagnostics
+
+`torch_remat` also comes with a number tools for understanding the recompute/save behavior
+and memory usage of your program.  Here are some things you can do:
+
+**Trace the configured region hierarchy to see what is being saved/recomputed:**
+
+```python
+with remat.collect_trace() as trace:
+    output = model(inputs)
+print(trace.format())
+```
+
+```text
+torch_remat trace
+scope [test_flag]
+  sin: save
+  cos: recompute
+```
+
+**Inspect retained activations from inside a checkpoint forward:**
+
+```python
+if not remat.is_recomputing():
+    remat.print_current_memory_report()
+```
+
+```text
+layers.0: 28 B resident in 2 storage(s)
+layers.0::attn.softmax: 28 B
+  12 B  lse    (3,)  float32
+  16 B  probs  (4,)  float32
+```
+
+**Inspect all live checkpoint regions plus saves reachable from a loss after
+the full forward:**
+
+```python
+remat.print_saved_tensors_report(loss)
+```
+
+```text
+saved for backward: 240 B resident -- 2 region(s) 192 B, outside regions 48 B
+
+regions:
+  96 B  x2  layer.0-1  (2 storages each)
+
+outside regions: 48 B in 1 storage
+       48 B  TanhBackward0 (x1)
+
+region detail:
+[x2: layer.0-1]
+layer.0: 96 B resident in 2 storage(s)
+layer.0::sq: 96 B
+  48 B  y (output at idx 0)  (3, 4)  float32
+  48 B  gf                   (3, 4)  float32
+```
+
+Allocation-site annotations and an attachable CUDA OOM observer are also
+provided. See the [Diagnostics API](docs/api/diagnostics.md) for more details.
+
+## License
+
+BSD 3-Clause License. See [LICENSE](LICENSE).

torch_remat-0.2.0/README.md

@@ -0,0 +1,184 @@
+# torch_remat
+
+`torch_remat` is a modern activation checkpointing implementation for PyTorch.
+It takes the classic recipe of non-reentrant `torch.utils.checkpoint` (recompute
+everything before backwards) and enhances it with the ability to mark regions
+inside the recompute region to save their activations for backward, so they
+don't have to be recomputed.
+
+Compared with selective activation checkpointing (SAC), `torch_remat` allows you
+to specify save/recompute decisions on regions of source code, instead of writing
+a generic policy function that operates on a per-ATen operator basis.  This
+allows for finer-grained recompute policies (e.g., you can easily express that
+one matmul should be saved while another should be recomputed).  It's also
+easier to use `torch_remat` with code that has many custom kernels (you don't
+have to custom op'ify them), and the fact that `torch_remat` doesn't use a
+`TorchDispatchMode` means CPU overhead is lower.
+
+## Quick start
+
+Put one checkpoint around the unit you want to replay, usually a transformer
+block:
+
+```python
+import torch_remat as remat
+
+# torch.utils.checkpoint.checkpoint compatible API:
+# remat.checkpoint(**remat_kwargs)(func)(*args, **kwargs)
+output = remat.checkpoint(region_name="layers.0")(block)(hidden_states)
+```
+
+(`remat.checkpoint(block)(hidden_states)` is intentionally not supported,
+because this phrasing of the API is ambiguous with `torch.utils.checkpoint`.)
+
+Inside the block's `forward` method, annotate calls whose backward activations
+should be saved with `remat.region(..., recompute=False)`, and annotate
+outputs of save regions which will be needed for recompute with
+`remat.recompute_needs_tensor` (you can also omit these and `torch_remat` will
+tell you which ones you need to mark):
+
+```python
+def forward(self, hidden_states):
+    x = self.attention_norm(hidden_states)
+    attn = remat.region(
+        self.attention,
+        "attention",
+        recompute=False,
+    )(x)
+    remat.recompute_needs_tensor(attn)
+    hidden_states = hidden_states + attn
+
+    x = self.ffn_norm(hidden_states)
+    moe = remat.region(
+        self.moe,
+        "moe",
+        recompute=False,
+    )(x)
+    remat.recompute_needs_tensor(moe)
+    return hidden_states + moe
+```
+
+The interaction between recompute and save regions is somewhat subtle;
+check [Mental model](docs/mental_model.md) for more details.
+
+## State and side effects
+
+Code that is replayed must behave consistently with the original forward, even
+if saved region bodies are skipped during replay.  In particular, if you rely
+on mutable state in forwards (e.g., for RNG), you need to ensure you can
+snapshot and restore this state.  Code `remat.RecomputeStateHook` and pass it
+with `recompute_state_hooks=`. The hook restores state at checkpoint entry and
+at every non-recomputed function.  Here is an example that takes care of
+setting both a custom user RNG counter as well as standard PyTorch RNG state.
+
+```python
+from contextvars import ContextVar
+
+rng_counter = ContextVar("rng_counter", default=0)
+
+
+class CudaRNGStateHook:
+    def __init__(self, device):
+        self.device = device
+
+    def snapshot(self):
+        return torch.cuda.get_rng_state(self.device), rng_counter.get()
+
+    def restore(self, state):
+        cuda_rng_state, counter = state
+        torch.cuda.set_rng_state(cuda_rng_state, self.device)
+        rng_counter.set(counter)
+
+
+rng_hook = CudaRNGStateHook(hidden_states.device)
+output = remat.checkpoint(
+    region_name="layers.0",
+    recompute_state_hooks=(rng_hook,),
+)(block)(hidden_states)
+```
+
+We don't provide a "stock" save/restore hook; in particular, `torch_remat`
+doesn't support the `preserve_rng_state=True` kwarg that
+`torch.utils.checkpoint` supports.  The primary reason for this is that
+`preserve_rng_state` is documented to also save CPU RNG state, but in modern
+PyTorch code this is unnecessary (RNG should be sampled on-device) and
+expensive (a 5KB allocation is needed to snapshot the CPU MT19937 state)--and
+unlike `torch.utils.checkpoint`, we will repeatedly save/load RNG state many
+times per a `remat.checkpoint`.
+
+Separately, you can check whether other code is being replayed with
+`remat.is_recomputing()`, for example to suppress forward-only logging or
+metrics:
+
+```python
+if not remat.is_recomputing():
+    record_metric(value)
+```
+
+TODO: We should offer a simple way of checking that the replay is bitwise
+equivalent to the original.
+
+## Diagnostics
+
+`torch_remat` also comes with a number tools for understanding the recompute/save behavior
+and memory usage of your program.  Here are some things you can do:
+
+**Trace the configured region hierarchy to see what is being saved/recomputed:**
+
+```python
+with remat.collect_trace() as trace:
+    output = model(inputs)
+print(trace.format())
+```
+
+```text
+torch_remat trace
+scope [test_flag]
+  sin: save
+  cos: recompute
+```
+
+**Inspect retained activations from inside a checkpoint forward:**
+
+```python
+if not remat.is_recomputing():
+    remat.print_current_memory_report()
+```
+
+```text
+layers.0: 28 B resident in 2 storage(s)
+layers.0::attn.softmax: 28 B
+  12 B  lse    (3,)  float32
+  16 B  probs  (4,)  float32
+```
+
+**Inspect all live checkpoint regions plus saves reachable from a loss after
+the full forward:**
+
+```python
+remat.print_saved_tensors_report(loss)
+```
+
+```text
+saved for backward: 240 B resident -- 2 region(s) 192 B, outside regions 48 B
+
+regions:
+  96 B  x2  layer.0-1  (2 storages each)
+
+outside regions: 48 B in 1 storage
+       48 B  TanhBackward0 (x1)
+
+region detail:
+[x2: layer.0-1]
+layer.0: 96 B resident in 2 storage(s)
+layer.0::sq: 96 B
+  48 B  y (output at idx 0)  (3, 4)  float32
+  48 B  gf                   (3, 4)  float32
+```
+
+Allocation-site annotations and an attachable CUDA OOM observer are also
+provided. See the [Diagnostics API](docs/api/diagnostics.md) for more details.
+
+## License
+
+BSD 3-Clause License. See [LICENSE](LICENSE).

torch_remat-0.2.0/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch_remat"
+version = "0.2.0"
+description = "Explicit activation rematerialization helpers for PyTorch"
+readme = "README.md"
+license = {text = "BSD-3-Clause"}
+requires-python = ">=3.10"
+authors = [
+    {name = "Meta Platforms, Inc."},
+]
+dependencies = [
+    "torch>=2.10.0",
+]
+
+[tool.setuptools.packages.find]
+include = ["torch_remat*"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["*_test.py"]

torch_remat-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_remat-0.2.0/torch_remat/__init__.py

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from torch_remat._annotate import (  # noqa: F401
+    clear_memory_snapshot_cache,
+    memory_snapshot_annotate,
+)
+from torch_remat._api import (  # noqa: F401
+    checkpoint,
+    current_saved_tensor_info,
+    is_placeholder,
+    recompute_needs_tensor,
+    region,
+    save_for_backward,
+    saved_tensors_hooks,
+)
+from torch_remat._oom import (  # noqa: F401
+    discover_autograd_roots,
+    format_oom_saved_tensors_report,
+    oom_observer,
+)
+from torch_remat._region import is_recomputing, RecomputeStateHook  # noqa: F401
+from torch_remat._reporting import (  # noqa: F401
+    format_current_memory_report,
+    print_current_memory_report,
+)
+from torch_remat._saved_report import (  # noqa: F401
+    format_saved_tensors_report,
+    print_saved_tensors_report,
+)
+from torch_remat._trace import collect_trace, trace_scope  # noqa: F401
+from torch_remat._types import SavedTensorInfo, SavedTensorKind  # noqa: F401

torch_remat-0.2.0/torch_remat/_annotate.py

@@ -0,0 +1,111 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+"""A default :data:`~torch_remat._reporting.Annotate` backed by the CUDA memory snapshot.
+
+The saved-tensors report can tag each storage row with a caller-supplied label. By far the
+most common label is *where the tensor was allocated*, which the CUDA caching allocator
+already records (when memory-history recording is on). This module turns that snapshot into a
+ready-made annotator -- the "which tensor is this" that a positional ``saved.<i>`` row cannot
+give -- so callers don't reinvent the storage-pointer -> allocation-site mapping.
+
+The snapshot walk is expensive, so the fetch is memoized (:func:`clear_memory_snapshot_cache`
+resets it) and the built annotator indexes the snapshot once, reusing it across every
+per-tensor lookup -- a whole report costs a single snapshot fetch.
+"""
+
+from __future__ import annotations
+
+import bisect
+from typing import Any, Mapping
+
+import torch
+from torch_remat._reporting import Annotate
+
+_snapshot_cache: Mapping[str, Any] | None = None
+
+
+def _memory_snapshot(*, refresh: bool = False) -> Mapping[str, Any]:
+    """Return the CUDA allocator memory snapshot, memoized after the first fetch.
+
+    ``torch.cuda.memory._snapshot()`` walks the whole allocator and is expensive, so the
+    result is cached. Pass ``refresh=True`` (or call :func:`clear_memory_snapshot_cache`) to
+    re-fetch after more allocation. The snapshot carries allocation backtraces only while
+    memory-history recording is enabled.
+    """
+    global _snapshot_cache
+    if refresh or _snapshot_cache is None:
+        _snapshot_cache = torch.cuda.memory._snapshot()
+    return _snapshot_cache
+
+
+def clear_memory_snapshot_cache() -> None:
+    """Drop the memoized snapshot so the next annotator re-fetches (see :func:`_memory_snapshot`)."""
+    global _snapshot_cache
+    _snapshot_cache = None
+
+
+def _alloc_site_from_frames(frames: list[dict[str, Any]] | None) -> str | None:
+    """Pick a short ``file:line name`` label from an allocation backtrace (innermost first).
+
+    Prefer the innermost frame outside library internals (torch and torch_remat), which is
+    usually the model/kernel code that produced the tensor; fall back to the innermost frame
+    of any kind.
+    """
+    for frame in frames or []:
+        filename = frame.get("filename", "")
+        if "/torch/" not in filename and "/torch_remat/" not in filename:
+            base = filename.rsplit("/", 1)[-1]
+            return f"{base}:{frame.get('line')} {frame.get('name', '')}".rstrip()
+    if frames:
+        frame = frames[0]
+        return f"{frame.get('filename', '?').rsplit('/', 1)[-1]}:{frame.get('line')}"
+    return None
+
+
+def memory_snapshot_annotate(snapshot: Mapping[str, Any] | None = None) -> Annotate:
+    """Return an :data:`Annotate` mapping a saved tensor's storage to its allocation site.
+
+    Each live allocator block carries its address and allocation backtrace, so a saved
+    tensor's storage pointer (which falls inside exactly one live block) resolves to the code
+    that allocated it. Pass the result as ``annotate=`` to
+    :func:`torch_remat.format_saved_tensors_report`.
+
+    Args:
+        snapshot: A CUDA memory snapshot (``torch.cuda.memory._snapshot()``); defaults to the
+            cached :func:`_memory_snapshot`. The block index is built once here and reused by
+            every per-tensor lookup, so a whole report costs a single snapshot fetch.
+    """
+
+    if snapshot is None:
+        snapshot = _memory_snapshot()
+
+    blocks: list[tuple[int, int, str]] = []  # (address, size, site), sorted by address
+    for segment in snapshot.get("segments", []):
+        for block in segment.get("blocks", []):
+            if block.get("state") != "active_allocated":
+                continue
+            site = _alloc_site_from_frames(block.get("frames"))
+            if site is not None:
+                blocks.append((block["address"], block["size"], site))
+    blocks.sort()
+    addrs = [address for address, _, _ in blocks]
+
+    def annotate(tensor: torch.Tensor) -> str | None:
+        try:
+            ptr = tensor.untyped_storage().data_ptr()
+        except Exception:
+            return None
+        index = bisect.bisect_right(addrs, ptr) - 1
+        if 0 <= index < len(blocks):
+            address, size, site = blocks[index]
+            if address <= ptr < address + size:
+                return site
+        return None
+
+    return annotate