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

lhb8125 · 2026-02-03T04:36:38Z

What does this PR do ?

⚠️ 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.

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>

…ithub.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_cuda_graph

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

2. refine offloading docs; 3. remove TransformerLayer.cuda_graph_stream and cuda_graph_event Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

…graph

lhb8125 · 2026-03-02T03:51:08Z

/ok to test 4bf0085

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

…ttps://github.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_refactor_cuda_graph

lhb8125 · 2026-03-02T04:31:43Z

/ok to test cd84623

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

lhb8125 · 2026-03-02T07:23:10Z

/ok to test 61b589a

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

lhb8125 · 2026-03-02T09:26:26Z

/ok to test 0200121

buptzyb · 2026-03-04T08:47:28Z

docs/api-guide/fine_grained_activation_offloading.md

+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`.


Can I use "moe" scope if I'm in a drop-pad MoE?

Can I offload attention part modules if my cuda graph scope is only "moe_router"? This may be needed since some cases have dynamic-shaped attention so only the router part can be captured.

I removed this hard limitation, now the scope could be moe_router alone or moe.

buptzyb · 2026-03-04T08:54:45Z

docs/api-guide/fine_grained_activation_offloading.md

+
+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).


unless using "moe" cudagrpah scope in a drop-pad or sync-free MoE.

what if we only capture moe_router or moe_preprocess? Is it still true?

I think so. If we only capture moe_router, mlp_norm works as the input buffer of the graph, so not offloadable. The only exception is that we use attn+moe scope for drop-pad MoE, then the mlp_norm is totally inside the graph, so offloadable.

btw you cannot only capture moe_preprocess . moe_preprocess must go together with moe_router .

jiemingz · 2026-03-04T16:02:57Z

megatron/core/transformer/transformer_layer.py

        hidden_states: Tensor,
        inference_context: BaseInferenceContext | None = None,
        padding_mask: Tensor | None = None,
+        flush_delayed_groups: bool = True,


since flush_delayed_groups is cudagraph specific can it be moved to cudagraph specific code? If this just needs to run after warmup can it be passed as TE's make_graphed_callable(post_warmup_hook=) ?

Thanks for the comments, I removed the function call in forward() and kept only the call in _te_cuda_graph_replay. Now in the warmup iterations, we launch the offloading immediately and in the replay iterations, we delay the offloading and flush it after graph replay.

The "warmup" here is a little vague:

The first several iterations are warmup iterations, after which we start graph capturing.

Before capturing cuda graph, TE runs several fprop&bprop steps to "warmup";

In previous code, the flush_delayed_groups is executed in the end of forward() in the warmup iterations(1st case).

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>

…graph

lhb8125 · 2026-03-05T09:06:42Z

/ok to test b481fa9

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

…ttps://github.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_refactor_cuda_graph

…graph

lhb8125 · 2026-03-05T13:21:16Z

/ok to test ce84682

buptzyb · 2026-03-05T13:37:58Z

docs/api-guide/fine_grained_activation_offloading.md

+3. **Apply fraction**: Only a fraction of eligible groups are actually offloaded (controlled by `activation_offload_fraction`).
+4. **Print summary table**: An ASCII table of per-rank offload bytes is printed for debugging.
+
+### CPU Tensor Pool


GPU Tensor Pool?

buptzyb · 2026-03-05T13:40:30Z

docs/api-guide/fine_grained_activation_offloading.md

+
+### 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:


So we cannot capture cudagraphs on the first training iteration? If so, we should assert cuda_graph_warmup_steps>0 when offloading is enabled.

lhb8125 and others added 30 commits October 29, 2025 02:46

renaming golden values

1219a26

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

fix bug: accuracy issu because of recomputing and offloading same module

ce6e661

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

Merge branch 'dev' into hongbinl/activation_offloading_fix

d04d741

format

2fe4aeb

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

update golden values

fb3f7c3

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

Merge branch 'dev' into hongbinl/activation_offloading_fix

5001e2b

update golden values

9937890

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

update model_config and golden values

6c83118

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

format

33a38f5

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

update golden values

6c76b07

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

Merge branch 'dev' into hongbinl/activation_offloading_fix

4d83f69

temp save

8e72b44

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

support offloading+cuda graph

1646f04

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

Merge branch 'dev' into hongbinl/activation_offloading_cuda_graph

43973a7

support PP=1

a177cf5

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

support VPP

f7cfbba

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

bug fix

6d475ad

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

support VPP

089da6c

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

code refactor

35b0f97

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

big code refactor and format

df09b85

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

Merge branch 'dev' into hongbinl/activation_offloading_cuda_graph

06ef4e2

minor fix

12cb8de

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

minor fix

3cf19b7

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

dump offloading information

d0fc888

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

Merge branch 'dev' into hongbinl/activation_offloading_cuda_graph

bc47650

fix ut

b7c0fba

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

Merge branch 'hongbinl/activation_offloading_cuda_graph' of https://g…

b18e69b

…ithub.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_cuda_graph

format

ae4e2b5

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

fit ut

b797438

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

delay d2h copies until finishing cuda graph

6cec22f

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

copy-pr-bot bot temporarily deployed to test February 26, 2026 10:16 Inactive

1. replace hasattr+delattr with None;

998d1b0

2. refine offloading docs; 3. remove TransformerLayer.cuda_graph_stream and cuda_graph_event Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

lhb8125 force-pushed the hongbinl/activation_offloading_refactor_cuda_graph branch from d34be1d to 998d1b0 Compare March 2, 2026 03:50

Merge branch 'dev' into hongbinl/activation_offloading_refactor_cuda_…

4bf0085

…graph

lhb8125 added 2 commits March 1, 2026 20:30

format

716e12a

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

Merge branch 'hongbinl/activation_offloading_refactor_cuda_graph' of h…

cd84623

…ttps://github.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_refactor_cuda_graph

copy-pr-bot bot temporarily deployed to test March 2, 2026 04:32 Inactive

bug fix

61b589a

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

copy-pr-bot bot temporarily deployed to test March 2, 2026 07:24 Inactive

add flag to control flush_delayed_groups in fine_grained_callables.py

0200121

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

copy-pr-bot bot temporarily deployed to test March 2, 2026 09:27 Inactive

buptzyb reviewed Mar 4, 2026

View reviewed changes

lhb8125 changed the title ~~Support CUDA Graph capture offloading modules~~ [Dev][feat] Support CUDA Graph capture offloading modules Mar 4, 2026

jiemingz reviewed Mar 4, 2026

View reviewed changes

lhb8125 and others added 4 commits March 5, 2026 01:05

1. move backward_record() to te_cuda_graph_capture()

c8bd90d

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

format

ddd67d2

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

remove the knob forward_only when executing reset()

e989b95

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

Merge branch 'dev' into hongbinl/activation_offloading_refactor_cuda_…

b481fa9

…graph

copy-pr-bot bot temporarily deployed to test March 5, 2026 09:07 Inactive

lhb8125 added 3 commits March 5, 2026 05:18

fix ut and reviewer's comments

19fe6b3

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

Merge branch 'hongbinl/activation_offloading_refactor_cuda_graph' of h…

cf04b4a

…ttps://github.com/lhb8125/Megatron-LM into hongbinl/activation_offloading_refactor_cuda_graph

Merge branch 'dev' into hongbinl/activation_offloading_refactor_cuda_…

ce84682

…graph

copy-pr-bot bot deployed to test March 5, 2026 13:22 Active

buptzyb reviewed Mar 5, 2026

View reviewed changes


		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).


		### 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:

Conversation

lhb8125 commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

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

Merging your PR

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhb8125 Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhb8125 commented Mar 5, 2026

Uh oh!

lhb8125 commented Mar 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lhb8125 commented Feb 3, 2026 •

edited

Loading

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

lhb8125 Mar 5, 2026 •

edited

Loading