Skip to content

[AMD] Drop inert DSv4 rollout knobs and add an MTP recipe - #1733

Merged
guapisolo merged 8 commits into
radixark:mainfrom
XinyuJiangCMU:pr/dockerfile-rm-dead-aiter-args-20260627
Jul 29, 2026
Merged

[AMD] Drop inert DSv4 rollout knobs and add an MTP recipe#1733
guapisolo merged 8 commits into
radixark:mainfrom
XinyuJiangCMU:pr/dockerfile-rm-dead-aiter-args-20260627

Conversation

@XinyuJiangCMU

@XinyuJiangCMU XinyuJiangCMU commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Co-authored-with: @JessicaJiang-123

Most of the rollout knobs the DeepSeek-V4 AMD script sets are inert. For DeepseekV4ForCausalLM on HIP, ServerArgs force-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 sets SGLANG_USE_AITER_AG=false to 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.

XinyuJiangCMU and others added 5 commits July 18, 2026 06:39
…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>
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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

XinyuJiangCMU and others added 3 commits July 28, 2026 09:31
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 XinyuJiangCMU changed the title [AMD] Clean up obsolete dsv4 rollout workarounds [AMD] Drop inert DSv4 rollout knobs and add an MTP recipe Jul 29, 2026

@guapisolo guapisolo left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@guapisolo
guapisolo merged commit 6bbc317 into radixark:main Jul 29, 2026
36 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants