WIP: port MI355X DeepSeek-R1 MXFP4 disagg to srt-slurm - #2633

Every override block in fp4-disagg-fixed-seq.yaml is supposed to be an independently editable unit — that's the whole point of the override_* naming convention, and the file's own base section demonstrates the right pattern: the large cuda-graph-bs list is hoisted once to &cg128 at the top level and then aliased with *cg128 wherever it's needed (prefill and decode within base).

Two override blocks break that pattern by defining a YAML anchor that is then consumed by other, unrelated override blocks rather than just within themselves:

1. override_mtp3_1p2d_tp8_wide (prefill block, line 299) defines &mtp3_tp8. The very next sibling block, override_mtp3_1p2d_tp8_narrow (lines 325-326), doesn't define its own speculative-decoding config — it aliases *mtp3_tp8 for both its prefill and decode keys, reaching into the internals of a block that has nothing to do with it structurally.
2. override_stp_1p1d_dep8's decode block (line 222) defines &cg512, the 512-entry cuda-graph-bs list. Four other override blocks — override_stp_2p1d_dep8 (261), override_mtp3_1p1d_dep8 (393), override_mtp1_1p1d_dep8 (434), and override_mtp1_2p1d_dep8 (480) — all alias *cg512 even though none of them share any other config with override_stp_1p1d_dep8.

Because this is a single YAML document, anchors are document-scoped and resolved in document order — *cg512/*mtp3_tp8 work today only because their defining block happens to appear earlier in the file than every block that aliases it. Nothing in the file signals this dependency locally: reading override_stp_2p1d_dep8 in isolation gives no hint that its cuda-graph-bs list actually lives inside override_stp_1p1d_dep8. This is fragile in a way the file's own convention (base-level &cg128) was clearly designed to avoid.

Concrete failure scenario: a future contributor decides override_mtp3_1p2d_tp8_wide is no longer needed and deletes it (or simply renames its anchor while touching MTP3-wide tuning). YAML anchors are undefined outside their defining document position, so override_mtp3_1p2d_tp8_narrow's *mtp3_tp8 alias on lines 325-326 would immediately fail to parse — with an error pointing at line 325, giving no indication that the real cause is the deletion of an anchor 26 lines earlier in a "sibling" block. The same applies to &cg512: reordering or deleting override_stp_1p1d_dep8 silently breaks parsing for the 4 dependent blocks (up to 5 total counting override_stp_2p1d_dep8), scattered across unrelated topology/tuning families (STP-dep8, MTP3-dep8, MTP1-dep8 ×2).

This is purely a maintainability/DRY concern — the YAML parses correctly today because every defining anchor happens to precede its uses in document order, so nothing is broken at merge time. The fix that avoids the fragility mirrors what the file already does for &cg128: hoist the shared cuda-graph-bs 512-entry list to a base/top-level anchor (e.g. &cg512 defined once in base, or promoted alongside &cg128), and likewise hoist the mtp3_tp8 speculative-decoding config to a shared top-level anchor, then have every dependent override alias the shared, top-level definition instead of a sibling override's internals. That keeps every override block independently addable, removable, or reorderable without risk of silently breaking unrelated blocks.

runners/launch_mi355x-amds-srt.sh (lines 89-108) adds a standalone Python heredoc, PYMODEL, that validates a local checkpoint directory before it is used for MI355X MXFP4 disaggregated serving: it checks that config.json, tokenizer_config.json, and model.safetensors.index.json exist, then loads the index file, computes the sorted set of shard filenames from weight_map, and confirms every named shard is present on disk — raising SystemExit on any gap.

This is functionally the same check already implemented as a bash helper, checkpoint_is_complete(), in benchmarks/single_node/agentic/glm5.2_fp4_b200_sglang_mtp.sh (lines 56-72) and duplicated again in minimaxm3_fp4_b200_mtp.sh / minimaxm3_fp8_h200.sh. That helper checks tokenizer_config.json + (tokenizer.json or tokenizer.model) + model.safetensors.index.json, then runs an inline Python block that loads the index, computes sorted(set(weight_map.values())), and exits nonzero if any shard file is missing — the identical core validation the new PR code performs. The existing helper explicitly exists to catch aborted or metadata-only checkpoint pulls (its usage is tied to CI run 30729467646, where a checkpoint directory existed with valid metadata but incomplete shard data, which caused a serving failure that this check now guards against).

None of these implementations — the two/three existing copies or this PRs new one — are shared through benchmarks/benchmark_lib.sh or any common location. grep -r "checkpoint_is_complete\|weight_map" benchmarks/ confirms they are textually independent copies. Concretely: if a future run reveals that the shard-existence check also needs to handle sharded tokenizer.model files (as flagged by the differences between the two existing copies, one of which checks tokenizer.model/tokenizer.json while the new code omits that entirely and only checks config.json+tokenizer_config.json), that fix has to be manually propagated across three-plus files, and it is easy to update one copy and forget the others, silently leaving some launchers vulnerable to the exact aborted-pull failure mode the check was designed to prevent.

The mitigating factor, which keeps this from being a "must-fix," is that the new copy runs inside a generated sbatch heredoc that is srun-dispatched to a remote compute node and does not source benchmark_lib.sh or any repo-relative bash file, so directly reusing the existing bash function is not a drop-in replacement — some plumbing (e.g., embedding a shared Python script, or writing it out at generation time) would be needed to de-duplicate. That makes this a design/DRY suggestion rather than a functional defect: the new code is correct as written, it just reimplements a well-tested pattern instead of centralizing it.

Proof of duplication: (1) Read benchmarks/single_node/agentic/glm5.2_fp4_b200_sglang_mtp.sh:56-72 — defines checkpoint_is_complete(), which loads model.safetensors.index.json, computes shards = sorted(set(json.load(fh)["weight_map"].values())), and exits 1 if any shard file is missing. (2) Read runners/launch_mi355x-amds-srt.sh:89-108 (the new PYMODEL heredoc) — computes shards = sorted(set(index.get("weight_map", {}).values())) and raises SystemExit if any shard file (root / shard) is missing. (3) These are the same algorithm on the same JSON structure, confirming a duplicate rather than a coincidental similarity. (4) grep across benchmarks/ for weight_map and checkpoint_is_complete shows no shared module defines this logic once for reuse.

Recommended fix: extract the shared JSON-index + shard-existence walk into a small standalone validate_checkpoint.py (or equivalent) that both the bash helpers and this generated heredoc can invoke/embed, so the completeness contract (including the tokenizer file check the new code currently omits) lives in one place.

The bug: runners/launch_mi355x-amds-srt.sh builds model_repos as a bash array and joins it with commas (IFS=,; echo "${model_repos[*]}") to produce model_repo_list, then passes it to srun as:

--export=ALL,HF_HOME=...,MODEL_REPOS="$model_repo_list"

Slurm's srun --export documents its argument as a comma-separated list of NAME or NAME=value tokens, with no escaping mechanism for a comma embedded inside a value. So once model_repos has two or more entries, e.g. ("deepseek-ai/DeepSeek-R1-0528" "SGLang/DeepSeek-R1-NextN"), the resulting string MODEL_REPOS=repoA,repoB is not parsed as one NAME=value pair — srun splits on the comma first, yielding two tokens: MODEL_REPOS=repoA (a valid export) and a bare repoB (interpreted as "propagate an existing env var named repoB"). Since repoB contains a / (HF repo ids always do) it is not a legal environment variable name and does not exist in the parent shell anyway, so srun silently drops it.

Where this fires: Inside the container, os.environ["MODEL_REPOS"] therefore only ever contains the first repo. The Python one-liner [snapshot_download(repo) for repo in os.environ["MODEL_REPOS"].split(",") if repo] then only prefetches repo[0]; every repo after the first is silently skipped from the shared-cache prefetch step.

Why the surrounding code doesn't catch this: The array/loop refactor in this PR was written specifically to support multiple repos — seed_legacy_cache is called in a for repo in "${model_repos[@]}" loop, and the Python download step is a list comprehension precisely to prefetch more than one model. The one link in that chain that can't carry more than one element is the comma-joined --export transport, because Slurm's own comma-delimited parsing collides with the comma-joined value.

Reachability today: I checked every entry in configs/amd-master.yaml that reaches this launcher. Every MTP arm sets SRT_MODEL_LOCAL_PATH alongside SRT_DRAFT_MODEL (see e.g. override_mtp3_1p1d_tp8 selectors), so model_repos becomes [SRT_DRAFT_MODEL] — a single entry, because the script's else branch that appends $MODEL is skipped when SRT_MODEL_LOCAL_PATH is set. Every non-MTP (STP) arm sets only SRT_MODEL_LOCAL_PATH, giving an empty model_repos (early exit at ${#model_repos[@]} == 0). So no config shipped in this PR triggers the multi-entry path — the bug is latent.

Step-by-step proof of the failure mode (hypothetical future recipe): Suppose a future disagg-MTP recipe uses a remote (non-local) main model plus a draft model — i.e. it sets SRT_DRAFT_MODEL=SGLang/DeepSeek-R1-NextN but does not set SRT_MODEL_LOCAL_PATH, and MODEL=deepseek-ai/DeepSeek-R1-0528.

1. SRT_MODEL_LOCAL_PATH is empty -> the else branch runs -> model_repos+=("$MODEL") -> model_repos=("deepseek-ai/DeepSeek-R1-0528").
2. SRT_DRAFT_MODEL is set -> model_repos+=("$SRT_DRAFT_MODEL") -> model_repos=("deepseek-ai/DeepSeek-R1-0528" "SGLang/DeepSeek-R1-NextN").
3. model_repo_list="deepseek-ai/DeepSeek-R1-0528,SGLang/DeepSeek-R1-NextN".
4. srun receives --export=ALL,...,MODEL_REPOS=deepseek-ai/DeepSeek-R1-0528,SGLang/DeepSeek-R1-NextN.
5. srun's own comma-split parses this as two tokens: MODEL_REPOS=deepseek-ai/DeepSeek-R1-0528 (accepted) and SGLang/DeepSeek-R1-NextN (a bare token — not a legal var name because of the /, and not present in the environment, so it's dropped with no error).
6. Inside the container, os.environ["MODEL_REPOS"] == "deepseek-ai/DeepSeek-R1-0528" only.
7. snapshot_download runs once, for the main model only. The NextN draft model is never pre-seeded into the shared cache.

Impact: This degrades gracefully rather than crashing — the skipped draft model just falls back to being downloaded independently by each P/D worker at startup, which is the exact pre-refactor behavior the PR set out to eliminate for shared-cache efficiency. It does not corrupt results or fail the job.

Fix: Use a delimiter that cannot appear inside an HF repo id and that Slurm's --export parser will not itself split on — e.g. a space or newline — when joining model_repo_list, and split on that same delimiter in the Python os.environ["MODEL_REPOS"].split(...) call. Since --export values also generally can't contain characters like spaces safely across all Slurm versions/shells, a newline joined via printf '%s\n' (or passing repos through a small temp file instead of --export) is the more robust choice.