Add DeepSeek-V4.1-Flash miles recipe with upstream source overlays

[joyliu-q]

Onboards DeepSeek-V4.1-Flash (deepseek-ai/DeepSeek-V4.1-Flash: 40L / 5120h / 384 routed experts top-6, DSA indexer + CSA + Engram) as a miles model, ported from scripts/run_deepseek_v41.py in radixark/miles#3179, and made to train on amd64 H200s from the public fp8/fp4 checkpoint.

Image: amd64 nightly + radixark's V4.1 sglang tree + miles#3179

V4.1 support is unmerged on both sides (radixark/miles#3179, sgl-project/sglang#38798) and the only image built from them, radixark/miles:deepseek-v41, is arm64/GB300-only — Modal's builder rejects it. Neither public sglang ref works alone either: miles drives its engines through the weight-update session API (/begin_weight_update … /end_weight_update) that exists only on sglang-miles, and #38798 branches from upstream without it. The merged tree radixark actually trains on exists nowhere public except inside that arm64 image.

So the recipe composes the same thing for amd64 at image build:

radixark/miles:dev-202609050049          # last amd64 nightly on the same lmsysorg/sglang:v0.5.18 base
  + miles_git_ref = pull/3179/head        # apply_source_overlays (launcher.py), gym patches reapplied after
  + patch_deepseek_v41_sglang_tree        # unpack the sglang python layer of radixark/miles:deepseek-v41
                                          # (pinned by digest) over /sgl-workspace/sglang/python
  + 5 patches for what that tree/miles still get wrong on H200 (below)

MilesRecipe gains two launcher-only fields, miles_git_ref / sglang_git_ref, skipped in cli_args like local_miles. The overlay re-applies the gym's reporting / substep-timing / sglang-abort patches, since the checkout reverts them. All of this collapses to one published tag once the PRs land in an amd64 nightly.

Build-time patches (frameworks/miles/modal_helpers/patches/patch_deepseek_v41_*)

Each is idempotent, asserts on the exact text it expects, and documents the failure it fixes:

patch	fixes
sglang_tree	pulls the V4.1 sglang tree described above
processor_tokenizer	UnboundLocalError: get_tokenizer — the merge left the import local to a branch the V4.1 early-return skips
fp8_hopper_gemm	KeyError: SWAP_AB — the merge dropped the sm90 dispatch for #38798's Hopper-tuned [32,32] block-FP8 kernel
chat_template	registers deepseek_v41 on miles' DeepSeekFamily bridge (the checkpoint ships no jinja template; sglang's encoding_dsv41 renders prompts)
fp4_dequant_block	FP4→FP8 expert dequant hardcoded 128×128 scale blocks (V4's); V4.1 declares [32, 32], so the MoE kernel asserted on the first forward and miles' weight sync would have pushed mismatched scales
vision_topk_capture	sglang's V4.1 router goes through dsv41/vl_routing.vision_topk, which never calls capture_routed_experts_if_allowed, so --use-rollout-routing-replay replayed an all-zero expert buffer and Megatron's dropless dispatcher died with Split sizes doesn't match total dim 0 size

Checkpoint conversion (HF fp8/fp4 → bf16 torch_dist)

The public release is fp8 dense + packed-e2m1 fp4 experts with e8m0 block scales and no bf16 export. hf_block_dequant.py wraps mbridge's safetensor reader (CONVERT_DEQUANT_HF_WEIGHTS=1) so every weight with a sibling .scale comes back dequantized to bf16. Tensors are read with positional file reads rather than safe_open: under the sandbox, mmap'd shard pages are charged to the container and the 485 GB read got ranks SIGKILLed long before host RAM was exhausted (confirmed with the memory logger the wrapper now runs on local rank 0).

Layout is upstream's TP4/PP1 but at EP16 over two nodes instead of EP8 on one — the bf16 model is ~970 GB, ~121 GiB/rank at EP8, which fits a 288 GiB GB300 but not an H200. get_checkpoint_conversion_policy now widens the conversion world to lcm(tp*pp, etp*ep*pp) (data-parallel replicas) instead of rejecting EP > tp*pp; CONVERT_KEEP_PP1=1 stops the converter from inflating PP to the rank count. _is_complete_torch_dist_checkpoint no longer requires common.pt, which newer megatron-core folds into .metadata. dsv4_impl is forwarded to the converter (default megatron rejects TP>1). Output is ~1.1 TB, staged on 2 TB ephemeral disk.

Recipe

8×8 H200 colocated, TP4/PP1/CP1/EP64/ETP1 + sequence parallel, --dsv4-impl miles, --sglang-attention-backend dsv4, bf16 + transformer_engine, full uniform recompute, frozen router gate and e-score correction bias, deterministic mode, R3 routing replay, CPU-offloaded optimizer. Model args come from miles_model_name = "deepseek-v4.1" (upstream's Python model script, which carries the V4.1 --spec). Differences from upstream's GB300 config, each commented in the recipe:

• SGLANG_DSV4_FP4_EXPERTS=1 + SGLANG_DSV4_FP4_DEQUANT=1: load the public packed-fp4 experts and dequantize to fp8 at load (the mxfp4 MoE kernels are Blackwell-only).
• sglang_mem_fraction_static=0.9: dequantized experts + engram tables take ~112 GB/H200; upstream's 0.6 left no KV pool.
• extra_config: sglang_device=cuda (miles renders device and newer sglang no longer defaults it), loader num_threads=2 + drop_cache_after_load (8 TP ranks × 9 mmap'd 10 GB shards in flight exhausted the 2 TB node and got a scheduler SIGKILLed; mmap stays on because the eager path can't decode F8_E8M0).
• use_dynamic_batch_size=False (conflicts with --qkv-format bshd), rollout_health_check_first_wait=3600 (first deepgemm JIT).

Registered run_on_pr=False: 64 H200s plus a ~485 GB download is not a per-PR shape.

Validation

Run 21 (before the vision_topk_capture patch) got through image build, checkpoint conversion, engine load (load_weight=550s), a 128-sample rollout with correct math answers and raw_reward=0.46, and the miles weight sync (update_weights … ok=true 39.1s), then died in compute_log_prob on the MoE all-to-all split mismatch above. A standalone engine probe with --enable-return-routed-experts confirmed the returned buffer was [tokens, 40, 6] but all zeros; with the patch it holds 6 distinct valid ids in [0, 384) for every token/layer.

One-step run with the fix (Run 22) is in progress — this section will be updated with the step result.

Checklist

• Example is documented with comments throughout, in a Literate Programming style.
• Example does not require third-party dependencies to be installed locally
• Example pins its dependencies
  • Example pins container images to a stable tag, not a dynamic tag like latest
  • Example specifies a python_version for the base image, if it is used
  • Example pins all dependencies to at least minor version, ~=x.y.z or ==x.y
  • Example dependencies with version < 1 are pinned to patch version, ==0.y.z

Link to Devin session: https://modal.devinenterprise.com/sessions/f090ce96b33d49d0978b71260859cf55
Open in Devin Desktop: https://modal.devinenterprise.com/desktop/session/f090ce96b33d49d0978b71260859cf55?variant=devin
Requested by: @joyliu-q

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR that start with 'DevinAI' or '@devin'.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[github-actions]

Documentation preview: https://modal-labs-training-gym--training-gym-previews-preview-r-78da35.modal.run/530/docs