Add FSDP2 activation_checkpointing_offload: offload checkpointed layer inputs to pinned CPU memory

src/accelerate/commands/launch.py

@@ -602,6 +602,12 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="Decides Whether (true|false) intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder. (useful only when `use_fsdp` flag is passed).",
     )
+    fsdp_args.add_argument(
+        "--fsdp_activation_checkpointing_offload",
+        default="false",
+        type=str,
+        help="Decides Whether (true|false) each checkpointed layer's input activation is offloaded to pinned CPU memory during the forward pass and restored on demand during the backward pass. Requires `fsdp_activation_checkpointing` and FSDP2. (useful only when `use_fsdp` flag is passed).",
+    )
 
     # megatron_lm args
     megatron_lm_args = parser.add_argument_group("Megatron-LM Arguments", "Arguments related to Megatron-LM.")

src/accelerate/utils/dataclasses.py

@@ -1792,6 +1792,14 @@ class FullyShardedDataParallelPlugin:
             "for reduced memory usage. Defaults to `False`"
         },
     )
+    activation_checkpointing_offload: bool = field(
+        default=None,
+        metadata={
+            "help": "Whether to offload each checkpointed layer's input activation to pinned CPU memory during the "
+            "forward pass and restore it on demand during the backward pass. Bounds activation memory at long "
+            "sequence lengths. Requires `activation_checkpointing=True` and `fsdp_version=2`. Defaults to `False`"
+        },
+    )
     cpu_ram_efficient_loading: bool = field(
         default=None,
         metadata={
@@ -1952,6 +1960,15 @@ def __post_init__(self):
                 str_to_bool(os.environ.get(env_prefix + "ACTIVATION_CHECKPOINTING", "False")) == 1
             )
 
+        if self.activation_checkpointing_offload is None:
+            self.activation_checkpointing_offload = (
+                str_to_bool(os.environ.get(env_prefix + "ACTIVATION_CHECKPOINTING_OFFLOAD", "False")) == 1
+            )
+        if self.activation_checkpointing_offload and not self.activation_checkpointing:
+            raise ValueError("`activation_checkpointing_offload=True` requires `activation_checkpointing=True`.")
+        if self.activation_checkpointing_offload and self.fsdp_version != 2:
+            raise ValueError("`activation_checkpointing_offload=True` requires `fsdp_version=2`.")
+
         if self.ignored_modules is None:
             self.ignored_modules = os.environ.get(env_prefix + "IGNORED_MODULES", None)
 

src/accelerate/utils/fsdp_utils.py

@@ -17,13 +17,15 @@
 import re
 import shutil
 import warnings
+import weakref
 from collections import defaultdict
 from collections.abc import Iterable
 from contextlib import nullcontext
 from pathlib import Path
 from typing import Callable, Union
 
 import torch
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ActivationWrapper
 
 from ..logging import get_logger
 from .constants import FSDP_MODEL_NAME, OPTIMIZER_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_NAME
@@ -687,6 +689,61 @@ def fsdp2_switch_optimizer_parameters(optimizer: torch.optim.Optimizer, mapping:
         )
 
 
+class _OffloadedCheckpointWrapper(ActivationWrapper):
+    """Non-reentrant activation checkpointing with the layer input offloaded to pinned CPU memory.
+
+    Non-reentrant `torch.utils.checkpoint.checkpoint` runs the forward grad-enabled (so tensors
+    captured out of the region — e.g. MoE router logits recorded for the aux loss — keep real
+    gradients) and stashes the layer input in the recompute closure, where it stays on GPU for
+    the whole forward+backward: `num_layers x seq_len x hidden` bytes, the dominant activation
+    cost at long sequence lengths. This wrapper frees that memory by swapping the input's
+    storage out to pinned host memory after the layer's forward
+    (`untyped_storage().resize_(0)`) and refilling it just-in-time when the recompute calls
+    back into the layer during backward.
+
+    Only the first positional tensor argument (`hidden_states`) is offloaded: it is uniquely
+    owned by its layer's closure, while other tensor arguments (rope embeddings, masks) are
+    shared across layers and must stay resident during the forward pass.
+    """
+
+    def __init__(self, module: torch.nn.Module):
+        # `ActivationWrapper` sets `_checkpoint_wrapped_module` and registers the state-dict hooks
+        # that hide this wrapper from parameter names, so checkpoints stay loadable by an
+        # unwrapped model.
+        super().__init__(module)
+        # (weakref to the gpu tensor, cpu copy, original storage size) per in-flight forward.
+        # The reference is weak so that an entry whose forward is never followed by a backward
+        # (an aborted step) is collected with the tensor instead of pinning its host copy: while
+        # the checkpoint frame is alive it holds the tensor, so the weakref stays valid.
+        self._stash = []
+
+    def _run(self, hidden_states, *args, **kwargs):
+        storage = hidden_states.untyped_storage()
+        if storage.size() == 0:
+            # Recompute path: refill the freed storage from the host copy, then drop the entry.
+            for i, (ref, cpu, size) in enumerate(self._stash):
+                if ref() is hidden_states:
+                    storage.resize_(size)
+                    hidden_states.copy_(cpu)
+                    del self._stash[i]
+                    break
+        return self._checkpoint_wrapped_module(hidden_states, *args, **kwargs)
+
+    def forward(self, hidden_states, *args, **kwargs):
+        output = torch.utils.checkpoint.checkpoint(
+            self._run, hidden_states, *args, use_reentrant=False, preserve_rng_state=False, **kwargs
+        )
+        if torch.is_grad_enabled() and hidden_states.is_cuda:
+            self._stash = [entry for entry in self._stash if entry[0]() is not None]
+            storage = hidden_states.untyped_storage()
+            size = storage.size()
+            cpu = torch.empty_like(hidden_states, device="cpu", pin_memory=True)
+            cpu.copy_(hidden_states, non_blocking=False)
+            storage.resize_(0)
+            self._stash.append((weakref.ref(hidden_states), cpu, size))
+        return output
+
+
 def fsdp2_apply_ac(accelerator, model: torch.nn.Module):
     """
     Applies the activation checkpointing to the model.
@@ -707,7 +764,11 @@ def fsdp2_apply_ac(accelerator, model: torch.nn.Module):
 
     for layer_name, layer in get_module_children_bottom_up(model, return_fqns=True)[:-1]:
         if auto_wrap_policy_func(layer):
-            model.set_submodule(layer_name, checkpoint_wrapper(layer, preserve_rng_state=False))
+            if accelerator.state.fsdp_plugin.activation_checkpointing_offload:
+                wrapped = _OffloadedCheckpointWrapper(layer)
+            else:
+                wrapped = checkpoint_wrapper(layer, preserve_rng_state=False)
+            model.set_submodule(layer_name, wrapped)
 
     return model
 
@@ -940,10 +1001,7 @@ def policy(module: torch.nn.Module) -> bool:
                 return False
             # Activation checkpointing (applied before sharding) wraps matched layers in
             # `CheckpointWrapper`; look through it so such layers still get their own FSDP group.
-            from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper
-
-            if isinstance(module, CheckpointWrapper):
-                module = module._checkpoint_wrapped_module
+            module = getattr(module, "_checkpoint_wrapped_module", module)
             return isinstance(module, tuple(transformer_cls_to_wrap))
 
     elif fn is size_based_auto_wrap_policy:

src/accelerate/utils/launch.py

@@ -328,6 +328,7 @@ def prepare_multi_gpu_env(args: argparse.Namespace) -> dict[str, str]:
         current_env["FSDP_CPU_RAM_EFFICIENT_LOADING"] = str(args.fsdp_cpu_ram_efficient_loading).lower()
         current_env["FSDP_SYNC_MODULE_STATES"] = str(args.fsdp_sync_module_states).lower()
         current_env["FSDP_ACTIVATION_CHECKPOINTING"] = str(args.fsdp_activation_checkpointing).lower()
+        current_env["FSDP_ACTIVATION_CHECKPOINTING_OFFLOAD"] = str(args.fsdp_activation_checkpointing_offload).lower()
         if getattr(args, "fsdp_ignored_modules", None) is not None:
             current_env["FSDP_IGNORED_MODULES"] = str(args.fsdp_ignored_modules)
 

tests/test_activation_checkpointing_offload.py

@@ -0,0 +1,96 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from torch import nn
+
+from accelerate.test_utils import require_non_cpu
+from accelerate.test_utils.testing import AccelerateTestCase
+from accelerate.utils.fsdp_utils import _OffloadedCheckpointWrapper
+
+
+class Block(nn.Module):
+    def __init__(self, dim=64):
+        super().__init__()
+        self.up = nn.Linear(dim, 4 * dim)
+        self.down = nn.Linear(4 * dim, dim)
+
+    def forward(self, hidden_states, scale=1.0):
+        return hidden_states + self.down(torch.nn.functional.silu(self.up(hidden_states))) * scale
+
+
+class CapturingBlock(Block):
+    """Mimics models that record an intermediate out of the layer (e.g. MoE router logits)."""
+
+    def forward(self, hidden_states, capture_list=None, **kwargs):
+        inner = self.up(hidden_states)
+        if capture_list is not None:
+            capture_list.append(inner)
+        return hidden_states + self.down(torch.nn.functional.silu(inner))
+
+
+@require_non_cpu
+class OffloadedCheckpointWrapperTester(AccelerateTestCase):
+    def _grads(self, model, x):
+        out = model(x).square().mean()
+        out.backward()
+        return {n: p.grad.clone() for n, p in model.named_parameters()}
+
+    def test_matches_unwrapped_gradients(self):
+        torch.manual_seed(0)
+        ref = nn.Sequential(*[Block() for _ in range(4)]).cuda()
+        wrapped = nn.Sequential(*[_OffloadedCheckpointWrapper(b) for b in ref])
+        x = torch.randn(2, 128, 64, device="cuda")
+
+        ref_grads = self._grads(ref, x)
+        for p in ref.parameters():
+            p.grad = None
+        off_grads = self._grads(wrapped, x)
+
+        for name, ref_grad in ref_grads.items():
+            wrapped_name = name.replace(".", "._checkpoint_wrapped_module.", 1)
+            torch.testing.assert_close(off_grads[wrapped_name], ref_grad)
+
+    def test_input_storage_freed_and_restored(self):
+        torch.manual_seed(0)
+        block = _OffloadedCheckpointWrapper(Block().cuda())
+        x = torch.randn(2, 128, 64, device="cuda", requires_grad=True)
+        hidden = x * 1.0  # non-leaf boundary tensor, like a previous layer's output
+        out = block(hidden)
+        assert hidden.untyped_storage().size() == 0  # offloaded after forward
+        out.square().mean().backward()
+        assert hidden.untyped_storage().size() != 0  # restored by the recompute
+        assert x.grad is not None
+
+    def test_state_dict_hides_the_wrapper(self):
+        # A checkpoint written from a wrapped model has to load into an unwrapped one, so the
+        # wrapper must not appear in parameter names.
+        ref = nn.Sequential(*[Block() for _ in range(2)])
+        wrapped = nn.Sequential(*[_OffloadedCheckpointWrapper(Block()) for _ in range(2)])
+        # Same keys as torch's own `checkpoint_wrapper`, so a checkpoint written from a wrapped
+        # model loads into an unwrapped one.
+        assert list(wrapped.state_dict().keys()) == list(ref.state_dict().keys())
+        ref.load_state_dict(wrapped.state_dict())
+
+    def test_captured_intermediate_keeps_gradient(self):
+        # Reentrant-style checkpointing would detach tensors captured out of the region;
+        # this wrapper must preserve their gradient path (grad-enabled forward).
+        torch.manual_seed(0)
+        block = _OffloadedCheckpointWrapper(CapturingBlock().cuda())
+        x = torch.randn(2, 128, 64, device="cuda", requires_grad=True)
+        captured = []
+        out = block(x * 1.0, capture_list=captured)
+        loss = out.square().mean() + captured[0].square().mean()
+        loss.backward()
+        assert x.grad is not None
+        assert captured[0].requires_grad