support cuda graph capture offloading module - #2435

lhb8125 · 2025-12-01T05:34:12Z

Description

This PR supports offloading modules captured by partial cuda graph in Megatron-LM.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Megatron's offloading relys on cpu_offload_v1 but we need to reuse mark_not_offload() so that the weights won't be offloaded.
Do not offload the output of core_attention
Refine the allocation strategy of fp8&fp4 tensors because the previous impl allocates tensor by from_blob(), which is not compatitable with record_stream()
Add pre_warmup_hook and post_warmup_hook
Passing cuda_graph_stream and cuda_graph_event to user_kwargs so that the fwd&bwd replay runs at a side stream, where the cuda_graph_event records on current stream after finishing computing.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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

for more information, see https://pre-commit.ci

greptile-apps · 2025-12-01T05:39:26Z

Greptile Summary

Adds CUDA graph capture offloading support for Megatron-LM integration. The PR introduces cuda_graph_stream and cuda_graph_event parameters to enable graph replay on side streams, allowing overlap with main stream work. It also adds pre_warmup_hook and post_warmup_hook callbacks for setup/teardown during graph capture warmup.

Key changes:

Modified autograd function signatures to accept cuda_graph_stream and cuda_graph_event parameters
Added conditional stream/event synchronization logic in forward and backward passes
Hooks are called once before/after warmup loop (not on each iteration)
Updated backward return signature to match new forward parameters

Previously raised concerns that remain valid:

cuda_graph_event is silently ignored if cuda_graph_stream is not different from current stream, which could lead to synchronization bugs
Stream/event extraction from user_kwargs uses pop(), though this is intentional to prevent passing them to underlying functions

The implementation follows the stated goal of supporting offloading with partial CUDA graphs, and the synchronization logic appears sound when parameters are used as intended.

Confidence Score: 3/5

Acceptable with some concerns - the stream/event synchronization logic needs validation improvements
The implementation adds complex CUDA graph stream/event handling that appears logically sound, but lacks validation for edge cases (e.g., event provided without different stream). Previous review threads identified legitimate concerns about synchronization. The core functionality follows established patterns for CUDA graph capture, and documentation has been added for new parameters.
Pay attention to transformer_engine/pytorch/graph.py lines 869-878 (stream/event parameter extraction) and lines 804-812, 830-839 (stream/event synchronization logic)

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Adds CUDA graph capture offloading support with custom stream/event handling, pre/post warmup hooks, and modified forward/backward signatures to support side-stream execution

_{Last reviewed commit: 5115a1c}

greptile-apps

_{1 file reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

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

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

Signed-off-by: root <root@eos0046.eos.clusters.nvidia.com>

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

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

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

for more information, see https://pre-commit.ci

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

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

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

for more information, see https://pre-commit.ci

lhb8125 · 2026-02-05T03:26:46Z

/te-ci pytorch L1

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

lhb8125 · 2026-02-05T05:29:26Z

@buptzyb @zhongbozhu @pggPL Could you review this PR?

lhb8125 · 2026-02-05T05:30:46Z

        """
        return self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute()

+    def trigger_backward_dw(self):


Please ignore this method, which will be removed after https://github.com/NVIDIA/TransformerEngine/pull/2614/files merged

lhb8125 · 2026-02-05T05:31:04Z

                bwd_dw_graphs[graph_idx].replay()
+                for module in te_modules:
+                    if hasattr(module, "trigger_backward_dw"):
+                        module.trigger_backward_dw()


Please ignore this code block, which will be removed after https://github.com/NVIDIA/TransformerEngine/pull/2614/files are merged

buptzyb · 2026-02-05T07:29:08Z

    pool: Optional[Tuple[int, ...]] = None,
    retain_graph_in_backward: bool = False,
    _reuse_graph_input_output_buffers: bool = False,
+    pre_warmup_hook: Optional[Callable] = None,


Just to confirm: are the hooks used to disable and re-enable offloading for warmup? Could you point me to the implementation in MCore?

Yes, https://github.com/lhb8125/Megatron-LM/blob/hongbinl/activation_offloading_refactor_cuda_graph/megatron/core/transformer/cuda_graphs.py#L2184

and https://github.com/lhb8125/Megatron-LM/blob/hongbinl/activation_offloading_refactor_cuda_graph/megatron/core/pipeline_parallel/fine_grained_activation_offload.py#L1353

Yes, https://github.com/lhb8125/Megatron-LM/blob/hongbinl/activation_offloading_refactor_cuda_graph/megatron/core/pipeline_parallel/fine_grained_activation_offload.py#L1353

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

greptile-apps

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

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

lhb8125 · 2026-02-25T07:13:34Z

/te-ci pytorch L1

greptile-apps

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

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

lhb8125 · 2026-02-26T02:55:36Z

/te-ci pytorch L1

lhb8125 · 2026-02-27T02:07:19Z

/te-ci pytorch L1

lhb8125 · 2026-03-02T03:49:50Z

/te-ci pytorch L1

lhb8125 · 2026-03-02T06:36:21Z

/te-ci pytorch L1

yaox12

LGTM

* support cuda graph capture offloading module Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove reset_hook and init_chunk_handler_hook Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * remove reset_hook and init_chunk_handler_hook Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: root <root@eos0046.eos.clusters.nvidia.com> * temp fix overlap-grad-reduce Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reuse mark_not_offload() and do not offload scale_inv Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * temp fix for mxfp8 Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * fix bug for record_stream and from_blob Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * disable offloading core_attn_out and refine cpu overhead of at::empty Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * return ptr of whole buffer and offload the whole buffer Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestions from code revie Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * remove code changes of offloading and quantizer Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * minor fix Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> * add docstring Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> --------- Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> Signed-off-by: root <root@eos0046.eos.clusters.nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <root@eos0046.eos.clusters.nvidia.com> Co-authored-by: root <root@eos0022.eos.clusters.nvidia.com>

support cuda graph capture offloading module

5fc9857

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

lhb8125 marked this pull request as draft December 1, 2025 05:34

[pre-commit.ci] auto fixes from pre-commit.com hooks

913fbe8

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Dec 1, 2025

View reviewed changes

Comment thread transformer_engine/pytorch/graph.py Outdated

Comment thread transformer_engine/pytorch/graph.py Outdated

Comment thread transformer_engine/pytorch/graph.py Outdated

lhb8125 and others added 17 commits December 8, 2025 06:35

remove reset_hook and init_chunk_handler_hook

e04bc00

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

remove reset_hook and init_chunk_handler_hook

dda34c2

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

Merge branch 'hongbinl/offload_activation_cuda_graph' of https://gith…

6ed4b91

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

2f61c00

minor fix

88295b4

Signed-off-by: root <root@eos0046.eos.clusters.nvidia.com>

Merge branch 'hongbinl/offload_activation_cuda_graph' of https://gith…

ed2ee6a

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

temp fix overlap-grad-reduce

09d0801

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

Merge branch 'hongbinl/offload_activation_cuda_graph' of https://gith…

8641228

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

[pre-commit.ci] auto fixes from pre-commit.com hooks

c3e341a

for more information, see https://pre-commit.ci

reuse mark_not_offload() and do not offload scale_inv

6cd4af9

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

Merge branch 'hongbinl/offload_activation_cuda_graph' of https://gith…

b54e77c

…ub.com/lhb8125/TransformerEngine into hongbinl/offload_activation_cuda_graph

temp fix for mxfp8

ba065fc

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

fix bug for record_stream and from_blob

e00db5e

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

disable offloading core_attn_out and refine cpu overhead of at::empty

f47b543

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

minor fix

7ca3618

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

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

12cf77b

[pre-commit.ci] auto fixes from pre-commit.com hooks

8c8fe59

for more information, see https://pre-commit.ci

lhb8125 marked this pull request as ready for review February 5, 2026 03:37

greptile-apps Bot reviewed Feb 5, 2026

View reviewed changes

lhb8125 commented Feb 5, 2026

View reviewed changes

buptzyb reviewed Feb 5, 2026

View reviewed changes

return ptr of whole buffer and offload the whole buffer

8421cf9

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

greptile-apps Bot reviewed Feb 25, 2026

View reviewed changes

lhb8125 force-pushed the hongbinl/offload_activation_cuda_graph branch from a181176 to aeb3ecb Compare February 25, 2026 07:08

greptile-apps Bot reviewed Feb 25, 2026

View reviewed changes

lhb8125 added 8 commits February 24, 2026 23:10

Apply suggestions from code revie

75458d9

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

remove code changes of offloading and quantizer

294ca41

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

minor fix

1caf359

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

minor fix

473d4e3

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

minor fix

2133d10

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

minor fix

2cb76a5

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

minor fix

df4a4c1

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

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

60ad9a7

lhb8125 force-pushed the hongbinl/offload_activation_cuda_graph branch from aeb3ecb to 60ad9a7 Compare February 25, 2026 07:12

greptile-apps Bot reviewed Feb 25, 2026

View reviewed changes

lhb8125 requested a review from buptzyb February 25, 2026 07:19

buptzyb reviewed Feb 26, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/graph.py

lhb8125 and others added 2 commits February 25, 2026 18:35

add docstring

46c8dc0

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

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

aa64663

buptzyb approved these changes Feb 26, 2026

View reviewed changes

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

37de9cc

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

945bc86

Merge branch 'main' into hongbinl/offload_activation_cuda_graph

5115a1c

yaox12 approved these changes Mar 2, 2026

View reviewed changes

yaox12 merged commit bba7bf6 into NVIDIA:main Mar 2, 2026
29 of 32 checks passed

nvMelissa added the 2.14.0 label Mar 3, 2026

Conversation

lhb8125 commented Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Uh oh!

greptile-apps Bot left a comment • 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!

Uh oh!

Uh oh!

Uh oh!

lhb8125 commented Feb 5, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

lhb8125 commented Feb 5, 2026

Uh oh!

lhb8125 Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

lhb8125 Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

buptzyb Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

lhb8125 Feb 25, 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!

lhb8125 Feb 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

lhb8125 commented Feb 25, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lhb8125 commented Feb 26, 2026

Uh oh!

lhb8125 commented Feb 27, 2026

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

lhb8125 commented Mar 2, 2026

Uh oh!

yaox12 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

lhb8125 commented Dec 1, 2025 •

edited

Loading

greptile-apps Bot commented Dec 1, 2025 •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

lhb8125 Feb 25, 2026 •

edited

Loading