[jinyouzhi]

Summary

Keep the continuous-decode graph alive across RL rollout rounds, and add an explicit way to release it.

Change	Before	After
Graph signature	Includes prompt_len and every parameter's (data_ptr, _version).	Keeps the decode batch size, cache length, device and dtype, plus each parameter's and buffer's address, device, dtype, shape and strides. Captured kernels read these tensors through the address and layout fixed at capture, so in-place value updates are replayed correctly. Replacing a parameter or buffer (for example a re-registered rotary inv_freq or quant scale), or changing its layout at the same address (for example transpose_()), triggers a recapture. An unindexed device is normalized to the current device index.
Width changes	Every new prompt width recaptures, even when continuous_cache_capacity pins the cache length.	With a pinned capacity, one graph serves all widths that fit. The capacity check still runs on every call.
Rebuild order	The new static cache and graph pool are allocated while the old state is still alive.	Drop the old state before allocating its replacement, avoiding overlapping ownership of both states.
HybridEngineRollout.shutdown()	No API; the state lives until the rollout object is collected.	Synchronizes and releases the graph and static cache. The rollout stays usable and recaptures on the next graph call.

shutdown() implements the existing RolloutEngine.shutdown() lifecycle hook. It matters for TP>1: in the NCCL reproducer below, a persistent graph holding captured collectives makes destroy_process_group() hang in ProcessGroupNCCL.shutdown(). Call shutdown() on each rank before destroying the process group. The docs describe reuse and release.

Why dropping _version is safe

• With bf16, DeepSpeed's optimizer writes updated weights through flat buffers that the parameters view via .data, which does not bump param._version. On master, single-width bf16 rollouts already reused the graph after engine.step(), and the replayed outputs changed with the weights. The version check gave no protection there.
• On paths where a plain torch optimizer updates parameters in place (for example fp32 ZeRO-0), _version changes on every step, so master recaptured on every rollout round.
• The cases where storage really moves (AutoTP gather/partition or a .data = swap) change data_ptr, which stays in the signature.
• A same-address layout change also invalidates the captured kernels. Shape and strides are in the signature, so this recaptures too. This also covers param.data = param.data.t(), which keeps both data_ptr and _version and was therefore missed by the old signature.

Related work / scope

• Depends on Size HybridEngine rollout KV caches from local prefill shapes under AutoTP #9008 for TP>1 continuous graph. Without it, TP>1 continuous graph fails at the KV copy before this code matters. The two PRs touch neighbouring lines in _get_continuous_graph_state; whichever merges second needs a trivial rebase. TP2 results below use Size HybridEngine rollout KV caches from local prefill shapes under AutoTP #9008 + this PR.
• Builds on Enable CUDA Graph for continuous rollouts #8732 (persistent continuous graph), reusing its existing fixed-capacity setting rather than adding another configuration option.
• Measure first-call and post-update costs of continuous rollout CUDA Graphs #8853 only measures graph capture and keeps the version check. Its test_continuous_graph_first_call_benchmark_on_cuda asserts capture.call_count == 2 because an in-place SGD step currently forces a recapture. With this PR that step reuses the graph, so whichever of the two merges second needs that expectation changed to 1 (and the benchmark's post-update number no longer includes a recapture). perf(rollout): select generation and prefill buckets with a tunable cost model #8786 buckets prefill and does not change graph reuse or release.
• Not covered: ZeRO-3, offload, HybridEngine kernel injection, LoRA fuse/unfuse inside rollout.

Validation

After merging master

Merge dedd6c840 includes master a4c989b76. On 1× NVIDIA L20 (driver 580.173.02, CUDA 13.0, torch 2.14.1+cu130, Transformers 5.17.0), the tested files matched the merged worktree by SHA-256:

unset DEEPSPEED_PRETRAINED_ROLLOUT_MODEL
CUDA_VISIBLE_DEVICES=0 DS_ACCELERATOR=cuda OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
  HF_HUB_OFFLINE=1 timeout 600 python -m pytest -q -p no:cacheprovider -o addopts= \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py tests/unit/utils/test_static_cache.py
# 104 passed, 2 skipped: optional pretrained model and native inference extension

CUDA_VISIBLE_DEVICES=0 DS_ACCELERATOR=cuda OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
  HF_HUB_OFFLINE=1 DEEPSPEED_PRETRAINED_ROLLOUT_MODEL=Qwen/Qwen3-0.6B \
  timeout 300 python -m pytest -q -p no:cacheprovider -o addopts= \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py::test_continuous_graph_matches_pretrained_model_when_enabled
# 1 passed, using cached Qwen3-0.6B weights

On CPU (AMD EPYC 9575F, 256 logical CPUs visible; Python 3.12.3, torch 2.14.0.dev20260804+xpu):

DS_ACCELERATOR=cpu OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python -m pytest -q -p no:cacheprovider -o addopts= \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py tests/unit/utils/test_static_cache.py \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py
# 78 passed, 28 skipped (accelerator-only cases and the optional pretrained test)

Pre-commit passed on the four PR files after resolving the conflict.

Earlier validation

CUDA: 8× NVIDIA L20 available (1/2/4 devices used below), driver 580.173.02, CUDA 13.0, torch 2.14.1+cu130, Transformers 5.17.0. Base: d8241a9ba. The three tested Python files on L20 have identical SHA-256 hashes to the then-current PR head 55384d86d.

With DEEPSPEED_PRETRAINED_ROLLOUT_MODEL unset:

CUDA_VISIBLE_DEVICES=0 python -m pytest -q -p no:cacheprovider \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py tests/unit/utils/test_static_cache.py
# 61 passed, 2 skipped; exit 0
# Skips: optional pretrained model and native inference kernel unavailable in this invocation.
# Without the layout fix in 55384d86d: the two lifecycle tests fail (unit: state reused after
# transpose_(); v1: graph tokens differ from eager after transpose_()).

The earlier run (before the layout fix) with the optional pretrained model enabled and the native inference kernel available:

CUDA_VISIBLE_DEVICES=0 DEEPSPEED_PRETRAINED_ROLLOUT_MODEL=<Qwen3-0.6B> python -m pytest --forked \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/utils/test_static_cache.py
# master + this PR's test files: 4 failed, 59 passed (3 incompatible/new regression tests + 1 pre-existing)
# this PR:            1 failed, 62 passed  (pre-existing only)
# this PR + #9008:    1 failed, 71 passed  (pre-existing only)

The pre-existing failure on baseline d8241a9ba was the opt-in test_continuous_graph_matches_pretrained_model_when_enabled with Transformers 5.17 ('DeepSpeedStaticCache' object has no attribute 'get_query_offset', see #8852). The fix is now in master and the opt-in test passes after the merge, as shown above; the test remains skipped in CI without the model.

The unit test covers in-place updates, a same-address transpose_(), parameter replacement, capacity changes and shutdown() with graph capture mocked. On a real device the new test runs a tiny Llama with a fixed capacity and widths 8 → 8 → 12 → 8. It applies an in-place weight change before round 2 and a same-address lm_head.weight.transpose_() before round 4. Each round matches plain eager token for token, the output changes after the weight update, and there are exactly 2 captures (initial + after the layout change). shutdown() frees at least the cache size.

Review follow-up (2313079fa, rebased on master 84481c333). Buffers are now in the signature. The unit test also re-registers a buffer, and the device test adds a fifth round that re-registers model.model.rotary_emb.inv_freq at half its values: the output still matches eager, the tokens change, and there are 3 captures. On 1× NVIDIA L20 (torch 2.14.1+cu130): tests/unit/v1/test_hybrid_engine_rollout_v1.py + tests/unit/runtime/rollout/test_hybrid_engine_rollout.py give 102 passed, 6 skipped. With buffers removed from the signature, the device test fails on the token comparison after the buffer swap. On 1× Arc Pro B70 the same files give 71 passed, 37 skipped (graph-capture cases skip on this XPU stack). Multi-GPU TP and the RL loop below were not rerun for this change, which only adds buffers to the signature.

RL loop on L20. Qwen3-0.6B, TP1, 8 prompts, 4 slots, 32 new tokens, greedy, 6 rounds. Each round does a graph rollout, then a real engine.backward() + engine.step() with AdamW (lr 1e-4). Times are wall-clock per rollout, including prefill.

Setup	Captures (6 rounds)	Rollout ms (rounds 2–6)	Peak MiB (rounds 2–6)
fp32 ZeRO-0, width 64 — master	6	677–703	7069 → 7107
fp32 ZeRO-0, widths 64/96, capacity 160 — this PR	1	546–553	7107 / 7070 alternating
bf16, widths 64/96, capacity 160 — master	6	625–1331	8188 → 8220
bf16, widths 64/96, capacity 160 — this PR	1	401–542	8155 / 8137 alternating
bf16, width 64, no capacity — master / this PR	1 / 1	386–388 / 390–396	unchanged

After the layout fix, the patched fp32 TP1, bf16 TP1 and fp32 TP2×DP2 runs were repeated: still 1 capture, the same per-round token-match rates and the same shutdown() release. Optimizer steps do not change parameter shape or strides.

These are exploratory short runs, not a controlled end-to-end speedup claim. In particular the two fp32 rows use different prompt schedules and cache capacities. No long-run convergence or training-throughput claim is made.

Correctness against plain eager generation.

• fp32: every round is token-identical in the TP1 master and patched runs, and in the patched TP2×DP2 run (with Size HybridEngine rollout KV caches from local prefill shapes under AutoTP #9008).
• bf16, TP1: rounds 4–6 differ from plain eager after training steps. The per-round token-match rates are the same on master and this PR (0.9961 / 0.8984 / 0.9336). This is consistent with precision-sensitive generation, but does not prove the root cause or bitwise equivalence of the two graph runs.
• bf16, TP2×DP2: the patched run also differs from eager (minimum match 0.5742 / 0.9805 / 0.9883 / 0.8047). There is no matched unpatched TP2×DP2 run in this experiment, so this cell is not validated for token-exact correctness.

Tiny Llama fp32 runs:

• ZeRO-2, widths 16/24: 1 capture, token-identical on every round.

Release with TP2 (tiny Llama, #9008 + this PR):

Exit path	Ranks that returned from destroy_process_group()
No shutdown()	0/2 (killed by the 600 s timeout)
shutdown() first	2/2

shutdown() released 116.5 MiB (Qwen3-0.6B bf16, TP1, capacity 160) and 43.7 MiB on rank 0 (TP2×DP2).

Intel XPU (1× Arc Pro B70, Intel compute runtime package 26.27.39122.11, Level Zero loader 1.28.2, torch 2.14.1+xpu, Transformers 5.22.1, on top of #9007, which adds XPU graph capture):

DS_ACCELERATOR=xpu python -m pytest tests/unit/v1/test_hybrid_engine_rollout_v1.py \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
# #9007 + this PR: 52 passed, 3 skipped; exit 0
# Without the layout fix, the v1 lifecycle test fails: graph tokens differ from eager after transpose_().

The earlier torch 2.15.0+xpu run also reported 52 passed / 3 skipped, but aborted during interpreter teardown (exit 134). A minimal torch.ones(16, device="xpu"); torch.xpu.synchronize() program reproduces that abort without importing DeepSpeed. It exits normally on torch 2.14.1+xpu; the clean-exit test run above uses that version. The earlier #9007-only comparison reached the expected assertion failure (3 captures rather than 1), but its process also used the affected 2.15 runtime.

CPU: Python 3.12.3, torch 2.14.0.dev20260804+xpu selected with DS_ACCELERATOR=cpu.

DS_ACCELERATOR=cpu python -m pytest tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/utils/test_static_cache.py tests/unit/v1/test_hybrid_engine_rollout_v1.py
# 57 passed, 6 skipped
pre-commit run --files <changed files>   # passed

Compatibility: fp32 ZeRO-0 at TP1/TP2/TP2×DP2 and fp32 ZeRO-2 at TP1 are token-exact on CUDA. Single-device fp16 graph inference matches eager for the tested Qwen3 and tiny Llama cases after the merge. Single-device fp32 graph tests pass on XPU with #9007. bf16 runs establish capture/release behavior only; their eager-output discrepancies are reported above. Not validated: ZeRO-1, ZeRO-3, offload, fp16 training, autocast, quantization, HybridEngine kernel injection, multi-device XPU graph, pipeline/expert/sequence parallelism, or training convergence.

AI assistance: GitHub Copilot helped with implementation, documentation, tests, review and validation. The original signed code commit is preserved; this PR also includes signed documentation and layout-signature follow-ups.

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-09T04:43:27.382515Z	dedd6c8	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[delock]

Hi @jinyouzhi , my agent found these issues, can you take a look? Thanks!

1. [blocking] The graph-reuse signature only walks module.parameters() — a replaced buffer (e.g. re-registered inv_freq, quant scale, LoRA fuse/unfuse) leaves the captured kernels reading the old address silently; please include named_buffers() (data_ptr, dtype, shape, stride) in the signature plus a unit test.
2. Pre-existing nits, feel free to ignore: the capacity check is duplicated (_validate_continuous_graph_options vs max_cache_len < minimum_cache_len in _generate_continuous_graph), and the signature's device comparison would be more robust normalized (torch.device('cuda') != torch.device('cuda', 0)).
3. Nice work overall — the _version removal argument and the eager-comparison test (with the outputs[0] != outputs[1] stale-weight guard) are solid; just needs the buffer gap closed before merge.

[jinyouzhi]

@delock Thanks for the review. Addressed in 2313079 (branch rebased onto current master):

• Buffers in the signature (blocking): the signature now covers module.buffers() with the same (data_ptr, device, dtype, shape, stride) as parameters. The CPU unit test re-registers a buffer and expects a recapture. The device test adds a round that re-registers rotary_emb.inv_freq with different values: the output matches eager, the tokens change, and the capture count goes from 2 to 3. If buffers are removed from the signature, that test fails on the eager comparison.
• Device normalization: an unindexed device of the accelerator type is mapped to current_device(). In this path the device comes from tensor.device, which already carries an index, so this is defensive.
• Duplicated capacity check: removed the copy in _generate_continuous_graph; _validate_continuous_graph_options still raises the same error before any work is done.

Tested on 1× L20: rollout unit + v1 tests give 102 passed, 6 skipped. B70 runs the same files, but the graph-capture cases skip there. Details are in the updated description. #9077 builds on this and will be rebased.