[Xdydy]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Enable local whole-layer CUDA Graph together with fine-grained CPU activation offload. Activations live at a fixed GPU virtual address (VA); physical backing can be released or swapped outside Graph replay, so the captured Graph never sees a new pointer.

Why Graph and CPU offload currently conflict

Upstream already has CUDA Graph and activation offload as separate features. Combining them fails on pointer semantics, not on a missing flag.

Graph replay  requires: captured tensor addresses never change
CPU offload   requires: physical device memory is released after forward

Naive offload/reload:

offload: cpu_backup.copy_(src_tensor, non_blocking)   # D2H to pinned CPU
         release GPU tensor (storage returns to the caching allocator)
reload:  gpu_tensor = torch.empty(...)                 # new allocation, address not guaranteed
         gpu_tensor.copy_(cpu_backup, non_blocking)    # H2D

Capture recorded src_tensor.data_ptr(). After reload, gpu_tensor != src_tensor, so replay reads the old address. That can be silent corruption if the allocator reuses the captured VA, or a miss if the data landed elsewhere.

Related: NVIDIA/Megatron-LM#3697. Community therefore either used partial Graph (e.g. Graph + MoE-only offload) or made Graph and activation offload mutually exclusive.

Core idea: keep the VA, replace the physical pages

VMM (CUDA/MUSA Virtual Memory Management) splits “address the Graph captured” from “pages currently mapped there”:

Naive:  act = torch.empty(...)     # address comes from the caching allocator
VMM:    slot = VmmSlot(...)        # reserve a VA first
        act  = slot.tensor(...)    # act.data_ptr() == VA, forever

1. Before capture, reserve a fixed VA and map physical backing.
2. Forward and backward Graphs both record that VA.
3. After forward replay, copy activations to pinned CPU outside the Graph.
4. After D2H, unmap/release physical backing or keep two resident banks; the VA reservation stays.
5. Before backward replay, map new (or reuse resident) backing at the same VA and H2D.
6. Backward Graph still reads the captured address. No recapture.

The Graph column never changes. Graph does not know about P1/P2; it only reads the VA, and we swap what lives behind that VA outside replay.

---

Group lifecycle

LocalCudaGraphOffloadGroup has five life-cycle phases and three explicit state values.

Phase	Entry	state	Role
1. Discovery	discover_tensor()	discovering	Eager warmup records eligible saved-tensor shape/stride/dtype/device; autograd unchanged
2. Allocation	allocate()	discovering → mapped	Outside capture: fixed-VA VMM slot, initial physical backing, pinned CPU buffer
3. Capture	begin_capture() / save_tensor() / finish_capture()	mapped	D2D into the slot; backward Graph records that address
4. Forward offload	offload_and_release()	mapped → unmapped (or resident pending)	After forward replay: D2H; then unmap/release or keep resident bank
5. Backward restore	remap_and_reload()	unmapped → mapped	Same VA: new backing (or resident bank) + H2D; compute stream waits on the restore event

Discovery and capture run once at graph build. Training replay then cycles mapped → unmapped → mapped (or the resident equivalent) every iteration. Empty groups no-op.

---

Build-time flow: discovery, allocate, capture

Slot reserve/map is not done inside Graph capture.

create_cudagraphs()
└─ _CudagraphGlobalRecord.create_cudagraphs()
   ├─ _link_backward_runners()
   │  ├─ collect bwd runners from cudagraph_record
   │  ├─ wire next_bwd_runner in actual backward consumption order
   │  └─ for local Graph + CPU offload: assign slot bank 0/1
   ├─ runner._local_graph_offload_manager()
   │  └─ PipelineOffloadManager.get_instance()
   └─ runner warmup / capture
      ├─ manager.begin_local_graph_discovery(runner)
      │  └─ saved_tensors_hooks(pack_hook=save_tensor)
      │     └─ discover_tensor(shape, stride, dtype, device)
      ├─ manager.end_local_graph_discovery()
      │  └─ group.allocate(manager, slot_bank, group_index)
      │     ├─ reserve fixed VA
      │     ├─ create/map physical backing
      │     └─ allocate pinned CPU buffer
      ├─ manager.begin_local_graph_capture(runner)
      │  └─ save_tensor(tensor)
      │     └─ group.save_tensor(): D2D copy into fixed-VA slot
      └─ manager.end_local_graph_capture()
         └─ capture snapshot / check descriptor order

Capture checks: group count, descriptor count, shape/stride/dtype/device, and data_ptr() == reserved VA.

Resident ping-pong is not an extra env var. _link_backward_runners() assigns banks when _uses_local_graph_activation_offload() is true. Runners without local Graph offload keep local_graph_slot_bank = None. Eager offload never enters this chain.

---

Forward replay: Graph on compute stream, D2H on a copy stream

Principle: replay stays on the compute/replay stream; D2H is queued on d2h_stream and only waits for the forward-done event. No host sync on the hot path.

sequenceDiagram
    autonumber
    participant S as training schedule
    participant R as _CudagraphReplayNode.forward
    participant G as forward CUDA Graph
    participant C as compute/replay stream
    participant M as PipelineOffloadManager
    participant D as d2h_stream
    participant H as pinned CPU buffer

    S->>R: replay runner.fwd_graph
    R->>M: local_graph_forward_wait_ready(runner, replay_stream)
    M->>C: wait previous D2H event (resident bank reuse)
    R->>G: fwd_graph.replay()
    G->>C: execute forward kernels
    R->>M: local_graph_forward_replay(runner)
    M->>D: wait forward completion event
    M->>D: group.enqueue_d2h()
    D->>H: asynchronous device-to-host copy
    D-->>M: record bank D2H completion event

Loading

Resident path does not MemUnmap/MemRelease on the forward hot path. Banks stay mapped after capture. Forward only saves the current activation into that runner’s pinned host buffer and records a D2H event so the bank can be reused later.

---

Backward: look-ahead so reload overlaps the current Graph

Autograd consumption order can be the reverse of forward save order, so runner links follow the recorded backward order. Current runner waits for its own reload, then immediately prefetches next_bwd_runner. Next-layer H2D can then run on a copy stream while the current backward Graph runs.

The first serial restore waited until the current backward Graph finished before starting the next layer’s H2D. That left a host/driver bubble between consecutive Graphs:

Look-ahead moves next-layer H2D (and, on the fallback path, VMM remap) under the current Graph:

---

Resient Threads offload

This is one method to implement the vmm slot.

VMM driver APIs cannot run inside a CUDA host callback (deadlock with the driver). First look-ahead version sent unmap/release/remap to resident worker threads:

---

Ping-pong slots: two mapped banks, no hot-path remap

While there exists some latency when using MemUnmap / MemSetAccess et all, we provide ping-pong slots, which reduces the calls of these driver apis. In some scene, this will reduce some latency.

Capture-time allocate two banks and keep them mapped for the whole run. Replay no longer calls MemUnmap / MemRelease / MemCreate / MemMap / MemSetAccess.

bank 0: fixed VA + resident physical backing
bank 1: fixed VA + resident physical backing

backward runner 0 → bank 0
backward runner 1 → bank 1
backward runner 2 → bank 0
...

Two banks are required because the Graph captured a VA, not an abstract buffer. One shared slot:

current backward graph → reads the shared bank
next runner H2D         → writes the same bank

Next-layer H2D would overwrite the activation still being read (silent numerical error), or a page fault if unmapped at the wrong time.

Host buffers stay per-runner. Ping-pong only solves GPU concurrent read/write:

current backward graph → reads bank 0
next runner H2D         → writes bank 1

before reusing bank 0:
  wait current backward graph consumed event

Resident state:

mapped → resident_d2h_pending → resident_reload_pending → mapped

---

Here is an example to show the steps how ping-pong slots work.

Forward ping-pong (three layers)

Forward order F0 → F1 → F2. H0/H1/H2 are three independent pinned host buffers. Slot A and slot B are two GPU banks (fixed VA + resident backing). Graph writes the layer into its bank, then D2H saves it to that layer’s host buffer. Slot A is reused by F2 only after F0’s D2H has finished, so H0 is not overwritten.

Backward ping-pong (three layers)

Backward order B2 → B1 → B0. Restore G2 from H2 into slot A; B2 reads slot A. While B2 runs, H2D restores G1 from H1 into slot B. B2 reads A, reload writes B — different physical backing, safe overlap. After B2, B1 reads slot B; slot A can then restore G0 from H0. B0 reads slot A.

Two safety conditions:

1. Forward reuse: F2 must not write slot A until F0’s activation has fully D2H’d into H0.
2. Backward reuse: H0 must not write back into slot A until B2 has finished consuming G2.

Ping-pong does not give each layer a permanent GPU slot. It only guarantees that “the bank the current backward Graph is reading” and “the bank the next layer is reloading” are different. Three host buffers stay independent; two GPU banks rotate under D2H-complete and consumed events.

---

Results

Memory (MB, after init / capture / warmup)

Config	allocated	max allocated	reserved	max reserved
No Graph + no offload	24129.93	32944.00	35748	35748
No Graph + CPU offload	24129.03	31663.10	35394	35394
Local Graph + no offload	24720.52	33259.46	37012	37012
Local Graph + CPU offload	24624.48	28366.72	36286	36286

Vs no-Graph/no-offload:

Config	allocated Δ	reserved Δ
No Graph + CPU offload	−0.90 MB	−354 MB
Local Graph + no offload	+590.59 MB	+1264 MB
Local Graph + CPU offload	+494.55 MB	+538 MB

Same Graph mode, offload vs not: no-Graph reserved −354 MB (~1.0%); local Graph reserved −726 MB (~2.0%).

allocated / reserved are PyTorch caching-allocator stats. VMM physical pages from MemCreate are not proven returned to the driver by these two numbers alone. Driver-level MemGetInfo (or vendor tools) is required for that claim.

Performance

Config	n iters	iteration time (ms)	TFLOP/s/GPU
No Graph + no offload	12	1559.1	179.1
No Graph + CPU offload	11	1609.2	173.5
Local Graph + no offload	18	1525.1	183.1
Local Graph + CPU offload	53	1602.9	174.2

Vs no-Graph/no-offload: local Graph alone +~2.2% throughput; CPU offload −~2.7–3.1%. Graph+offload reserved is still +538 MB vs baseline (Graph static memory + offload storage), even though it is −726 MB vs Graph-only.

These are short windows (14 / 13 / 20 / 763 memory-report iterations). Treat as trend, not a ship gate. Formal numbers need matched train_iters, warmup, and an explicit resident-ping-pong vs fallback A/B with driver-level memory.

⚠️ 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 @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: Related to NVIDIA/Megatron-LM#3697 (local Graph vs activation offload pointer conflict). No dedicated public issue yet.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Suggested coverage before Ready for Review:

• TE VMM slot: reserve/map/unmap, address stability across remap, D2H-then-unmap, H2D-then-replay.
• Megatron: discovery order vs capture, resident bank fencing (no cross-runner overwrite), backward look-ahead.
• Negative case: naive empty() reload either silently reads an interloper (allocator reuse of the captured VA) or restores bytes the Graph never reads. VMM unmap turns the same bug into a hard fault.

Code review

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

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

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.