Is this page helpful?

Fine-Grained Activation Offloading#

Contributed in collaboration with RedNote.

Fine-grained activation offloading reduces GPU memory by asynchronously transferring activations to CPU at the granularity of individual submodules within a transformer layer. Unlike layer-level offloading, it allows precise control over which activations to offload, enabling a tradeoff between memory savings and PCIe bandwidth overhead.

Supported offloading modules are "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp", and "shortcut_post_norm". They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device. fused_group_mlp requires --use-transformer-engine-op-fuser and offloads the whole fused grouped MLP, so it cannot be combined with expert_fc1 or moe_act. shortcut_post_norm requires --moe-shortcut-connection.

User Guide#

Basic Usage#

# Enable fine-grained activation offloading
--fine-grained-activation-offloading

# Modules whose inputs are offloaded (refer to your training script for list or delimiter syntax).
# Choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp", "shortcut_post_norm".
--offload-modules core_attn attn_proj expert_fc1

Offloadable Modules#

Each module offloads its input activation to CPU during forward and reloads it before backward:

Module

Description

Notes

attn_norm

Input layernorm of attention

Skipped if using IdentityOp

qkv_linear

QKV linear projection

core_attn

Core attention (softmax + matmul)

attn_proj

Output projection of attention

Must be used together with core_attn

mlp_norm

Pre-MLP layernorm

Skipped if using IdentityOp

expert_fc1

First FC layer in MoE experts

MoE models only

moe_act

Activation function in MoE experts

MoE models only

fused_group_mlp

Whole fused grouped MLP

Requires --use-transformer-engine-op-fuser; cannot be combined with expert_fc1 or moe_act

shortcut_post_norm

Shortcut post-combine normalization

Requires --moe-shortcut-connection

When selective residual_stream replay is active, its shared checkpoint manager owns the connected ShortcutMoE post-norm. In that case shortcut_post_norm offloading is skipped for the replay-owned forward so that replay and offloading do not release the same activation.

Tuning Parameters#

# Minimum tensor size (in elements) to offload. Smaller tensors are skipped.
# Default: 1048576 (1M elements)
--min-offloaded-tensor-size 1048576

# Fraction of activations to offload, range [0, 1]. Default: 1.0
# Useful for partial offloading when PCIe bandwidth is a bottleneck.
--activation-offload-fraction 0.8

# Reduce offload amount on higher PP ranks (in bytes). Default: 0
# Higher PP ranks have fewer microbatches in flight, so offloading less
# reduces overhead without increasing peak memory.
--delta-offload-bytes-across-pp-ranks 1073741824

# Optional: cap inflight D2H offloads per offload group to N (omit or None in most setups).
# Required as a non-None non-negative integer when fine-grained activation offloading is used with
# local full-iteration CUDA graphs (full_iteration in cuda_graph_scope); see prose below.
--fine-grained-offloading-max-inflight-offloads <N>

# Optional: pack all offloaded activations on a rank into one fixed pinned buffer of this
# many GiB instead of allocating pinned host memory per offload. Default: 0 (disabled).
--fine-grained-offloading-buffer-size-gib <GiB>

TransformerConfig.fine_grained_offloading_max_inflight_offloads caps, per offload group (for example moe_act, qkv_linear), how many D2H copies may be in flight before a main-stream wait_event. 0 waits after each offload; larger values allow more overlap; None skips these joins.

With full-iteration CUDA graphs (local graph impl, full_iteration in cuda_graph_scope) and fine-grained activation offloading enabled, set it to a non-None integer: that path does not rely on record_stream, so explicit joins are required.

Activation Offload Fraction#

--activation-offload-fraction (TransformerConfig.activation_offload_fraction) is a fraction over eligible offload groups, not a byte fraction and not a selector for which module names are enabled. It is used together with --offload-modules: all module names listed in --offload-modules still register their offload groups, and the steady-state fraction is applied once across the combined eligible groups from all configured modules. Warmup uses the per-module policy described in Warmup and Adaptive Offloading.

The manager keeps the first N% of eligible groups in forward execution order and leaves the later groups on GPU. For example, with --offload-modules core_attn attn_proj expert_fc1 --activation-offload-fraction 0.5, the eligible core_attn, attn_proj, and expert_fc1 groups are considered together in execution order, and the first 50% of that combined group list are offloaded. The fraction does not mean “offload 50% of the activation bytes” and does not mean “offload only the first 50% of the module names”.

The fraction is applied after other eligibility filters such as min_offloaded_tensor_size, the last-group margin used to avoid backward reload stalls, and delta_offload_bytes_across_pp_ranks. Therefore N% is computed over the remaining eligible groups from all configured offload modules after those filters.

Fixed Pinned Buffer#

By default, fixed-shape groups reuse pinned tensors cached by shape and dtype, while the variable-shape MoE groups (expert_fc1, moe_act, fused_group_mlp) allocate from the PyTorch pinned host allocator. That allocator rounds each request up to a power of two and keeps freed blocks for reuse, so as routing varies its cache keeps growing and new sizes pin more host memory during the step. --fine-grained-offloading-buffer-size-gib (TransformerConfig.fine_grained_offloading_buffer_size_gib) instead packs all backups on a rank into one pinned buffer:

  • The buffer is pinned with cudaHostAlloc from cuda-python at the requested size, rounded up only to whole pages. The size must be a positive multiple of 256 bytes, so the GiB value must be a multiple of 2-22, for example 152 or 152.5.

  • It is allocated once per rank, on the first forward pass with offloading, and never grows. Running out raises an error that reports the request and the capacity.

  • Space is reused only after every backup in the buffer has been reloaded. Size the buffer for everything offloaded between such points, with each tensor aligned to 256 bytes: one microbatch when each backward finishes before the next forward starts (PP=1, or the last pipeline stage without virtual pipeline stages, and no --overlap-moe-expert-parallel-comm), otherwise everything the rank offloads in an iteration, across all microbatches and model chunks.

  • Captured addresses stay reserved for the lifetime of the buffer, including after reset() or clear(). Each capture gets a separate range; it can reuse that range within the same capture after all its backups are reloaded. Eager offloads use the remaining tail, so size mixed-mode buffers for both graph reservations and eager backups. get_pool_status() reports graph_reserved_bytes separately from live bytes. To reclaim graph reservations, destroy the graphs and their buffer together.

  • The size bounds this buffer only, not other host allocations.

CUDA Graph Integration#

Fine-grained offloading is compatible with CUDA graphs. When CUDA graph is enabled, the following constraints apply:

  • attn_norm and mlp_norm cannot be offloaded (they cross CUDA graph boundaries).

  • cuda_graph_scope must include attn and moe_router.

  • cuda_graph_impl must be transformer_engine.

  • Requires torch >= 2.9.0 and transformer_engine >= 2.14.0.

# Optional: defer D2H enqueue for offloads *outside* cuda_graph_scope (MoE experts; see below)
--delay-offload-until-cuda-graph

--delay-offload-until-cuda-graph (TransformerConfig.delay_offload_until_cuda_graph)

Inside vs outside cuda_graph_scope. Offload boundaries that lie inside the captured cuda_graph_scope (for example qkv_linear, core_attn, and attn_proj when attn is in scope) are part of CUDA graph capture and replay. Their offload-related work is replayed with the graph rather than re-driven from Python each step, so they do not incur the same per-step CPU launch overhead as a purely eager path.

Boundaries that run outside the captured region still execute as normal eager PyTorch each forward—for the recommended MoE setup, that includes expert compute after a graphed moe_router (e.g. offloading expert_fc1 / moe_act). For those groups, each group_offload would otherwise submit D2H work from the host as soon as the forward hits the commit point.

What this flag does. It only affects offload commits that are explicitly wired with delayed group commit (currently the MoE expert path: expert_fc1, moe_act). Around each layer’s TransformerEngine CUDA graph replay, the offload manager enters replay mode; delayed commits enqueue (callback, group name, forced tensors) instead of launching D2H immediately, then flush_delayed_groups runs after that graph replay returns and issues the queued D2H copies in forward order, without changing the offload/reload semantics.

When this actually buys time (EP A2A after replay). The benefit assumes a real CPU/GPU synchronization gap right after graph replay—in the usual MoE training layout, expert parallel (EP) all-to-all and related dispatch follows the graphed moe_router region. That A2A path typically needs the host to coordinate collectives and to sync with the GPU (e.g. wait for graph work to finish or for communication staging), so the CPU is not fully overlapped with useful launch work during that interval. Scheduling flush_delayed_groups immediately after cudaGraphLaunch returns uses that window to issue D2H copies from the host: the enqueue cost is largely hidden in slack that EP A2A would already incur. If there were no such post-replay sync (or expert work were fully captured inside the graph with no host-visible gap), deferring commits would not provide the same “free” host time.

Behavioral notes

  • Does not replace or “delay” attention-side offloads inside the graphed attn region; those are not on the delayed path in the implementation.

  • Warmup and non-replay forwards still commit delayed-eligible groups immediately (no replay-mode deferral).

  • Must be used together with fine-grained activation offloading and CUDA graph under the same rules as this section (TE cuda_graph_impl, scope including attn and moe_router, etc.).

  • Stream ordering between the graph compute path and d2h_stream still uses the existing events (forward_record / backward_record); this option only changes when eligible D2H work is submitted from the host.

Combining with Fine-Grained Recomputation#

Offloading and recomputation are complementary:

  • Use recomputation for lightweight modules (e.g., layernorm, activation functions) with negligible compute overhead.

  • Use offloading for heavy modules (e.g., core_attn, expert_fc1) where recomputation would be too costly.

--recompute-granularity selective
--recompute-modules layernorm moe_act
--fine-grained-activation-offloading
--offload-modules core_attn attn_proj expert_fc1

Fine-grained Activation Offloading and Fine-grained Recomputation

Compatibility#

Feature

Supported

PP / Interleaved PP / PP=1

Yes

Fine-grained recomputation

Yes

FP8 training

Yes

MTP (Multi-Token Prediction)

Yes

Mixed dense & MoE layers

Yes

A2A overlap (EP)

Yes

CUDA Graph (TE impl)

Yes

How It Works#

Architecture Overview#

The implementation consists of three layers:

  1. PipelineOffloadManager (singleton): Global coordinator that manages CUDA streams, CPU tensor pools, and chunk lifecycle across pipeline stages.

  2. ChunkOffloadHandler: Per-microbatch handler that tracks tensor groups, executes D2H/H2D transfers, and decides which groups to actually offload.

  3. FineGrainedActivationOffloadingInterface: Lightweight interface used by transformer modules (attention, MoE, etc.) to mark offload boundaries.

Offload/Reload Flow#

Forward pass (Layer N):                    Backward pass (Layer N):
┌─────────────────────┐                    ┌───────────────────────┐
│ group_start(input)  │─── register ──►    │                       │
│                     │    tensor group    │ group_commit_backward │
│ module.forward()    │                    │   wait H2D complete   │
│                     │                    │   pop tensors from    │
│ group_offload(out)  │─── D2H async ──►   │   CPU → GPU           │
│   on d2h_stream     │    to pinned CPU   │   on h2d_stream       │
└─────────────────────┘                    └───────────────────────┘

  1. group_start: Registers a new tensor group and hooks into saved_tensors_hooks to intercept save_for_backward.

  2. Forward execution: All tensors saved by autograd within the group are captured.

  3. group_offload: Triggers asynchronous D2H copy on a dedicated CUDA stream (d2h_stream), optionally releases GPU storage of input tensors.

  4. Backward: Before the group’s backward, tensors are reloaded from CPU to GPU on h2d_stream, and the compute stream waits for the transfer to complete.

Committed forward groups occupy slots in the backward reload queue, including groups kept on GPU by the fraction or PP-rank delta policy. These resident groups are consumed as no-op slots without recording a reload event. The group-start autograd marker retains the corresponding forward group so backward can consume its current no-op slot before preloading the next group.

Groups explicitly reserved by the post-warmup margin policy are excluded from this queue. Their backward markers still trigger prefetches, providing compute time to reload earlier groups. This also handles delayed norm commits: although attn_norm starts before core_attn, it commits after the attention residual addition. Enqueuing both resident margin groups in commit order would spend the first two prefetch opportunities consuming their slots instead of reloading earlier activations. Warmup has no reserved margin yet, so all warmup groups retain their reload slots.

Warmup and Adaptive Offloading#

The first training iteration serves as a warmup phase where the manager records tensor groups, their candidate transfer sizes, and the execution order without requiring a CPU copy for every group. Because the final group count is not yet known, each chunk uses a separate online quota for each module name. Its kth occurrence is selected when ceil(k * activation_offload_fraction) increases. For example, fraction 0.5 selects occurrences 1, 3, 5, … of each configured module. Every group advances its module’s counter, including empty groups and groups with no eligible tensors; those groups perform no copy and do not shift the selection of later occurrences. This keeps corresponding core_attn and attn_proj groups aligned when both occur once per layer, so projection offload does not release an output still held on GPU for attention backward. Fraction 0 performs no copies; fraction 1 offloads every eligible warmup group. Each module rounds up independently, favoring GPU memory headroom; the sum can exceed a single quota rounded over all modules together.

This temporary warmup selection does not change the groups’ eligibility for later iterations. It also does not add explicit cache eviction. After warmup, post_warmup_callback computes the normal policy independently, preserving the preference for earlier forward groups:

  1. Reserve margin: The last group of each module name is kept on GPU and excluded from the reload queue, leaving its backward marker available to trigger prefetches.

  2. Apply PP rank delta: Higher PP ranks offload fewer bytes (controlled by delta_offload_bytes_across_pp_ranks).

  3. Apply fraction: Only the first N% of the remaining eligible groups are offloaded across all configured modules (controlled by activation_offload_fraction).

  4. Print summary table: An ASCII table of per-rank offload bytes is printed for debugging.

CPU Tensor Pool#

A ‘OffloadTensorPool(on CPU with pinned memory) caches allocated tensors by(shape, dtype). This avoids repeated cudaMallocHost/cudaFreeHostcalls and reduces D2H latency after the first iteration. With–fine-grained-offloading-buffer-size-gib`, the pool instead hands out views of one fixed pinned buffer (see Fixed Pinned Buffer).

Non-contiguous views covering at least half their storage are transferred as full storage and restored with their original shape, stride, and storage offset. Smaller views use a contiguous copy.

core_attn offloading does not support packed Q/K/V inputs with non-contiguous views sharing storage. Attention checks the actual inputs after Q/K normalization and rotary embedding and raises ValueError during training when the offload fraction is positive. Gathering these views would change the layout retained by fused attention backward. Use unpacked attention inputs or remove core_attn and its dependent attn_proj from offload_modules. This restriction does not apply when core attention offloading is disabled or the fraction is zero; it is not a blanket restriction on models without Q/K normalization.

TODO: Support packed Q/K/V inputs with core_attn fine-grained activation offloading (FGAO). Preserve the layout required by attention backward while transferring each shared QKV storage only once in each direction. Add fused/FlashAttention forward and gradient parity tests, including partial offload and CUDA graphs, before removing the current runtime restriction.

The determinism suite covers supported Q/K-normalized and RoPE inputs with partial and full offload. It compares logits and parameter gradients bit-for-bit against a resident baseline and across warmup/steady-state replays under side-stream contention.

CUDA Graph Support#

When offloading interacts with CUDA graphs:

  • A dedicated cuda_graph_stream runs the captured computation, while d2h_stream overlaps D2H transfers for regions that are inside the graph capture.

  • During CUDA graph warmup, offloading is disabled (pre_warmup_hook / post_warmup_hook).

  • The delay_offload_until_cuda_graph option defers D2H launches until graph replay, utilizing the CPU idle time during cudaGraphLaunch to issue offload commands with near-zero CPU overhead.