What does this PR do ?

PR to main branch

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: 18 files changed, +944 / -259 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:

• CudaGraphScope.attn sub-graph containing offloaded qkv_linear / core_attn / attn_proj
• CudaGraphScope.mlp sub-graph (dense layers) containing offloaded mlp_norm
• Incompatible combinations (e.g., attn sub-graph + attn_norm offload) are auto-disabled with warnings

5. New Configuration Options (TransformerConfig)

6. Code Cleanup

Replaced redundant from ... import FineGrainedActivationOffloadingInterface scattered across methods with a single _get_offloading_interface() helper (@lru_cache), accessed uniformly via self.off_interface.

Execution Flow

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.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

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.

Merging your PR

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