[AMD] Drop inert DSv4 rollout knobs and add an MTP recipe - #1733
Merged
guapisolo merged 8 commits intoJul 29, 2026
Merged
Conversation
…karound flags cuda-graph is graph-safe on rocm720 (standalone GSM8K 0.950, +2.5x rollout), so the gfx950 workarounds are no longer needed: drop --sglang-disable-custom-all-reduce, --sglang-dsa-topk-backend torch, SGLANG_DSA_TOPK_BROADCAST, SGLANG_OPT_USE_TOPK_V2, --use-miles-router. Add SGLANG_MEMORY_SAVER_CUDA_GRAPH=1 (required so cuda-graph capture routes aiter custom-AR through the unreg path under torch_memory_saver pause/resume, avoiding the aiter radixark#2061 wedge). Lower sglang-mem-fraction-static 0.7 -> 0.6 for 4-node colocate. Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
The loop offloads the sglang engine per --offload-rollout-level but the onload after update_weights was still unconditional, so a partial level left offload and onload asymmetric. Gate the loop onload on the same flag. WIP: cuda_graph still needs to be resumed unconditionally (it currently rides inside onload_kv) -- see PR description. Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
…-aiter-args-20260627
Revert the 4-node-colocate tuning back to the original values: re-enable eval, rollout/eval max-response-len 4096, and sglang-mem-fraction-static 0.7. Keep only the cuda-graph enable and the workaround-flag cleanup in this PR.
Contributor
There was a problem hiding this comment.
Code Review
This pull request cleans up the configuration in scripts/amd/run_deepseek_v4.py by removing several unused or obsolete CLI arguments and environment variables, including --sglang-dsa-topk-backend, --sglang-disable-custom-all-reduce, and --use-miles-router. It also introduces the SGLANG_MEMORY_SAVER_CUDA_GRAPH environment variable. There are no review comments, and I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
miles already turns this on for every rollout engine. server_group.py sets "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true" as the default in the env_vars dict, with no condition on model or colocate, and sglang parses "true" and "1" to the same boolean. Setting it again in the DeepSeek-V4 runner changes nothing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015dv3gbGviUMDoXHSqZHJS7
--enable-mtp serves the checkpoint's MTP block as an EAGLE draft. On gfx950 it also switches the all-gather to RCCL, since the aiter one can deadlock the rollout engines. Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
XinyuJiangCMU
added a commit
to XinyuJiangCMU/miles
that referenced
this pull request
Aug 1, 2026
Brings in 107 upstream commits, including the merged versions of the weight sync selector (radixark#1926), the DSv4 rollout knob cleanup and MTP recipe (radixark#1733), and the JIT norm pin (radixark#2040), all of which started here. Conflict resolution: - update_weight/common.py, update_weight_from_tensor.py, sglang_engine.py: take upstream. These conflicts were only placement, import wrapping and a docstring; upstream also carries review changes the local copies predate (the unmarked grouped-expert gather helper, tighter partition_dim checks, the check_equal parameter). All three files now match upstream exactly. - scripts/amd/run_deepseek_v4.py: keep enable_eval=False and the multinode parallel config, take upstream's extra_env_vars verbatim. Upstream's dict is the reviewed end state of radixark#1733 and radixark#2040: SGLANG_OPT_USE_COMPRESSOR_V2 is gone and SGLANG_OPT_USE_JIT_NORM=false is in. Keep the cuda-graph note, which has no upstream counterpart. Neither side sets --disable-cuda-graph, so cuda-graph stays on by default. Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
XinyuJiangCMU
added a commit
to XinyuJiangCMU/miles
that referenced
this pull request
Aug 1, 2026
…removed radixark#1733 audited these on DeepSeek-V4-Flash / gfx950 and dropped them from the launcher as inert -- force-set to the same value by the DeepseekV4 + is_hip block in server_args, equal to the platform default, on a CUDA-only path, or never read at all. That conclusion does not depend on which channel sets them, so they should not survive here either. Five more were duplicates of entries the launcher still sets, which is how SGLANG_OPT_USE_COMPRESSOR_V2 came to outlive its last reader. The launcher's extra_env_vars is the channel that reaches the sglang engine; this file should not be a second copy of it. NVTE_FP8_BLOCK_SCALING_FP32_SCALES stays: it is TransformerEngine on the training side, not a rollout knob. Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Co-authored-with: @JessicaJiang-123
Most of the rollout knobs the DeepSeek-V4 AMD script sets are inert. For
DeepseekV4ForCausalLMon HIP,ServerArgsforce-sets the same variables to the same values right after, so the script is writing environment variables that get overwritten a moment later. This drops those 13, plus three CLI workarounds that are no longer needed:--sglang-dsa-topk-backend torch,--sglang-disable-custom-all-reduce, and--use-miles-router.It also adds
--enable-mtp, which serves the checkpoint's own MTP block as an EAGLE draft with 3 speculative steps, top-k 1, and 4 draft tokens. The draft is not trained. On gfx950 the aiter all-gather deadlocks the rollout engines once speculative decoding is on, so the flag also setsSGLANG_USE_AITER_AG=falseto fall back to RCCL for the all-gather.I ran the MTP recipe on 4 nodes, 32 MI355X, DeepSeek-V4-Flash-FP8 colocate GRPO on dapo_aime. It completed 3 training steps with no engine restarts, at an accepted length of 2.75 to 2.79.