Repository navigation
Conversation
tohtana
added
the
needs-gpu-test
tohtana
removed
the
needs-gpu-test
29e7658 to
1ffa78a
Compare
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 968022173a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| signature = self._graph_signature(module, batch_size, max_len, device, model_dtype) | ||
| state = self._static_graph_state | ||
| if state is None or state.signature != signature: |
There was a problem hiding this comment.
When inference tensor parallelism is active and only one rank's tensor address or layout changes (for example, for a rank-local shard), this rank-local condition makes that rank execute the initialization, warmup, and capture forwards while its peers immediately replay their retained graph. Those forwards contain TP collectives, so the ranks execute different collective counts and generation can hang; all-reduce a recapture flag over the inference TP group before entering this branch, and apply the same fix to the continuous sibling.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d5c3d71. Both graph paths now all-reduce (MAX) the recapture decision over the AutoTP group, so if any rank needs to capture, the whole group captures together. The new TestAutoTPGraphRecapture moves one parameter on rank 1 only, on the static and continuous paths. With the fix, both ranks capture twice and match generate(). Without the all-reduce, the test hangs.
| state = self._static_graph_state | ||
| if state is None or state.signature != signature: | ||
| # Free the stale cache and graph pool before allocating replacements. | ||
| self._static_graph_state = None |
There was a problem hiding this comment.
When a same-device call changes the batch size or total decode length, the local state variable still references the old _StaticGraphState, so clearing only self._static_graph_state does not release its cache or graph. The replacement cache, static buffers, and graph are therefore captured while both the old state and the new prefill cache remain live, which can transiently double this large allocation and OOM even when either state fits alone; set the local to None before allocating, as _get_continuous_graph_state already does.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d5c3d71: the local state is cleared before the replacement is allocated, as in _get_continuous_graph_state.
| """ | ||
| from deepspeed.utils.static_cache import DeepSpeedStaticCache | ||
|
|
||
| signature = self._graph_signature(module, batch_size, max_len, device, model_dtype) |
There was a problem hiding this comment.
When the same rollout is invoked once in training mode and later in eval mode (or vice versa) with a model whose forward branches on module.training, such as one with nonzero dropout, this reuse call sees the same tensor signature and replays kernels captured in the previous mode while only the eager prefill reflects the current mode. Before this change every static call recaptured, so decode can now silently retain or omit training-only operations; include the execution mode in the shared graph signature so the static and continuous siblings both recapture.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d5c3d71: module.training is part of the shared graph signature, so both paths capture again when the mode changes. The reuse test now switches to training mode and checks for a third capture.
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Capture runs forwards that hold TP collectives, so a rank whose own graph signature changed must not capture while its peers replay. Both graph paths now all-reduce the recapture decision over the AutoTP group. The graph signature also records module.training, so a graph captured in one mode is not replayed in the other, and the static path drops its stale state before allocating the replacement. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Summary
The default graph path (
use_graph_capture=Truewithoutcontinuous_batch_size) currently rebuilds itsDeepSpeedStaticCache, runs three warmup forwards and captures a new graph on everygenerate()call. In an RL loop that is a fixed setup cost per rollout round, paid even though the decode graph depends only on the decode shape, not on the prompt contents or the weight values.This PR keeps the static path's graph and cache across calls, the same way #8732 / #9009 already do for the continuous path.
generate()callDeepSpeedStaticCache, 3 warmups, capturemax_new_tokens, device, dtype, or a parameter's address / shape / strides / dtype / deviceshutdown()Implementation
_continuous_graph_signature→_graph_signature, shared by both paths (no behaviour change).DeepSpeedStaticCache(module.config, max_cache_len=...)and, as on the continuous path, capture first runs one forward on the current stream so each layer is sized from the model's local (AutoTP-sharded) KV tensors before the side-stream warmup._graph_signaturekeeps Reuse continuous rollout graphs and release captured resources on shutdown #9009's signature (parameters and buffers, normalized device index)._StaticGraphState/_get_static_graph_state/_capture_static_graph. The warmup + capture code is moved unchanged out of_generate_graph. Capture now runs on a zeroed cache before the prefill KV is loaded; the replayed values never depend on capture-time contents.shutdown()is called. Documented ininference-engine.rst.Tests
New accelerator test
test_static_graph_is_reused_across_rollouts_with_the_same_decode_shape(Llama, sdpa). Every call is compared againstmodel.generate():Then it asserts exactly 2 captures, and that
shutdown()frees at least the cache bytes.Mutation check (each must fail the new test): not zeroing the mask on reuse → fails; recapturing every call → fails (capture count); skipping the prefill KV copy → fails (all three static tests).
Validation
1× NVIDIA H20 (driver 580.178.04), torch 2.13.0+cu130, Transformers 5.15.1:
CUDA_VISIBLE_DEVICES=7 DS_ACCELERATOR=cuda OMP_NUM_THREADS=1 HF_HUB_OFFLINE=1 \ 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 # 105 passed, 2 skipped (optional pretrained model, native inference extension)After the rebase, 2× NVIDIA L20 (driver 580.173.02), torch 2.14.1+cu130, Transformers 5.17.0, same files: 147 passed, 4 skipped, 1 failed. The failure is
TestZero3SynchronizedEarlyStop, which fails the same way on master on this machine (JIT build oftransformer_inferencefails).CPU (
DS_ACCELERATOR=cpu, same files): 78 passed, 29 skipped (accelerator-only cases).Pre-commit passes on the three changed files.
Not covered
continuous_cache_capacity); can follow if wanted.