[zhihengy]

ci-sglang-pr: dmd2/h3-dmd2-e2e
ci-sglang-repo: zhihengy/sglang

On top of #343. The H3 DMD2 recipe gets the knobs a CI proxy needs, and a 2-GPU e2e runs it nightly against a recorded standard.

What

• scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py: --num-nodes / --num-gpus-per-node (one node starts its own ray cluster, as the other recipes do), --critic-steps-per-rollout, --eval-interval 0 and --save-interval 0 to run without eval or checkpoints (0 used to reach should_run_periodic_action's modulo), --fsdp-cpu-offload (also offloads the teacher). The 16-GPU defaults are unchanged.
• tests/e2e/short/test_h3_t2va_dmd2_16xGPU_single_node_2xGPU_proxy.py in stage-c-3-gpu-h200, label e2e: 2 rollouts of 4 prompts (4 s, 16:9), 2 critic steps of 2 samples and 1 student step on the last 2, --deterministic-mode. Checks train/critic-loss, train/critic-grad_norm, train/grad_norm and the DMD2 metrics (dmd2_grad_abs_mean, dmd2_score_gap_abs_mean, dmd2_teacher_dist_abs_mean, dmd2_x0_mean, dmd2_x0_std, x0_mean_abs_diff) bit for bit.
• docs/models/h3/h3.md: the single-node form and why it needs CPU offload.

Stack fixes this run surfaced (on the branches below, not in this PR's diff)

• dmd2/02-external-data ([DMD2 3/10] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330): train returns Box(ray.put(output)), but the driver discarded its ray.get and passed the task refs on, so the consumer's box.inner hit an ObjectRef. The driver now keeps the boxes, as RolloutManager.generate keeps the shard boxes.
• dmd2/05-dmd2-loss ([DMD2 7/10] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334): the DMD2 defaults still set custom_train_input_function_path after [DMD2 4/10] feat(fsdp): add a train-input hook that preprocesses train pairs before micro-batching #338 renamed the hook, so --loss-type dmd2 failed in parse_args.

Why this shape

• H3 is ~31B parameters and --loss-type dmd2 trains the student and the critic in full, so student + critic + teacher only fit 2 GPUs with --fsdp-cpu-offload: parameters, gradients and optimizer states live in host memory (about 1.1 TB) and the optimizer steps on the CPU.
• The master stays fp32. H3 keeps proj_in / time_embedder / proj_out in fp32 (_keep_in_fp32_modules), and FSDP2 wants one original dtype across a group's trainable parameters; a bf16 master trips FSDP expects uniform original parameter dtype for full-parameter training (the H3 SFT e2e gets away with bf16 because LoRA freezes those modules).
• 2 GPUs rather than 4: with one sample per rank per step the per-rank work is the same, the host memory is the same, and the 3gpu suite has the shorter queue.
• The engine needs [Feat] diffusion-rl: stop a rollout at a requested step sgl-project/sglang#45230 (stop step) and #46545 (skip decode), which are not in sglang-miles yet, so the PR body points CI at zhihengy/sglang:dmd2/h3-dmd2-e2e (#45230 + #46544 + #46545 on sglang-miles), the branch [DMD2 9/10] feat(dmd2): add the MiniMax H3 4-step DMD2 recipe with eval media and DMD2 metrics #343 names.

Local run (2 x H200, this node)

The CI arg set, in the e2e tree (torch 2.13 with the ad-hoc param-dtype patch port, sglang dmd2/h3-dmd2-e2e), on an otherwise idle pair of H200s. It passes end to end in 50 min; every registered metric is emitted at the expected steps (critic steps 1-4, student steps 1-2).

Phase	Time
Engine up (tp2)	1.5 min
Trainer init: 3 x 62 GB checkpoints into pinned host shards (critic, its teacher, student)	21.5 min
Rollout 0: generate 4 samples at 4 steps / critic 2 steps + scoring / student 1 step	0.5 / 8.3 / 5.0 min
Rollout 1: same	0.7 / 7.3 / 3.6 min
GPU, each rank	33 GB allocated peak, 48 GB used
Host memory for the job	about 0.9 TB

The time is the two offloaded 31B optimizers, so est_time=3300 lifts run_suite's per-file cap to 4125 s (max(timeout_per_file, 1.25 x est_time)); the nightly 3gpu job grows by about 50 min. Smaller is not available: sglang's H3 pins 1344x768 and >= 4 s, and the checkpoint materialization is 20 of the 50 minutes regardless of the rollout count.

Standard

The standard is recorded on the CI runner with record-e2e-standards.yml (runs_on: ["h200","3gpu"], the same sglang ref), not locally: the local tree runs torch 2.13 with an ad-hoc port of the FSDP param-dtype patch and would not match bit for bit.

Record run: https://github.com/radixark/miles_diffusion/actions/runs/42809959116 (50 min on the 3gpu runner; the artifact is committed as tests/ci/fixtures/e2e_standards/test_h3_t2va_dmd2_16xGPU_single_node_2xGPU_proxy.json). The test file is named ..._16xGPU_single_node_2xGPU_proxy.py like the wan22 and krea2 proxies.

Reproducibility check: a second record-e2e-standards.yml run on the same runner and engine, https://github.com/radixark/miles_diffusion/actions/runs/42875943793 (53 min), produced all 22 metric points bit for bit identical to the committed standard.

Note on run-ci-e2e: with the ci-sglang-pr override the whole suite runs on dmd2/h3-dmd2-e2e, whose base is the 2026-10-13 sglang-miles; the sd3 NFT and wan22 standards were recorded on the current sglang-miles and drift at step 2+ (run 42819564862), and the suite stops at the first failure. Like #343, this test only goes green on nightly once sgl-project/sglang#45230, #45523, #46545 and #46544 land in sglang-miles; the standard is then re-recorded there.

Checklist

• pre-commit hooks pass on the changed files
• Added/updated tests for new behaviour (the e2e)
• python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --help still parses; the CI arg set and the 16-GPU defaults both go through parse_args
• No public train flag added; the recipe knobs are documented in the H3 model page