diff --git a/docs/user-guide/features/fine_grained_activation_offloading.md b/docs/user-guide/features/fine_grained_activation_offloading.md index 915926a6b9b..d264d4d7201 100644 --- a/docs/user-guide/features/fine_grained_activation_offloading.md +++ b/docs/user-guide/features/fine_grained_activation_offloading.md @@ -11,22 +11,13 @@ Contributed in collaboration with RedNote. -Memory is often the limiting factor for very large sparse MoE models such as DeepSeek-V3 and Qwen3-235B. Fine-grained recomputation lowers activation memory at the cost of extra compute. Offloading can use host-device bandwidth so that reload overlaps compute and keeps overhead small in many setups. Fine-grained activation offloading moves activations at module granularity so you can tune how much activation memory leaves the device and adjust training throughput. +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"`, and `"fused_group_mlp"`. 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`. -## Features +## User Guide -- Pipeline parallelism: PP=1, PP, and interleaved PP -- Compatible with fine-grained recomputation -- FP8 training -- MTP -- Mixed dense and MoE layers -- A2A overlap -- CUDA graphs - - **Note:** A CUDA graph capture cannot include the offloading modules (temporary limitation). - -## Usage +### Basic Usage ```bash # Enable fine-grained activation offloading @@ -34,26 +25,177 @@ Supported offloading modules are `"attn_norm"`, `"qkv_linear"`, `"core_attn"`, ` # 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". ---offload-modules expert_fc1 +--offload-modules core_attn attn_proj expert_fc1 ``` -## Max inflight offloads +### 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` | + +### Tuning Parameters ```bash +# 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 ``` -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. +`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 fraction is applied once across +the combined eligible groups from all configured modules. + +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. + +### 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`. + +```bash +# 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. + +```bash +--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](../../images/fine_grained_activation_offloading/offloading_and_recomputing.png) + + +### 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. + +### Warmup and Adaptive Offloading + +The first training iteration serves as a **warmup phase** where the manager records tensor groups, their sizes, and the execution order. After warmup, a `post_warmup_callback` runs to: + +1. **Reserve margin**: The last N groups (by deduplication count) are kept on GPU to avoid reload blocking the compute stream. +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 -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. +A 'OffloadTensorPool` (on CPU with pinned memory) caches allocated tensors by `(shape, dtype)`. This avoids repeated `cudaMallocHost` / `cudaFreeHost` calls and reduces D2H latency after the first iteration. -## Compatible With Fine-Grained Recomputation +### CUDA Graph Support -- For low-overhead modules such as LayerNorm or `moe_act`, use recomputation to save activation memory. -- For other modules, use offloading to save activation memory. -- Overlap offload and reload with compute when possible. +When offloading interacts with CUDA graphs: -![Diagram comparing fine-grained activation offloading and fine-grained recomputation across a transformer layer](../../images/fine_grained_activation_offloading/offloading_and_recomputing.png) +- 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. diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index c6a67e1094d..32dec66a2cd 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import weakref from contextlib import nullcontext @@ -537,18 +537,16 @@ def forward_func( ) if not isinstance(layer.mlp, MoELayer): return hidden_states, None, None, None + mlp_norm_manager = off_interface(layer.offload_mlp_norm, hidden_states, "mlp_norm") + node.layer_state.mlp_norm_manager = mlp_norm_manager if layer.recompute_pre_mlp_layernorm: layer.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface( - layer.offload_mlp_norm, hidden_states, "mlp_norm" - ) as hidden_states: + with mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = layer.pre_mlp_norm_checkpoint.checkpoint( apply_module(layer.pre_mlp_layernorm), hidden_states ) else: - with off_interface( - layer.offload_mlp_norm, hidden_states, "mlp_norm" - ) as hidden_states: + with mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = apply_module(layer.pre_mlp_layernorm)( hidden_states ) @@ -662,10 +660,12 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): ) # Delay the offload of the mlp norm until after the mlp_bda has been computed # because the residual is needed in the mlp_bda. - if layer.offload_mlp_norm: - hidden_states = off_interface.group_commit( - hidden_states, name="mlp_norm", forced_released_tensors=[residual] + mlp_norm_manager = getattr(node.layer_state, 'mlp_norm_manager', None) + if mlp_norm_manager is not None: + hidden_states = mlp_norm_manager.group_offload( + hidden_states, forced_released_tensors=[residual] ) + node.layer_state.mlp_norm_manager = None output = make_viewless_tensor( inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True ) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 605ae3b02ee..d2582a3f353 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from collections import OrderedDict from typing import Any, Callable, Dict, Literal, Optional @@ -474,20 +474,23 @@ def _preprocess( def preprocess_for_fine_grained_offloading(self): """Preprocess for fine-grained activation offloading.""" off_interface.init_chunk_handler( + pp_rank=self.pg_collection.pp.rank(), vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage, min_offloaded_tensor_size=self.config.min_offloaded_tensor_size, + delta_offload_bytes_across_pp_ranks=self.config.delta_offload_bytes_across_pp_ranks, + activation_offload_fraction=self.config.activation_offload_fraction, max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads, ) if self.disable_param_offloading: for param in self.decoder.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) if self.mtp_process: for param in self.mtp.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) if self.post_process: for param in self.output_layer.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) self.disable_param_offloading = False def preprocess_for_paged_stash(self): diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 1637c9909f1..84e5d078554 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from typing import Literal, Optional @@ -351,20 +351,23 @@ def set_input_tensor(self, input_tensor: Tensor) -> None: def preprocess_for_fine_grained_offloading(self): """Preprocess for fine-grained activation offloading.""" off_interface.init_chunk_handler( + pp_rank=self.pg_collection.pp.rank(), vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage, min_offloaded_tensor_size=self.config.min_offloaded_tensor_size, + delta_offload_bytes_across_pp_ranks=self.config.delta_offload_bytes_across_pp_ranks, + activation_offload_fraction=self.config.activation_offload_fraction, max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads, ) if self.disable_param_offloading: for param in self.decoder.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) if self.mtp_process: for param in self.mtp.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) if self.post_process: for param in self.output_layer.parameters(): - off_interface.mark_not_offloadable(param) + off_interface.mark_not_offload(param) self.disable_param_offloading = False def preprocess_for_paged_stash(self): diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index e5c82876516..bcc8c14ebb6 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -5,6 +5,7 @@ from typing import Any, Dict, Optional, Tuple import torch +from torch.autograd.graph import saved_tensors_hooks # CPU offload implementation for pipeline parallelism DEBUG = False @@ -111,9 +112,9 @@ def print_offload_summary_table(total_offload_bytes: Dict[str, int]): torch.distributed.barrier() -class GPUTensorPool: +class OffloadTensorPool: """ - GPU memory pool for efficient allocation and deallocation of tensors. + Memory pool for efficient allocation and deallocation of tensors. Features: - Supports multiple tensor shapes and dtypes, each with its own pool @@ -122,7 +123,7 @@ class GPUTensorPool: - Uses queue-based management for O(1) allocation and deallocation Example: - pool = GPUTensorPool(device='cuda:0') + pool = OffloadTensorPool(device='cuda:0') tensor = pool.allocate((128, 512), dtype=torch.float32) # ... use tensor ... pool.free(tensor, (128, 512), dtype=torch.float32) @@ -130,10 +131,10 @@ class GPUTensorPool: def __init__(self, device: str = 'cuda', pin_memory: bool = False): """ - Initialize GPU tensor pool. + Initialize offload tensor pool. Args: - device: GPU device, default 'cuda' + device: Device, default 'cuda' pin_memory: Whether to use pinned memory (mainly for CPU tensors) """ self.device = torch.device(device) @@ -153,7 +154,7 @@ def __init__(self, device: str = 'cuda', pin_memory: bool = False): 'pool_misses': 0, # Number of times a new tensor was created } - debug_rank("GPUTensorPool: Initialized with dynamic allocation") + debug_rank("OffloadTensorPool: Initialized with dynamic allocation") def _get_pool_key(self, shape: Tuple, dtype: torch.dtype) -> Tuple: """Generate a unique key for the pool based on shape and dtype.""" @@ -198,7 +199,7 @@ def allocate(self, shape: Tuple, dtype: torch.dtype = torch.float32) -> torch.Te tensor = pool['free'].popleft() self._stats['pool_hits'] += 1 debug_rank( - f"GPUTensorPool.allocate: Reused tensor from pool, " + f"OffloadTensorPool.allocate: Reused tensor from pool, " f"shape={shape}, dtype={dtype}, " f"remaining in pool={len(pool['free'])}" ) @@ -211,7 +212,7 @@ def allocate(self, shape: Tuple, dtype: torch.dtype = torch.float32) -> torch.Te memory_mb = self._calculate_memory_size(shape, dtype) / (1024**2) debug_rank( - f"GPUTensorPool.allocate: Created new tensor, " + f"OffloadTensorPool.allocate: Created new tensor, " f"shape={shape}, dtype={dtype}, " f"memory={memory_mb:.2f} MB, " f"total_created={len(pool['all'])}" @@ -261,7 +262,7 @@ def free(self, tensor: torch.Tensor): self._stats['current_in_use'] -= 1 debug_rank( - f"GPUTensorPool.free: shape={shape}, dtype={dtype}, " + f"OffloadTensorPool.free: shape={shape}, dtype={dtype}, " f"available in pool={len(pool['free'])}" ) @@ -310,7 +311,7 @@ def get_pool_status(self, shape: Tuple = None, dtype: torch.dtype = None) -> Dic def reset(self): """Reset the pool, marking all tensors as available.""" - debug_rank("GPUTensorPool: Resetting pool...") + debug_rank("OffloadTensorPool: Resetting pool...") for pool_key, pool in self._pools.items(): # Clear and refill the free queue @@ -320,11 +321,11 @@ def reset(self): pool['allocated_count'] = 0 self._stats['current_in_use'] = 0 - debug_rank("GPUTensorPool: Reset complete") + debug_rank("OffloadTensorPool: Reset complete") def clear(self): """Clear the pool and release all GPU memory.""" - debug_rank("GPUTensorPool: Clearing pool...") + debug_rank("OffloadTensorPool: Clearing pool...") for pool_key, pool in self._pools.items(): # Clear all references, allowing PyTorch GC to reclaim memory @@ -338,7 +339,7 @@ def clear(self): if torch.cuda.is_available(): torch.cuda.empty_cache() - debug_rank("GPUTensorPool: Clear complete") + debug_rank("OffloadTensorPool: Clear complete") def __del__(self): """Destructor to ensure resources are released.""" @@ -427,11 +428,18 @@ def __init__(self): # allocate streams and events for synchronization self._d2h_stream = torch.cuda.Stream() self._h2d_stream = torch.cuda.Stream() + # TE CUDA graph offload paths need a stream/event pair that lives outside + # individual layer objects so capture, replay, and backward hooks order + # the same D2H/H2D work with the same synchronization primitives. + self._cuda_graph_stream = torch.cuda.Stream() + self._cuda_graph_event = torch.cuda.Event(external=True) # Shared CPU tensor pool for all chunks to improve reuse efficiency - self._cpu_tensor_pool = GPUTensorPool(device="cpu", pin_memory=True) + self._cpu_tensor_pool = OffloadTensorPool(device="cpu", pin_memory=True) # Whether the manager is in warmup phase. self._is_warmup = True + # Whether the manager is in CUDA graph replay phase. + self._in_replay = False # Cache OffloadChunkHandler objects for each virtual pipeline stage and each forward pass. self._cached_chunks_forward = [] # Cache OffloadChunkHandler objects for each virtual pipeline stage and each backward pass. @@ -450,6 +458,12 @@ def __init__(self): self._delayed_offload_groups = [] self.reset() + # Keep the hook context object around so each offload scope can enter/exit + # the same autograd saved-tensor hooks without touching private torch APIs. + self._saved_tensors_hooks = saved_tensors_hooks( + self.on_save_for_backward, self.on_get_saved_tensor + ) + @property def d2h_stream(self): """Get the device-to-host (GPU to CPU) transfer stream.""" @@ -460,22 +474,35 @@ def h2d_stream(self): """Get the host-to-device (CPU to GPU) transfer stream.""" return self._h2d_stream + @property + def cuda_graph_stream(self): + """Get the CUDA graph stream.""" + return self._cuda_graph_stream + + @property + def cuda_graph_event(self): + """Get the CUDA graph event.""" + return self._cuda_graph_event + @property def cpu_tensor_pool(self): """Get the shared CPU tensor pool.""" return self._cpu_tensor_pool - def push_offload_groups(self, group_hook, forced_released_tensors): + def push_offload_groups(self, group_hook, name, forced_released_tensors): """Push the offload groups to the delayed queue.""" debug_rank(f"pushing offload groups to the delayed queue") - self._delayed_offload_groups.append((group_hook, forced_released_tensors)) + # Store the group name because delayed CUDA graph replay flushes later, + # after the original group-start site has already moved on. + self._delayed_offload_groups.append((group_hook, name, forced_released_tensors)) def flush_delayed_groups(self): """Flush the delayed groups.""" debug_rank("flushing delayed groups") - # Flush the delayed groups in reverse order to maintain the order of the groups. - for group_hook, forced_released_tensors in reversed(self._delayed_offload_groups): - group_hook(forced_released_tensors) + # Preserve the original forward commit order; reload scheduling still + # relies on the same group order discovered during warmup. + for group_hook, name, forced_released_tensors in self._delayed_offload_groups: + group_hook(name, forced_released_tensors) self._delayed_offload_groups = [] def reset(self): @@ -566,13 +593,44 @@ def post_warmup_callback(self): debug_rank(f"setting offload to false for group {name} at chunk index {chunk_idx}") else: break - debug_rank(f"offload margin {self._offload_margin}") assert self._offload_margin == 0, "Offload margin is not 0" + # Disable the groups to meet the delta offload bytes across PP ranks. + keep_on_gpu_bytes = self._pp_rank * self._delta_offload_bytes_across_pp_ranks + for chunk in self._cached_chunks_backward: + for group in chunk.offload_groups: + if group.offload and keep_on_gpu_bytes > 0: + debug_rank( + f"group {group._name} offload {group.offload} \ + keep_on_gpu_bytes {keep_on_gpu_bytes}" + ) + keep_on_gpu_bytes -= group.total_offload_bytes + group.offload = False + # Disable the later groups to meet the activation offload fraction. + for chunk in self._cached_chunks_backward: + eligible_offload_groups = [ + group + for group in chunk.offload_groups + if group.offload and group.total_offload_bytes > 0 + ] + offloaded_groups_count = len(eligible_offload_groups) + disabled_groups_count = int( + offloaded_groups_count * (1 - self._activation_offload_fraction) + ) + debug_rank(f"Disabled {disabled_groups_count}/{offloaded_groups_count} groups") + # Prefer keeping earlier forward groups offloaded because releasing + # those activations sooner gives the longest memory-pressure relief. + for group in reversed(eligible_offload_groups): + if disabled_groups_count > 0: + disabled_groups_count -= 1 + group.offload = False + else: + break # Dump the offload information total_tensor_count = {} total_offload_bytes = {} for chunk in self._cached_chunks_forward: for group in chunk.offload_groups: + debug_rank(f"chunk {chunk} group {group} offload {group.offload}") if group.offload: if group._name not in total_tensor_count: total_tensor_count[group._name] = 0 @@ -584,6 +642,8 @@ def post_warmup_callback(self): # where the memory cost will not increase anymore. if chunk is self._cached_chunks_backward[0]: break + debug_rank(f"total_tensor_count {total_tensor_count}") + debug_rank(f"total_offload_bytes {total_offload_bytes}") # Cache summary for downstream consumers (e.g., unit tests). self._offload_summary_bytes = dict(total_offload_bytes) self._offload_summary_total_bytes = int(sum(total_offload_bytes.values())) @@ -625,18 +685,25 @@ def front_backward_chunk(self, name=None): def init_model_chunk_offload_handler( self, + pp_rank, vp_size, vp_stage, min_offloaded_tensor_size=1024 * 1024, + delta_offload_bytes_across_pp_ranks=0, + activation_offload_fraction: float = 1.0, max_inflight_offloads: Optional[int] = None, ): """ Initialize a chunk offload handler for a model chunk (microbatch). Args: + pp_rank: Pipeline parallel rank vp_size: Virtual pipeline size vp_stage: Virtual pipeline stage index (None means stage 0) min_offloaded_tensor_size: Minimum tensor size (in elements) to offload + delta_offload_bytes_across_pp_ranks: + Difference of offload bytes across PP ranks to balance the offload load. + activation_offload_fraction: Fraction of eligible groups to offload, in range [0, 1]. max_inflight_offloads: If set, cap pending offloads per group name before main wait_event; see ``fine_grained_offloading_max_inflight_offloads`` on ``TransformerConfig``. @@ -649,6 +716,10 @@ def init_model_chunk_offload_handler( self._vpp = vp_size self._stages = [[] for _ in range(vp_size)] + self._delta_offload_bytes_across_pp_ranks = delta_offload_bytes_across_pp_ranks + self._pp_rank = pp_rank + self._activation_offload_fraction = activation_offload_fraction + if vp_stage is None: cur_vpp_rank = 0 else: @@ -698,10 +769,12 @@ def cur_backward_chunk(self): """Get the current backward pass chunk handler.""" return self._cur_backward_chunk - def mark_not_offloadable(self, tensor: torch.Tensor): + def mark_not_offload(self, tensor: torch.Tensor): """Mark the current forward chunk as not offloadable.""" if tensor is not None: - tensor.offloading_activation = False + # TE marks some tensors with _TE_do_not_offload; this local flag + # gives Megatron-owned tensors the same opt-out path. + tensor._do_not_offload = True def __enter__(self): """Enter context manager to enable activation offloading hooks.""" @@ -715,10 +788,7 @@ def __enter__(self): else: raise RuntimeError("TE CPU offload is not available") self.inside_context = True - - torch._C._autograd._push_saved_tensors_default_hooks( - self.on_save_for_backward, self.on_get_saved_tensor - ) + self._saved_tensors_hooks.__enter__() def __exit__(self, *args: Any): """Exit context manager and restore original tensor saving behavior.""" @@ -732,7 +802,7 @@ def __exit__(self, *args: Any): else: raise RuntimeError("TE CPU offload is not available") self.inside_context = False - torch._C._autograd._pop_saved_tensors_default_hooks() + self._saved_tensors_hooks.__exit__() def on_save_for_backward(self, tensor: torch.Tensor) -> Any: """ @@ -834,17 +904,17 @@ def reset(self): # an event recorded in a previous (non-captured) iteration. self._offload_pending_by_name.clear() - def find_group_with_name(self, name: str, start_index: int = 0): + def find_group_with_name( + self, groups: list[OffloadTensorGroup], name: str, start_index: int = 0 + ): """Find the group with the given name starting from the given index.""" - return next( - (group for group in self.offload_groups[start_index:] if group._name == name), None - ) + return next((group for group in groups[start_index:] if group._name == name), None) def is_empty_chunk(self, name=None): """Check if this chunk has no tensors to manage.""" debug_rank(f"------is_empty_chunk {self._max_group_size}") if name is not None: - return self.find_group_with_name(name) is None + return self.find_group_with_name(self.offload_groups, name) is None return self._max_group_size == 0 def finish_all_groups(self, name=None) -> bool: @@ -861,12 +931,15 @@ def finish_all_groups(self, name=None) -> bool: ): return True assert name is not None, "Name is required" - return self.find_group_with_name(name, self._offloaded_group_index) is None + return ( + self.find_group_with_name(self.offload_groups, name, self._offloaded_group_index) + is None + ) def find_next_group(self, name=None): """Find the next group with the given name.""" assert name is not None, "Name is required" - return self.find_group_with_name(name, self._offloaded_group_index) + return self.find_group_with_name(self.offload_groups, name, self._offloaded_group_index) @staticmethod def _can_manage_tensor_for_offload(tensor): @@ -912,9 +985,7 @@ def tensor_pop(self, tensor_tag): def tensor_need_offloading_checker(self, tensor): """Check if the tensor needs to be offloaded.""" - debug_rank( - f"tensor_need_offloading_checker {getattr(tensor, 'offloading_activation', None)}" - ) + debug_rank("tensor_need_offloading_checker") if not self._can_manage_tensor_for_offload(tensor): return False if _te_do_not_offload(tensor): @@ -922,14 +993,15 @@ def tensor_need_offloading_checker(self, tensor): if tensor.numel() < self.min_offloaded_tensor_size: return False # Respect tensor's offload preference if specified - if hasattr(tensor, "offloading_activation") and not tensor.offloading_activation: + if getattr(tensor, "_TE_do_not_offload", False) or getattr( + tensor, "_do_not_offload", False + ): return False return True - def bulk_offload_group(self): + def bulk_offload_group(self, group_to_offload): """offload a group of tensors recorded in tensor_push().""" debug_rank("------bulk_offload_group") - group_to_offload = self._groups_to_offload[-1] nvtx_msg = "activation offloading " + group_to_offload._name nvtx_range_push(nvtx_msg) with torch.cuda.stream(self.d2h_stream): @@ -943,7 +1015,6 @@ def bulk_offload_group(self): tensor_on_device.record_stream(self.d2h_stream) group_to_offload.push_tensor(tensor_tag, state) group_to_offload.record_offload_event(self.d2h_stream) - self._groups_to_offload.pop() nvtx_range_pop(nvtx_msg) # Under full-iteration CG capture, the main stream may not wait on d2h # events; optional max-inflight enqueues each group's offload event and @@ -992,10 +1063,9 @@ def pre_reload_last_layer(self): # Reload the last group (last layer) early self.bulk_reload_group() - def should_bulk_offload(self): + def should_bulk_offload(self, group): """Determine if the current group should be offloaded.""" - assert len(self._groups_to_offload) > 0, "No groups to offload" - group = self._groups_to_offload[-1] + assert group in self._groups_to_offload, f"Group {group} is not pending offload" debug_rank(f"should_bulk_offload {self.is_warmup} {group.offload}") # Don't offload if the chunk is not in warmup stage if self.is_warmup: @@ -1016,12 +1086,16 @@ def should_bulk_offload(self): return True - def bulk_offload(self, forced_released_tensors): + def bulk_offload(self, name, forced_released_tensors): """Offload a group of tensors and optionally release their GPU memory.""" debug_rank("----bulk_offload") - if self.should_bulk_offload(): - self._groups_to_reload.append(self._groups_to_offload[-1]) - self.bulk_offload_group() + # CUDA graph scoped modules can create several pending groups before a + # commit runs, so match by name instead of assuming LIFO order. + group_to_offload = self.find_group_with_name(self._groups_to_offload, name) + assert group_to_offload is not None, f"Group {name} not found in {self._groups_to_offload}" + if self.should_bulk_offload(group_to_offload): + self._groups_to_reload.append(group_to_offload) + self.bulk_offload_group(group_to_offload) # Manually release tensors not auto-freed by torch GC if len(forced_released_tensors) > 0: cur_stream = torch.cuda.current_stream() @@ -1030,6 +1104,8 @@ def bulk_offload(self, forced_released_tensors): # Ensure tensor is not in use before freeing release_tensor.record_stream(cur_stream) release_tensor.untyped_storage().resize_(0) + # A group commit is consumed even when policy keeps its tensors on GPU. + self._groups_to_offload.remove(group_to_offload) def _drain_offload_pending(self, group_name: str) -> None: """For ``group_name``, have the main stream wait on older D2H events @@ -1043,14 +1119,14 @@ def _drain_offload_pending(self, group_name: str) -> None: old_evt = q.popleft() cur.wait_event(old_evt) - def on_group_commit_forward(self, forced_released_tensors): + def on_group_commit_forward(self, name, forced_released_tensors): """Called at the end of a layer group's forward pass to trigger offloading.""" if not self.do_offload: return - debug_rank("--on_group_commit_forward") + debug_rank(f"--on_group_commit_forward {name}") # Wait for compute to finish before starting offload self.d2h_stream.wait_stream(torch.cuda.current_stream()) - self.bulk_offload(forced_released_tensors) + self.bulk_offload(name, forced_released_tensors) def bulk_reload(self): """Reload the next group of tensors from CPU to GPU.""" @@ -1149,12 +1225,14 @@ def forward(ctx, tensor, cur_forward_chunk, name, forced_released_tensors, delay # pylint: disable=missing-function-docstring debug_rank("FineGrainedOffloadingGroupCommitFunction forward") - if delay_offload: + if delay_offload and PipelineOffloadManager.get_instance()._in_replay: + # During TE CUDA graph replay, queue D2H work and launch it after + # replay returns, where CPU scheduling can overlap with graph/comm gaps. PipelineOffloadManager.get_instance().push_offload_groups( - cur_forward_chunk.on_group_commit_forward, forced_released_tensors + cur_forward_chunk.on_group_commit_forward, name, forced_released_tensors ) else: - cur_forward_chunk.on_group_commit_forward(forced_released_tensors) + cur_forward_chunk.on_group_commit_forward(name, forced_released_tensors) ctx.cpu_offload_handler = cur_forward_chunk ctx.name = name return tensor @@ -1169,7 +1247,7 @@ def backward(ctx, *grad_output): return grad_output + (None, None, None, None) -def fine_grained_offloading_group_commit( +def fine_grained_offloading_group_offload( tensor, name, forced_released_tensors=None, delay_offload=False ): """ @@ -1186,23 +1264,23 @@ def fine_grained_offloading_group_commit( if isinstance(tensor, tuple): if len(tensor) == 0: return tensor - committed0 = fine_grained_offloading_group_commit( + offloaded0 = fine_grained_offloading_group_offload( tensor[0], name=name, forced_released_tensors=forced_released_tensors, delay_offload=delay_offload, ) - return (committed0,) + tensor[1:] + return (offloaded0,) + tensor[1:] if isinstance(tensor, list): if len(tensor) == 0: return tensor - committed0 = fine_grained_offloading_group_commit( + offloaded0 = fine_grained_offloading_group_offload( tensor[0], name=name, forced_released_tensors=forced_released_tensors, delay_offload=delay_offload, ) - return [committed0] + tensor[1:] + return [offloaded0] + tensor[1:] cur_forward_chunk = PipelineOffloadManager.get_instance().cur_forward_chunk() if cur_forward_chunk is None: @@ -1251,13 +1329,6 @@ def fine_grained_offloading_group_start(tensor, name=None): return FineGrainedOffloadingGroupStartFunction.apply(tensor, cur_forward_chunk, name) -def fine_grained_offloading_forward_record(event: torch.cuda.Event) -> None: - """Record the forward event for cuda graph capture.""" - d2h_stream = PipelineOffloadManager.get_instance().d2h_stream - torch.cuda.current_stream().record_event(event) - torch.cuda.current_stream().wait_stream(d2h_stream) - - class FineGrainedOffloadingBackwardRecordFunction(torch.autograd.Function): """ Identity operation that marks the end of a layer group for offload synchronization. @@ -1265,23 +1336,21 @@ class FineGrainedOffloadingBackwardRecordFunction(torch.autograd.Function): """ @staticmethod - def forward(ctx, tensor, event: torch.cuda.Event) -> torch.Tensor: + def forward(ctx, tensor) -> torch.Tensor: """Forward pass for cuda graph capture.""" - ctx.event = event + debug_rank("FineGrainedOffloadingBackwardRecordFunction forward") return tensor @staticmethod def backward(ctx, grad_output): """Record the backward event and wait for the h2d stream on cuda graph stream.""" - h2d_stream = PipelineOffloadManager.get_instance().h2d_stream - torch.cuda.current_stream().record_event(ctx.event) - torch.cuda.current_stream().wait_stream(h2d_stream) - return grad_output, None - - -def fine_grained_offloading_backward_record(tensor, event: torch.cuda.Event) -> torch.Tensor: - """Record the backward event for cuda graph capture.""" - return FineGrainedOffloadingBackwardRecordFunction.apply(tensor, event) + debug_rank("FineGrainedOffloadingBackwardRecordFunction backward") + mgr = PipelineOffloadManager.get_instance() + # This event connects TE's graph stream with the reload stream so + # backward consumers do not race H2D reloads launched outside the graph. + torch.cuda.current_stream().record_event(mgr.cuda_graph_event) + torch.cuda.current_stream().wait_stream(mgr.h2d_stream) + return (grad_output,) class FineGrainedActivationOffloadingInterface: @@ -1304,15 +1373,34 @@ def __exit__(self, *args: Any): if self.offload: PipelineOffloadManager.get_instance().__exit__() + @staticmethod + def cuda_graph_stream(): + """Get the CUDA graph stream.""" + return PipelineOffloadManager.get_instance().cuda_graph_stream + + @staticmethod + def cuda_graph_event(): + """Get the CUDA graph event.""" + return PipelineOffloadManager.get_instance().cuda_graph_event + @staticmethod def init_chunk_handler( - vp_size, vp_stage, min_offloaded_tensor_size, max_inflight_offloads: Optional[int] = None + pp_rank, + vp_size, + vp_stage, + min_offloaded_tensor_size, + delta_offload_bytes_across_pp_ranks, + activation_offload_fraction, + max_inflight_offloads: Optional[int] = None, ): """Initialize the chunk handler, called at the start of a microbatch forward pass.""" PipelineOffloadManager.get_instance().init_model_chunk_offload_handler( + pp_rank, vp_size, vp_stage, min_offloaded_tensor_size, + delta_offload_bytes_across_pp_ranks, + activation_offload_fraction, max_inflight_offloads=max_inflight_offloads, ) @@ -1321,24 +1409,30 @@ def get_context(flag): """Get the fine-grained offload context""" return PipelineOffloadManager.get_instance() if flag else nullcontext() - @staticmethod - def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False): - """Group commit the tensors.""" - return fine_grained_offloading_group_commit( - tensor, name, forced_released_tensors, delay_offload - ) + def group_offload(self, tensor, forced_released_tensors=None, delay_offload=False): + """Group offload the tensors.""" + if self.offload: + return fine_grained_offloading_group_offload( + tensor, self.name, forced_released_tensors, delay_offload + ) + return tensor @staticmethod - def mark_not_offloadable(tensor: torch.Tensor): + def mark_not_offload(tensor: torch.Tensor): """Mark the tensor as not offloadable.""" - PipelineOffloadManager.get_instance().mark_not_offloadable(tensor) + PipelineOffloadManager.get_instance().mark_not_offload(tensor) @staticmethod - def forward_record(event: torch.cuda.Event) -> None: + def forward_record() -> None: """Record the forward event for cuda graph capture.""" - d2h_stream = PipelineOffloadManager.get_instance().d2h_stream - torch.cuda.current_stream().record_event(event) - torch.cuda.current_stream().wait_stream(d2h_stream) + mgr = PipelineOffloadManager.get_instance() + torch.cuda.current_stream().record_event(mgr.cuda_graph_event) + torch.cuda.current_stream().wait_stream(mgr.d2h_stream) + + @staticmethod + def backward_record(tensor) -> torch.Tensor: + """Record the backward event for cuda graph capture.""" + return FineGrainedOffloadingBackwardRecordFunction.apply(tensor) @staticmethod def reset(): @@ -1349,3 +1443,28 @@ def reset(): def reset_instance(): """Reset the singleton instance.""" PipelineOffloadManager.reset_instance() + + @staticmethod + def flush_delayed_groups(): + """Flush the delayed groups.""" + PipelineOffloadManager.get_instance().flush_delayed_groups() + + @staticmethod + def disable_offload(): + """Disable the offload.""" + PipelineOffloadManager.get_instance().disable_offload() + + @staticmethod + def enable_offload(): + """Enable the offload.""" + PipelineOffloadManager.get_instance().enable_offload() + + @staticmethod + def enter_replay(): + """Enter CUDA graph replay mode to enable delayed offloading.""" + PipelineOffloadManager.get_instance()._in_replay = True + + @staticmethod + def exit_replay(): + """Exit CUDA graph replay mode.""" + PipelineOffloadManager.get_instance()._in_replay = False diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index d875367e93e..92e2bccb8cf 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import copy @@ -1344,18 +1344,16 @@ def forward( self.config.fused_single_qkv_rope and split_qkv ), "fused_single_qkv_rope requested but not available/supported for the config." - with off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") as hidden_states: + qkv_linear_manager = off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") + with qkv_linear_manager as hidden_states: qkv_output = self.get_query_key_value_tensors( hidden_states, key_value_states, split_qkv=split_qkv, output_gate=self.config.attention_output_gate, ) - if self.offload_qkv_linear: - # `qkv_output` may be a tuple; commit supports tuple/list and will keep structure. - qkv_output = off_interface.group_commit( - qkv_output, name="qkv_linear", forced_released_tensors=[] - ) + # `qkv_output` may be a tuple; commit supports tuple/list and will keep structure. + qkv_output = qkv_linear_manager.group_offload(qkv_output, forced_released_tensors=[]) attn_mask_type = self.attn_mask_type block_table = None gate = None @@ -1502,6 +1500,9 @@ def forward( # ================================== nvtx_range_push(suffix="core_attention") + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) if self.checkpoint_core_attention and self.training: core_attn_out = self._checkpointed_attention_forward( query, @@ -1515,9 +1516,7 @@ def forward( else: if inference_context is None or inference_context.is_static_batching(): # Static batching attention kernel. - with off_interface( - self.offload_core_attention and self.training, query, "core_attn" - ) as query: + with core_attn_manager as query: core_attn_out = apply_module(self.core_attention)( query, key, @@ -1554,10 +1553,9 @@ def forward( if is_using_quantization_scales(self.config): core_attn_out[inference_context.padding_slice] = 0.0 - if self.offload_core_attention and self.training: - core_attn_out = off_interface.group_commit( - core_attn_out, name="core_attn", forced_released_tensors=[query, key, value] - ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': # reshape to same output shape as unpacked case # (t, np, hn) -> (t, b=1, h=np*hn) @@ -1576,12 +1574,10 @@ def forward( # Output. [sq, b, h] # ================= nvtx_range_push(suffix="linear_proj") - with off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") as core_attn_out: + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: output, bias = apply_module(self.linear_proj)(core_attn_out) - if self.offload_attn_proj: - output = off_interface.group_commit( - output, name="attn_proj", forced_released_tensors=[core_attn_out] - ) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) nvtx_range_pop(suffix="linear_proj") return output, bias diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 0de90c9cde4..2b2893d9350 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import dataclasses import gc @@ -2310,6 +2310,17 @@ def _get_fp8_enabled(): ) else: kwargs['fp8_enabled'] = False + + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, + ) + + # TE CUDA graph warmup should establish graph state without launching + # activation D2H copies; the post-warmup hook restores offloading for + # the measured/replay iterations. + if self.config.fine_grained_activation_offloading: + kwargs['pre_warmup_hook'] = off_interface.disable_offload + kwargs['post_warmup_hook'] = off_interface.enable_offload return kwargs kwargs = get_make_graphed_callables_kwargs() @@ -2357,6 +2368,12 @@ def _finish_capturing(self, start_time): ) _set_capture_end() + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, + ) + + if self.config.fine_grained_activation_offloading: + off_interface.reset() torch.cuda.synchronize() self._reset_after_capture() if FREEZE_GC: diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index c30c107e791..bae3c70cf9c 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron Module.""" from functools import partial @@ -318,6 +318,18 @@ def _get_te_cuda_graph_replay_args(self, *args, **kwargs): cudagraph_kwargs = kwargs.copy() cudagraph_kwargs['is_first_microbatch'] = getattr(self, 'current_microbatch', 0) == 0 + if self.config.fine_grained_activation_offloading and getattr( + self, 'offload_module_in_cuda_graph', False + ): + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, + ) + + # TE captures/replays the module on its own graph stream. Passing the + # offload stream/event in lets TE order graph compute with D2H/H2D + # transfers managed by the fine-grained offload manager. + cudagraph_kwargs['cuda_graph_stream'] = off_interface.cuda_graph_stream() + cudagraph_kwargs['cuda_graph_event'] = off_interface.cuda_graph_event() return cudagraph_args, cudagraph_kwargs def _should_call_local_cudagraph(self, *args, **kwargs): diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index bde1737fce0..6f71c83465f 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import inspect @@ -18,6 +18,7 @@ from megatron.core.activations import squared_relu from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding +from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.fusions.fused_bias_geglu import quick_gelu, weighted_bias_quick_geglu_impl from megatron.core.fusions.fused_bias_swiglu import weighted_bias_swiglu_impl @@ -265,7 +266,9 @@ def __init__( set_save_original_input(self.linear_fc2) # This is to avoid the CPU overhead of multiple d2h copies - if self.offload_expert_fc1: + use_mxfp8 = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.mxfp8 + use_nvfp4 = self.config.fp4 and self.config.fp4_recipe == Fp4Recipe.nvfp4 + if self.offload_expert_fc1 and not (use_mxfp8 or use_nvfp4): from megatron.core.extensions.transformer_engine import set_save_original_input set_save_original_input(self.linear_fc1) @@ -610,9 +613,10 @@ def _fused_forward( stash_context = nullcontext() fine_grained_activation_offloading = getattr(self, "offload_fused_group_mlp", False) offload_name = "fused_group_mlp" - with off_interface( + fused_group_mlp_manager = off_interface( fine_grained_activation_offloading, permuted_local_hidden_states, offload_name - ) as permuted_local_hidden_states: + ) + with fused_group_mlp_manager as permuted_local_hidden_states: forced_released_tensors = ( [permuted_local_hidden_states] if fine_grained_activation_offloading else [] ) @@ -624,10 +628,9 @@ def _fused_forward( permuted_probs, # Scaled activation tokens_per_expert, # FC2 ) - if fine_grained_activation_offloading: - output = off_interface.group_commit( - output, name=offload_name, forced_released_tensors=forced_released_tensors - ) + output = fused_group_mlp_manager.group_offload( + output, forced_released_tensors=forced_released_tensors + ) # Remove padding if needed if unpadded_tokens_per_expert is not None: output = self.quantization_unpadding(output, unpadded_tokens_per_expert) @@ -695,18 +698,20 @@ def forward( # Probs already applied, so reset to 1. permuted_probs = torch.ones_like(permuted_probs) - with off_interface( + expert_fc1_manager = off_interface( self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1" - ) as permuted_local_hidden_states: + ) + with expert_fc1_manager as permuted_local_hidden_states: fc1_output, bias_parallel = apply_module(self.linear_fc1)( permuted_local_hidden_states, tokens_per_expert ) - if self.offload_expert_fc1: - fc1_output = off_interface.group_commit( - fc1_output, - name="expert_fc1", - forced_released_tensors=[permuted_local_hidden_states], - ) + fc1_output = expert_fc1_manager.group_offload( + fc1_output, + forced_released_tensors=[permuted_local_hidden_states], + delay_offload=self.config.delay_offload_until_cuda_graph, + ) + + moe_act_manager = off_interface(self.offload_moe_act, fc1_output, "moe_act") def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs): @@ -785,12 +790,12 @@ def glu(x): if self.activation_recompute: self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: + with moe_act_manager as fc1_output: bias_act_output = self.activation_checkpoint.checkpoint( bias_act_func, fc1_output, bias_parallel, permuted_probs ) else: - with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: + with moe_act_manager as fc1_output: bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs) output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert) if self.activation_recompute: @@ -798,10 +803,11 @@ def glu(x): # Delay the offload of the moe act until after the linear_fc2 has been computed # to make sure the fc1_output is reloaded to GPU before recomputing moe_act. - if self.offload_moe_act: - output = off_interface.group_commit( - output, name="moe_act", forced_released_tensors=[fc1_output] - ) + output = moe_act_manager.group_offload( + output, + forced_released_tensors=[fc1_output], + delay_offload=self.config.delay_offload_until_cuda_graph, + ) output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs) # upad and concat the output diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index eb4e79a6c35..202034986db 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import math @@ -339,7 +339,8 @@ def forward( # Get the query, key and value tensors based on the type of attention - # self or cross attn. # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128] - with off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") as hidden_states: + qkv_linear_manager = off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") + with qkv_linear_manager as hidden_states: query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( hidden_states, key_value_states, @@ -347,10 +348,7 @@ def forward( packed_seq_params, inference_context=inference_context, ) - if self.offload_qkv_linear: - query = off_interface.group_commit( - query, name="qkv_linear", forced_released_tensors=[hidden_states] - ) + query = qkv_linear_manager.group_offload(query, forced_released_tensors=[hidden_states]) # =================================================== # Adjust key, value for inference @@ -378,6 +376,9 @@ def forward( # core attention computation # ================================== # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) needs_output_trim = False if self.checkpoint_core_attention and self.training: core_attn_out = self._checkpointed_attention_forward( @@ -390,9 +391,7 @@ def forward( ) else: if inference_context is None or inference_context.is_static_batching(): - with off_interface( - self.offload_core_attention and self.training, query, "core_attn" - ) as query: + with core_attn_manager as query: core_attn_out = self._run_core_attention( query, key, @@ -426,10 +425,9 @@ def forward( if not inference_context.is_decode_only(): core_attn_out = rearrange(core_attn_out, 's b h d -> s b (h d)') needs_output_trim = need_v_pad - if self.offload_core_attention and self.training: - core_attn_out = off_interface.group_commit( - core_attn_out, name="core_attn", forced_released_tensors=[query, key, value] - ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) # We are doing absorption with cache mla latents and decode mode. if self.cache_mla_latents and inference_context.is_decode_only(): @@ -460,12 +458,10 @@ def forward( # ================= # Output. [sq, b, h] # ================= - with off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") as core_attn_out: + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: output, bias = apply_module(self.linear_proj)(core_attn_out) - if self.offload_attn_proj: - output = off_interface.group_commit( - output, name="attn_proj", forced_released_tensors=[core_attn_out] - ) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) return output, bias diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 812470a73f4..bbcf413baee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math @@ -912,8 +912,8 @@ class TransformerConfig(ModelParallelConfig): advanced fused kernels.""" moe_expert_rank_capacity_factor: Optional[float] = None - """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens - exceeding this budget will be dropped. None means no token will be dropped. + """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens + exceeding this budget will be dropped. None means no token will be dropped. The default is None.""" ################## @@ -1180,6 +1180,24 @@ class TransformerConfig(ModelParallelConfig): min_offloaded_tensor_size: int = 1024 * 1024 """The minimum size of the tensor to be offloaded.""" + delay_offload_until_cuda_graph: bool = False + """If True, delay the offload until the CUDA graph is executed for minimal CPU overhead. + For more details, see the documentation: + https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#cuda-graph-integration. + """ + + delta_offload_bytes_across_pp_ranks: int = 0 + """Difference of offload bytes across PP ranks to balance the offload load. + For more details, see the documentation: + https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#tuning-parameters. + """ + + activation_offload_fraction: float = 1.0 + """Fraction of eligible activation offload groups to offload across configured modules. + For details, see: + https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#activation-offload-fraction. + """ + moe_paged_stash: bool = False """If True, enable paged stash for all routed-expert activations needed for backward""" @@ -1748,6 +1766,27 @@ def __post_init__(self): "because the input of attn_proj is the output of core_attn, " "which is needed in core_attn.backward()." ) + if self.recompute_granularity == "selective" and "moe" in self.recompute_modules: + offload_inside_moe = {"moe_act", "expert_fc1", "fused_group_mlp"} & set( + self.offload_modules + ) + assert not offload_inside_moe, ( + f"Cannot offload {offload_inside_moe} while recomputing the entire MoE layer. " + f"'moe' in recompute_modules wraps the full MoE forward in a checkpoint, " + f"so offloading activations inside it is redundant and will cause errors. " + f"Either remove 'moe' from --recompute-modules or remove " + f"{offload_inside_moe} from --offload-modules." + ) + assert ( + self.min_offloaded_tensor_size >= 0 + ), "min_offloaded_tensor_size must be non-negative." + assert ( + self.activation_offload_fraction >= 0 and self.activation_offload_fraction <= 1 + ), "activation_offload_fraction must be in range [0, 1]." + assert ( + self.delta_offload_bytes_across_pp_ranks >= 0 + ), "delta_offload_bytes_across_pp_ranks must be non-negative." + if "fused_group_mlp" in self.offload_modules: if not self.use_transformer_engine_op_fuser: raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.") @@ -2430,6 +2469,19 @@ def _scope_to_str(s): if self.fine_grained_activation_offloading: offload_modules = set(self.offload_modules or []) + if self.cuda_graph_impl == "local": + local_supported_offload_modules = {"expert_fc1", "moe_act", "fused_group_mlp"} + unsupported_offload_modules = offload_modules - local_supported_offload_modules + assert not unsupported_offload_modules, ( + "fine-grained activation offloading with cuda_graph_impl='local' " + "only supports offload_modules 'expert_fc1', 'moe_act', and " + "'fused_group_mlp'. " + f"Unsupported offload_modules: {sorted(unsupported_offload_modules)}." + ) + assert self.cuda_graph_modules, ( + "fine-grained activation offloading with cuda_graph_impl='local' " + "is not supported with whole-layer CUDA graph capture." + ) local_partial_moe_offload = ( self.cuda_graph_impl == "local" and bool(offload_modules) @@ -2442,7 +2494,7 @@ def _scope_to_str(s): ), ( "fine-grained activation offloading is only supported with " "transformer_engine CUDA graph implementation or local CUDA graph " - "implementation with full_iteration scope. Local partial CUDA graphs " + "implementation with partial MoE offload. Local partial CUDA graphs " "are supported only for expert_fc1, moe_act, or fused_group_mlp " "offload when the full MoE module is not captured." ) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..8aec8878b60 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import functools @@ -31,6 +31,7 @@ deprecate_inference_params, get_pg_rank, is_te_min_version, + is_torch_min_version, log_single_rank, make_viewless_tensor, nvtx_range_pop, @@ -43,6 +44,16 @@ logger = logging.getLogger(__name__) +def _get_offloading_interface(): + """Get the offloading interface for fine-grained activation offloading.""" + # Keep this import lazy to avoid a transformer/pipeline circular import. + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface, + ) + + return FineGrainedActivationOffloadingInterface + + def get_transformer_layer_offset( config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None ): @@ -501,17 +512,10 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph(): if "mlp" in self.config.recompute_modules: if not self.is_moe_layer: self.recompute_mlp = True - self.offload_attn_norm = ( - self.config.fine_grained_activation_offloading - and "attn_norm" in self.config.offload_modules - and not isinstance(self.input_layernorm, IdentityOp) - ) - self.offload_mlp_norm = ( - self.config.fine_grained_activation_offloading - and "mlp_norm" in self.config.offload_modules - and not isinstance(self.pre_mlp_layernorm, IdentityOp) - ) + self._set_offload_modules() + self.off_interface = _get_offloading_interface() + self.mlp_norm_manager = None # @jcasper how should we handle nvfuser? # Set bias+dropout+add fusion grad_enable execution handler. # TORCH_MAJOR = int(torch.__version__.split('.')[0]) @@ -607,21 +611,18 @@ def _forward_attention( context (Tensor): Updated context tensor if cross-attention is used, otherwise None. """ - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) - inference_context = deprecate_inference_params(inference_context, inference_params) # Optional Input Layer norm + attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") if self.recompute_input_layernorm: self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_attn_norm, hidden_states, "attn_norm") as hidden_states: + with attn_norm_manager as hidden_states: input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( apply_module(self.input_layernorm), hidden_states ) else: - with off_interface(self.offload_attn_norm, hidden_states, "attn_norm") as hidden_states: + with attn_norm_manager as hidden_states: input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) if isinstance(input_layernorm_output, tuple): @@ -687,10 +688,9 @@ def _forward_attention( # Delay the offload of the attention norm until after the self_attn_bda has been computed # because the residual is needed in the self_attn_bda. - if self.offload_attn_norm: - hidden_states = off_interface.group_commit( - hidden_states, name="attn_norm", forced_released_tensors=[residual] - ) + hidden_states = attn_norm_manager.group_offload( + hidden_states, forced_released_tensors=[residual] + ) # Optional Layer norm after self-attention pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states) @@ -746,18 +746,15 @@ def forward(self, *args, **kwargs): return output, context def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) - + self.mlp_norm_manager = self.off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") if self.recompute_pre_mlp_layernorm: self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + with self.mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( apply_module(self.pre_mlp_layernorm), hidden_states ) else: - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + with self.mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = apply_module(self.pre_mlp_layernorm)(hidden_states) return pre_mlp_layernorm_output @@ -904,9 +901,6 @@ def _forward_post_mlp( Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) using_fused_tp_inference_kernel = ( InferenceMode.is_active() and self.config.inference_fuse_tp_communication @@ -935,10 +929,11 @@ def _forward_post_mlp( nvtx_range_pop(suffix="mlp_bda") # Delay the offload of the mlp norm until after the mlp_bda has been computed # because the residual is needed in the mlp_bda. - if self.offload_mlp_norm: - hidden_states = off_interface.group_commit( - hidden_states, name="mlp_norm", forced_released_tensors=[residual] + if self.mlp_norm_manager is not None: + hidden_states = self.mlp_norm_manager.group_offload( + hidden_states, forced_released_tensors=[residual] ) + self.mlp_norm_manager = None # Jit compiled function creates 'view' tensor. This tensor # potentially gets saved in the MPU checkpoint function context, @@ -1094,6 +1089,18 @@ def _te_cuda_graph_capture(self, *args, **kwargs): attribute can be set to control the scope of the CUDA graph. 2. If context is None, it cannot be returned as output. """ + # Record the backward event on cuda graph stream in backward pass. + # This is to ensure the main stream waits for computing on cuda graph stream to complete, + # and overlaps with the H2D transfer on reload stream. + if self.offload_module_in_cuda_graph: + if len(args) > 0: + hidden_states = args[0] + hidden_states = self.off_interface.backward_record(hidden_states) + args = (hidden_states,) + args[1:] + else: + hidden_states = kwargs.pop("hidden_states") + hidden_states = self.off_interface.backward_record(hidden_states) + kwargs["hidden_states"] = hidden_states context = None if ( not self.config.cuda_graph_modules @@ -1124,6 +1131,11 @@ def _te_cuda_graph_capture(self, *args, **kwargs): cuda_graph_outputs = list(hidden_states) if context is not None: cuda_graph_outputs.append(context) + # Record the forward event on cuda graph stream for cuda graph capture. + # This is to ensure the main stream waits for computing on cuda graph stream to complete, + # and overlaps with the D2H transfer on offloading stream. + if self.offload_module_in_cuda_graph: + self.off_interface.forward_record() return tuple(cuda_graph_outputs) def _te_cuda_graph_replay(self, *args, **kwargs): @@ -1150,8 +1162,25 @@ def _te_cuda_graph_replay(self, *args, **kwargs): "For inference cuda graph, please use cuda_graph_impl=local instead." ) + if self.config.delay_offload_until_cuda_graph: + self.off_interface.enter_replay() + + try: + return self._te_cuda_graph_replay_impl(args, kwargs, context) + finally: + if self.config.delay_offload_until_cuda_graph: + self.off_interface.exit_replay() + + def _te_cuda_graph_replay_impl(self, args, kwargs, context): + """Implementation of _te_cuda_graph_replay, separated for replay mode cleanup.""" cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) + # Flush delayed offload groups from previous layers after graph replay. + # The CPU is idle during the sync between graph replay and a2a comm, + # so we use that time to execute the delayed offload operations. + if self.config.delay_offload_until_cuda_graph: + self.off_interface.flush_delayed_groups() + if kwargs.get('context') is not None: context = cuda_graph_output.pop() @@ -1341,6 +1370,81 @@ def _should_call_local_cudagraph(self, *args, **kwargs): return True return False + def _set_offload_modules(self): + """Set the offload modules for the transformer layer.""" + if self.config.fine_grained_activation_offloading: + self.offload_attn_norm = "attn_norm" in self.config.offload_modules and not isinstance( + self.input_layernorm, IdentityOp + ) + self.offload_qkv_linear = "qkv_linear" in self.config.offload_modules + self.offload_core_attn = "core_attn" in self.config.offload_modules + self.offload_attn_proj = "attn_proj" in self.config.offload_modules + self.offload_mlp_norm = "mlp_norm" in self.config.offload_modules and not isinstance( + self.pre_mlp_layernorm, IdentityOp + ) + else: + self.offload_attn_norm = False + self.offload_qkv_linear = False + self.offload_core_attn = False + self.offload_attn_proj = False + self.offload_mlp_norm = False + # Check the compatibility of fine-grained activation offloading and cuda graph. + if self.config.fine_grained_activation_offloading: + cuda_graph_modules = self.config.cuda_graph_modules or [] + if CudaGraphModule.attn in cuda_graph_modules: + self.offload_attn_norm = False + log_single_rank( + logger, + logging.WARNING, + "attn_norm offloading is not supported with attn cudagraph. " + "Disabling attn_norm offloading.", + ) + mark_mlp_norm_offloading_not_supported = False + # For moe layer, mlp_norm offloading isn't supported with attn or moe_router cudagraph. + if self.is_moe_layer: + if ( + CudaGraphModule.attn in cuda_graph_modules + or CudaGraphModule.moe_router in cuda_graph_modules + ): + mark_mlp_norm_offloading_not_supported = True + # For non-moe layer, mlp_norm is the boundary of attn or mlp cudagraph. + # The only case where mlp_norm offloading is supported is when whole layer is captured. + elif ( + CudaGraphModule.attn in cuda_graph_modules + and CudaGraphModule.mlp not in cuda_graph_modules + ) or ( + CudaGraphModule.attn not in cuda_graph_modules + and CudaGraphModule.mlp in cuda_graph_modules + ): + mark_mlp_norm_offloading_not_supported = True + if mark_mlp_norm_offloading_not_supported: + self.offload_mlp_norm = False + log_single_rank( + logger, + logging.WARNING, + "mlp_norm offloading is not supported with the current cudagraph scope. " + "Disabling mlp_norm offloading.", + ) + # Set the offload module in cuda graph flag. + self.offload_module_in_cuda_graph = False + cuda_graph_modules = self.config.cuda_graph_modules or [] + if CudaGraphModule.attn in cuda_graph_modules: + if self.offload_core_attn or self.offload_attn_proj or self.offload_qkv_linear: + self.offload_module_in_cuda_graph = True + if not self.is_moe_layer and CudaGraphModule.mlp in cuda_graph_modules: + if self.offload_mlp_norm: + self.offload_module_in_cuda_graph = True + if self.offload_module_in_cuda_graph: + assert is_torch_min_version( + "2.9.0a0" + ), "Offloading modules captured in cuda graph requires torch>=2.9.0." + assert is_te_min_version( + "2.14.0" + ), "Offloading modules captured in cuda graph requires TE>=2.14.0." + assert ( + self.config.cuda_graph_warmup_steps > 0 + ), "Fine-grained activation offloading needs cuda_graph_warmup_steps > 0." + def get_layer_norm_weights(self): """ Get the weights of all layernorms (attention and MLP) in the transformer layer. diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 55cb7a4e7d6..49aebe43b3f 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import hashlib import inspect @@ -310,6 +310,9 @@ "fine_grained_activation_offloading": False, "min_offloaded_tensor_size": 1024 * 1024, "offload_modules": [], + "delay_offload_until_cuda_graph": False, + "delta_offload_bytes_across_pp_ranks": 0, + "activation_offload_fraction": 1.0, "fine_grained_offloading_max_inflight_offloads": None, "hybrid_context_parallel": False, "max_seqlen_per_dp_cp_rank": None, diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index 2af0728a7fb..dce35bdd4fa 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -65,7 +65,7 @@ def test_chunk_offload_handler_skips_non_offloadable_tensor_types(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") -def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out(): +def test_chunk_offload_handler_respects_tensor_opt_out_flags(): handler = _make_chunk_handler_for_offload_checker() tensor = torch.empty(1024, device="cuda") @@ -74,10 +74,6 @@ def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out(): tensor._TE_do_not_offload = True assert not handler.tensor_need_offloading_checker(tensor) - tensor = torch.empty(1024, device="cuda") - tensor.offloading_activation = False - assert not handler.tensor_need_offloading_checker(tensor) - def _build_gpt_model( *, @@ -383,7 +379,6 @@ def test_gpt_fine_grained_activation_offloading_correctness_and_memory( ("alltoall", True, ["mlp_norm"]), ("alltoall", False, ["expert_fc1"]), ("alltoall", False, ["moe_act"]), - ("alltoall", False, ["mlp_norm", "expert_fc1", "moe_act"]), ( "alltoall", True, @@ -639,3 +634,337 @@ def _run_schedule_1f1b_two_microbatches( ) finally: Utils.destroy_model_parallel() + + +# ============================================================================= +# CUDA Graph + Fine-grained Activation Offloading Tests +# ============================================================================= + + +def _build_gpt_model_with_cuda_graph( + *, + seed: int, + num_layers: int, + hidden_size: int, + num_attention_heads: int, + vocab_size: int, + seq_length: int, + num_experts: Optional[int], + fine_grained_activation_offloading: bool, + offload_modules: Optional[List[str]], + min_offloaded_tensor_size: int, + is_mla: bool, + cuda_graph_impl: str, + cuda_graph_scope: Optional[List[str]], + cuda_graph_warmup_steps: int, + delay_offload_until_cuda_graph: bool = False, + activation_offload_fraction: float = 1.0, +) -> GPTModel: + """Build a GPTModel with CUDA Graph support and fine-grained activation offloading.""" + model_parallel_cuda_manual_seed(seed) + torch.manual_seed(seed) + ConfigClass = MLATransformerConfig if is_mla else TransformerConfig + transformer_config = ConfigClass( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + attention_backend=AttnBackend.unfused, + bf16=True, + # Recompute + recompute_modules=["layernorm", "moe_act"] if num_experts is not None else ["layernorm"], + recompute_granularity="selective", + # MoE + num_moe_experts=num_experts, + moe_grouped_gemm=(num_experts is not None), + # Fine-grained activation offloading + fine_grained_activation_offloading=fine_grained_activation_offloading, + offload_modules=offload_modules, + min_offloaded_tensor_size=min_offloaded_tensor_size, + delay_offload_until_cuda_graph=delay_offload_until_cuda_graph, + activation_offload_fraction=activation_offload_fraction, + # CUDA Graph settings + cuda_graph_impl=cuda_graph_impl, + cuda_graph_scope=cuda_graph_scope, + cuda_graph_warmup_steps=cuda_graph_warmup_steps, + use_te_rng_tracker=True, + ) + gpt_model = GPTModel( + config=transformer_config, + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec( + num_experts=num_experts, + moe_grouped_gemm=num_experts is not None, + multi_latent_attention=is_mla, + ), + vocab_size=vocab_size, + max_sequence_length=seq_length, + ).bfloat16() + return gpt_model + + +def _run_iters_with_cuda_graph( + model: GPTModel, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + num_warmup_iters: int, + num_measure_iters: int, + enable_offload_reset: bool, +) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], int]: + """ + Run multiple forward+backward iterations with CUDA graph capture. + + Returns: + - logits from last iteration (CPU float32) + - selected grads from last iteration (CPU float32) + - peak_memory_allocated (bytes) during measurement iterations + """ + from megatron.core.transformer.cuda_graphs import _CudagraphGlobalRecord, delete_cuda_graphs + + if enable_offload_reset: + off_interface.reset() + + # Warmup iterations (before CUDA graph capture) + for _ in range(num_warmup_iters): + if enable_offload_reset: + off_interface.reset() + logits = model( + input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask + ) + loss = logits.float().sum() + loss.backward() + # Zero grads for next iteration + for p in model.parameters(): + if p.grad is not None: + p.grad.zero_() + + # Trigger post-warmup offload decisions + if enable_offload_reset: + off_interface.reset() + + # Create CUDA graphs after warmup + _CudagraphGlobalRecord.create_cudagraphs() + + # Measurement iterations (with CUDA graph replay) + torch.cuda.reset_peak_memory_stats() + for i in range(num_measure_iters): + if enable_offload_reset: + off_interface.reset() + logits = model( + input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask + ) + loss = logits.float().sum() + loss.backward() + if i < num_measure_iters - 1: + for p in model.parameters(): + if p.grad is not None: + p.grad.zero_() + + torch.cuda.synchronize() + peak_bytes = int(torch.cuda.max_memory_allocated()) + + # Capture grads from last iteration + grads: Dict[str, torch.Tensor] = {} + for name, p in model.named_parameters(): + grads[name] = p.grad.detach().float().cpu() if p.grad is not None else None + + # Cleanup CUDA graphs + delete_cuda_graphs() + + return logits.detach().float().cpu(), grads, peak_bytes + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offloading tests.") +@pytest.mark.skipif( + not is_te_min_version("2.14.0"), reason="CUDA Graph with TE RNG tracker requires TE >= 2.13.0" +) +@pytest.mark.parametrize( + "is_mla, offload_modules, cuda_graph_scope, activation_offload_fraction, delay_offload", + [ + # MoE model with attention CUDA graph + attn offloading + (False, ["core_attn", "attn_proj"], ["attn", "moe_router"], 1.0, True), + (False, ["expert_fc1", "moe_act"], ["attn", "moe_router", "moe_preprocess"], 1.0, True), + (False, ["core_attn", "attn_proj", "expert_fc1"], ["attn", "moe_router"], 1.0, True), + ( + False, + ["core_attn", "attn_proj", "expert_fc1", "moe_act"], + ["attn", "moe_router"], + 1.0, + True, + ), + ( + False, + ["core_attn", "expert_fc1", "moe_act"], + ["attn", "moe_router", "moe_preprocess"], + 1.0, + True, + ), + ( + True, + ["core_attn", "attn_proj", "expert_fc1", "moe_act"], + ["attn", "moe_router", "moe_preprocess"], + 1.0, + True, + ), + # Test activation_offload_fraction parameter + (False, ["core_attn", "attn_proj", "expert_fc1"], ["attn", "moe_router"], 0.0, True), + (False, ["core_attn", "attn_proj", "expert_fc1"], ["attn", "moe_router"], 0.5, True), + # Test delay_offload_until_cuda_graph parameter + (False, ["core_attn", "attn_proj", "expert_fc1"], ["attn", "moe_router"], 1.0, False), + ], +) +def test_fine_grained_activation_offloading_with_cuda_graph( + is_mla: bool, + offload_modules: List[str], + cuda_graph_scope: List[str], + activation_offload_fraction: float, + delay_offload: bool, +): + """ + Test fine-grained activation offloading combined with CUDA graph capture. + + Verifies: + - Forward output correctness with CUDA graph + offloading + - Backward gradient correctness + - Memory savings from offloading are preserved with CUDA graphs + - Different activation_offload_fraction values work correctly + - Both delay_offload_until_cuda_graph=True/False produce correct results + """ + from megatron.core.tensor_parallel.random import initialize_rng_tracker + + os.environ.pop("NVTE_FUSED_ATTN", None) + os.environ.pop("NVTE_FLASH_ATTN", None) + os.environ.pop("NVTE_UNFUSED_ATTN", None) + + initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + Utils.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + + seed = 123 + num_experts = 4 # Always MoE model + num_layers = 4 # Smaller for faster test with CUDA graphs + hidden_size = 1024 + num_attention_heads = 8 + vocab_size = 512 + seq_length = 512 + micro_batch_size = 2 + device = torch.device("cuda") + cuda_graph_warmup_steps = 3 + + input_ids, position_ids, attention_mask = _make_gpt_inputs( + seq_length=seq_length, micro_batch_size=micro_batch_size, device=device + ) + + off_interface.reset_instance() + + try: + # 1) Baseline: CUDA graph enabled, offloading disabled + _reset_cuda_memory() + base_model = _build_gpt_model_with_cuda_graph( + seed=seed, + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + vocab_size=vocab_size, + seq_length=seq_length, + num_experts=num_experts, + fine_grained_activation_offloading=False, + offload_modules=None, + min_offloaded_tensor_size=1024 * 1024, + is_mla=is_mla, + cuda_graph_impl="transformer_engine", + cuda_graph_scope=cuda_graph_scope, + cuda_graph_warmup_steps=cuda_graph_warmup_steps, + ).cuda() + base_model.train() + + base_logits, base_grads, base_peak = _run_iters_with_cuda_graph( + base_model, + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + num_warmup_iters=cuda_graph_warmup_steps, + num_measure_iters=2, + enable_offload_reset=False, + ) + del base_model + _reset_cuda_memory() + + # 2) Test: CUDA graph enabled + offloading enabled + off_interface.reset_instance() + + off_model = _build_gpt_model_with_cuda_graph( + seed=seed, + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + vocab_size=vocab_size, + seq_length=seq_length, + num_experts=num_experts, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + min_offloaded_tensor_size=1024, # Force offloading for determinism + is_mla=is_mla, + cuda_graph_impl="transformer_engine", + cuda_graph_scope=cuda_graph_scope, + cuda_graph_warmup_steps=cuda_graph_warmup_steps, + delay_offload_until_cuda_graph=delay_offload, + activation_offload_fraction=activation_offload_fraction, + ).cuda() + off_model.train() + + off_logits, off_grads, off_peak = _run_iters_with_cuda_graph( + off_model, + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + num_warmup_iters=cuda_graph_warmup_steps, + num_measure_iters=2, + enable_offload_reset=True, + ) + del off_model + _reset_cuda_memory() + + # 3) Correctness checks + assert torch.allclose( + off_logits, base_logits, rtol=1e-2, atol=1e-2 + ), f"Logits mismatch: max_diff={torch.max(torch.abs(off_logits - base_logits))}" + assert set(off_grads.keys()) == set(base_grads.keys()) + for name, gb in base_grads.items(): + go = off_grads[name] + if gb is None or go is None: + assert gb is None and go is None, f"Grad None mismatch for {name}" + continue + assert torch.allclose( + go, gb, rtol=1e-2, atol=1e-2 + ), f"Grad mismatch for {name}: max_diff={torch.max(torch.abs(go - gb))}" + + # 4) Memory checks - offloading should still reduce memory with CUDA graphs + saved_mib = (base_peak - off_peak) / (1024**2) + print( + f"CUDA Graph + Offload test (fraction={activation_offload_fraction}, delay={delay_offload}): " + f"base_peak={base_peak/(1024**2):.2f}MiB, " + f"off_peak={off_peak/(1024**2):.2f}MiB, " + f"saved={saved_mib:.2f}MiB" + ) + + # Basic sanity checks + assert not torch.isnan(off_logits).any(), "NaN detected in logits" + assert not torch.isinf(off_logits).any(), "Inf detected in logits" + + # Check gradients are valid + for name, g in off_grads.items(): + if g is not None: + assert not torch.isnan(g).any(), f"NaN detected in grad for {name}" + assert not torch.isinf(g).any(), f"Inf detected in grad for {name}" + + # Note: With CUDA graphs, memory behavior may differ from eager mode. + # We check that offloading doesn't significantly increase memory. + # In some cases, graph capture overhead may offset offload savings. + assert saved_mib >= -DELTA, ( + f"Offloading with CUDA graph significantly increased memory: " + f"saved={saved_mib:.2f}MiB (negative means increase)" + ) + + finally: + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 31d7727c1eb..726507ea4ef 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc import os @@ -98,6 +98,62 @@ def test_local_impl_defaults_to_layer_scope(self): cfg = _base_cuda_graph_config(cuda_graph_impl='local') assert cfg.inference_cuda_graph_scope == InferenceCudaGraphScope.layer + def test_local_impl_allows_expert_activation_offload_scope(self): + cfg = _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[CudaGraphModule.attn, CudaGraphModule.moe_router], + fine_grained_activation_offloading=True, + offload_modules=['expert_fc1', 'moe_act'], + num_moe_experts=4, + ) + + assert cfg.cuda_graph_impl == 'local' + assert CudaGraphModule.attn in cfg.cuda_graph_modules + assert CudaGraphModule.moe_router in cfg.cuda_graph_modules + assert CudaGraphModule.moe_preprocess in cfg.cuda_graph_modules + + def test_local_impl_rejects_unsupported_activation_offload_scope(self): + with pytest.raises( + AssertionError, + match=( + "fine-grained activation offloading with cuda_graph_impl='local'.*" + "Unsupported offload_modules: \\['qkv_linear'\\]" + ), + ): + _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[CudaGraphModule.attn], + fine_grained_activation_offloading=True, + offload_modules=['qkv_linear'], + ) + + def test_local_impl_rejects_full_layer_graph_with_activation_offload(self): + with pytest.raises( + AssertionError, match="not supported with whole-layer CUDA graph capture" + ): + _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[], + fine_grained_activation_offloading=True, + offload_modules=['expert_fc1'], + ) + + def test_local_impl_rejects_moe_router_graph_with_mlp_norm_offload(self): + with pytest.raises( + AssertionError, + match=( + "fine-grained activation offloading with cuda_graph_impl='local'.*" + "Unsupported offload_modules: \\['mlp_norm'\\]" + ), + ): + _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[CudaGraphModule.moe_router], + fine_grained_activation_offloading=True, + offload_modules=['mlp_norm'], + num_moe_experts=4, + ) + def test_full_iteration_impl_requires_empty_scope(self): with pytest.raises( AssertionError,