From c4ec07a56b9a9111d84f74ee62a95bd0c1f7947b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Wed, 19 Aug 2026 21:39:24 +0000 Subject: [PATCH] Add FSDP2 activation_checkpointing_offload: offload checkpointed layer inputs to pinned CPU memory --- src/accelerate/commands/launch.py | 6 ++ src/accelerate/utils/dataclasses.py | 17 ++++ src/accelerate/utils/fsdp_utils.py | 68 ++++++++++++- src/accelerate/utils/launch.py | 1 + .../test_activation_checkpointing_offload.py | 96 +++++++++++++++++++ 5 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 tests/test_activation_checkpointing_offload.py diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py index b768861b6e2..7974b93aaf1 100644 --- a/src/accelerate/commands/launch.py +++ b/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.") diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py index a5143ce7f06..b112d6f4816 100644 --- a/src/accelerate/utils/dataclasses.py +++ b/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) diff --git a/src/accelerate/utils/fsdp_utils.py b/src/accelerate/utils/fsdp_utils.py index bfdf42b3344..6deca92a196 100644 --- a/src/accelerate/utils/fsdp_utils.py +++ b/src/accelerate/utils/fsdp_utils.py @@ -17,6 +17,7 @@ import re import shutil import warnings +import weakref from collections import defaultdict from collections.abc import Iterable from contextlib import nullcontext @@ -24,6 +25,7 @@ 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: diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py index a1801c8c433..f7fce3f486b 100644 --- a/src/accelerate/utils/launch.py +++ b/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) diff --git a/tests/test_activation_checkpointing_offload.py b/tests/test_activation_checkpointing_offload.py new file mode 100644 index 00000000000..0dfcbe2914b --- /dev/null +++ b/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