Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
guapisolo merged 8 commits into from
Jul 29, 2026
Merged

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

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d34d044
[AMD] run_deepseek_v4: enable cuda-graph on gfx950, drop obsolete wor…
XinyuJiangCMU Jul 18, 2026
5c03ec9
train: gate per-step rollout onload by offload-rollout-level
XinyuJiangCMU Jul 19, 2026
aa01c1c
Merge remote-tracking branch 'origin/main' into pr/dockerfile-rm-dead…
XinyuJiangCMU Jul 20, 2026
8d23189
1
XinyuJiangCMU Jul 20, 2026
cfe9b35
run_deepseek_v4: restore eval + response-len + mem-fraction to defaults
XinyuJiangCMU Jul 20, 2026
a851f1e
Merge branch 'main' into pr/dockerfile-rm-dead-aiter-args-20260627
XinyuJiangCMU Jul 28, 2026
39f3560
Drop the redundant SGLANG_MEMORY_SAVER_CUDA_GRAPH setting
XinyuJiangCMU Jul 28, 2026
1712877
Add an MTP speculative decoding recipe to the DeepSeek-V4 script
XinyuJiangCMU Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 13 additions & 24 deletions scripts/amd/run_deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ class ScriptArgs(U.ExecuteTrainConfig):

task: Literal["dapo_aime", "gsm8k"] = "dapo_aime"
enable_eval: bool = True
enable_mtp: bool = False

hf_checkpoint: str | None = None
data_dir: str = "/root/datasets"
Expand Down Expand Up @@ -406,34 +407,15 @@ def _train(args: ScriptArgs):
"--router-health-success-threshold 1 "
"--router-health-check-interval-secs 15 "
"--router-health-failure-threshold 40 " # TODO improve
# gfx950: DSv4 sgl-kernel topk_v2 is CUDA-only; route DSA top-k through torch.
"--sglang-dsa-topk-backend torch "
# AITER graph registration fails through HIP IPC on gfx950; use RCCL.
"--sglang-disable-custom-all-reduce "
)
extra_env_vars = {
"SGLANG_SKIP_CHECKPOINT_LOAD_CHECK": "1",
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_HEALTH_CHECK_TIMEOUT": "120",
"SGLANG_DG_CACHE_DIR_PER_PROCESS": "1",
"SGLANG_OPT_FP8_WO_A_GEMM": "0",
# ROCm/gfx950 rollout kernel knobs
"SGLANG_HACK_FLASHMLA_BACKEND": "triton",
"SGLANG_FP8_PAGED_MQA_LOGITS_TORCH": "1",
"SGLANG_DSA_TOPK_BROADCAST": "1",
"SGLANG_OPT_USE_TILELANG_INDEXER": "true",
"SGLANG_OPT_USE_AITER_INDEXER": "false",
"SGLANG_OPT_USE_TILELANG_MHC_PRE": "false",
"SGLANG_OPT_USE_TILELANG_MHC_POST": "false",
"SGLANG_OPT_DEEPGEMM_HC_PRENORM": "false",
"SGLANG_OPT_USE_FUSED_COMPRESS": "true",
"SGLANG_OPT_USE_FUSED_COMPRESS_TRITON": "true",
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "false",
"SGLANG_OPT_USE_TOPK_V2": "false",
"SGLANG_OPT_USE_COMPRESSOR_V2": "false",
"SGLANG_OPT_USE_MULTI_STREAM_OVERLAP": "false",
"SGLANG_ROCM_USE_MULTI_STREAM": "false",
"SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2": "0",
"SGLANG_OPT_USE_FUSED_COMPRESS": "true",
"SGLANG_HEALTH_CHECK_TIMEOUT": "120",
"AITER_BF16_FP8_MOE_BOUND": "0",
}

Expand Down Expand Up @@ -487,9 +469,6 @@ def _train(args: ScriptArgs):
misc_args += "--use-rollout-routing-replay "
# Skip indexer-replay for now
# misc_args += "--use-rollout-indexer-replay "
# Route replay through the miles python router: the Rust router drops return_routed_experts
# on /generate passthrough, so routed_experts never reaches the scheduler.
misc_args += "--use-miles-router "

if args.train_deterministic:
misc_args += "--deterministic-mode "
Expand All @@ -506,6 +485,16 @@ def _train(args: ScriptArgs):
# ROCm TE MoE FP8 lacks fused wgrad accumulation; disable the fusion.
misc_args += "--no-gradient-accumulation-fusion "

if args.enable_mtp:
sglang_args += (
"--sglang-speculative-algorithm EAGLE "
"--sglang-speculative-num-steps 3 "
"--sglang-speculative-eagle-topk 1 "
"--sglang-speculative-num-draft-tokens 4 "
)
# gfx950: use RCCL all-gather for speculative decoding; aiter can deadlock.
extra_env_vars |= {"SGLANG_USE_AITER_AG": "false"}

train_args = (
f"{ckpt_args} "
f"{rollout_args} "
Expand Down
Loading