[Main][feat] Support CUDA Graph capture offloading modules by lhb8125 · Pull Request #3697 · NVIDIA/Megatron-LM

lhb8125 · 2026-03-04T13:17:30Z

What does this PR do ?

This PR enables Fine-Grained Activation Offloading to work seamlessly with Transformer Engine CUDA Graph capture and replay. Previously, these two features were mutually exclusive — CUDA Graph captures a fixed sequence of GPU operations, while activation offloading involves dynamic D2H/H2D memory copies that conflict with graph semantics. This PR resolves the conflict by introducing dedicated CUDA stream/event synchronization and an optional deferred-commit strategy.
Scope: 15 files changed, +1033 / -212 lines

Key Changes

1. In-Graph Offload Synchronization (`transformer_layer.py`)

_te_cuda_graph_capture(): When offload_module_in_cuda_graph=True, inserts backward_record() at the sub-graph entry (synchronizes the compute stream with the H2D reload stream during backward) and calls forward_record() at the sub-graph exit (synchronizes the compute stream with the D2H offload stream during forward).
_te_cuda_graph_replay(): Supports the delay_offload_until_cuda_graph mode — during replay, enter_replay() / exit_replay() cause offload groups to be enqueued without immediate execution; after replay, flush_delayed_groups() issues the D2H copies during the CPU-idle window between graph launch and subsequent communication.

2. `PipelineOffloadManager` Extensions (`fine_grained_activation_offload.py`)

Added cuda_graph_stream / cuda_graph_event (external event) dedicated to synchronizing in-graph captured modules with the D2H/H2D offload streams.
Deferred offload commit: FineGrainedOffloadingGroupCommitFunction pushes offload groups into a queue when delay_offload=True and the manager is in replay state; flush_delayed_groups() drains the queue in batch during CPU-idle gaps.
Warmup hook integration: pre_warmup_hook / post_warmup_hook temporarily disable/enable offloading around TE's warmup phase to avoid state-machine conflicts.

3. `GraphableMegatronModule` Integration (`module.py`)

When fine_grained_activation_offloading and offload_module_in_cuda_graph are both active, _get_te_cuda_graph_replay_args() injects cuda_graph_stream and cuda_graph_event into TE's replay kwargs, bridging the TE-side synchronization.

4. Automatic Offload-in-Graph Detection (`_set_offload_modules`)

Added the offload_module_in_cuda_graph flag, automatically determined by:

CudaGraphModule.attn sub-graph containing offloaded qkv_linear / core_attn / attn_proj
CudaGraphModule.mlp sub-graph (dense layers) containing offloaded mlp_norm
Incompatible combinations (e.g., attn sub-graph + attn_norm offload, or unsupported mlp_norm graph boundaries) are auto-disabled with warnings
Local CUDA graph activation offloading is limited to routed expert scopes (expert_fc1 / moe_act)

5. New Configuration Options (`TransformerConfig`)

Config	Description
`delay_offload_until_cuda_graph`	Defer offload commits until after CUDA Graph replay to minimize CPU overhead
`activation_offload_fraction`	Fraction of activations to offload, range [0, 1]
`delta_offload_bytes_across_pp_ranks`	Differential offload bytes across PP ranks
Validation: fine-grained activation offloading is supported with `cuda_graph_impl="transformer_engine"`, `"local"`, or `"full_iteration"`. The local implementation is limited to `expert_fc1` / `moe_act` offload scopes and does not support whole-layer capture; `full_iteration` requires `fine_grained_offloading_max_inflight_offloads >= 0`. `CudaGraphModule.moe` is temporarily unsupported; `cuda_graph_warmup_steps > 0` is required. Fine-grained activation offloading remains mutually exclusive with `cpu_offloading` and `mhc` recompute.

6. Code Cleanup and Compatibility

Replaced repeated method-local imports in transformer_layer.py with _get_offloading_interface() (@lru_cache) and self.off_interface.
Updated the offload context manager API from a static group_commit(...) call to per-group manager objects (group_offload(...)) so callers can carry group identity through delayed commits.
Treat MXFP8 and NVFP4 similarly for expert_fc1 input saving, avoiding extra set_save_original_input(self.linear_fc1) setup for those quantized recipes.

Execution Flow

Phase	Behavior
Warmup	Offloading is temporarily disabled via hooks during TE's multi-step eager warmup before graph construction
Capture	`backward_record` + `forward_record` link the compute stream with D2H/H2D streams through a shared `cuda_graph_event`
Replay	(Optional deferred mode) `group_commit` only enqueues; `flush_delayed_groups()` issues D2H during the CPU-idle window between graph launch and subsequent communication
Iteration End	`off_interface.reset()` in `schedules.py` and `cuda_graphs._finish_capturing` ensures clean state

Tests

test_fine_grained_activation_offloading_with_cuda_graph: covers multiple combinations of cuda_graph_scope, offload_modules, activation_offload_fraction, and delay_offload_until_cuda_graph (True/False); validates numerical correctness against baseline logits/gradients and performs peak memory sanity checks. Requires TE >= 2.14.
TestCudaGraphConfigAndArguments local-scope validation: covers allowed local CUDA graph activation-offload scopes, rejects unsupported offload modules, and rejects local whole-layer capture with fine-grained activation offloading.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

copy-pr-bot · 2026-03-04T13:17:35Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

lhb8125 · 2026-03-04T13:29:35Z

/ok to test 8a889c0

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

…_graph_main

lhb8125 · 2026-03-04T13:42:39Z

/ok to test ecaa319

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

lhb8125 · 2026-03-04T13:45:47Z

/ok to test 5591584

Phlip79 · 2026-03-04T23:46:11Z

We are changing our review process and marking all open, unlabeled PRs as draft. This change will go in effect starting once #3659 is merged.

Moving forward, all PRs will be required to start as draft PRs. If you wish to get your PR merged, mark your PR as “Ready for review”. Read more about the new process at submit.md.

2. remove flush_delayed_groups() when the training is not in replay mode Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>