DeepSpeed

f94aa5c5 - Reuse continuous rollout graphs and release captured resources on shutdown (#9009)

Reuse continuous rollout graphs and release captured resources on shutdown (#9009) ## 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 #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 #9008 + this PR. - Builds on #8732 (persistent continuous graph), reusing its existing fixed-capacity setting rather than adding another configuration option. - #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). #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: ```bash 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): ```bash 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: ```bash 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: ```bash 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 #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): ```bash 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`. ```bash 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. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>