[garrett361]

Adds a nightly suite that runs trainer-only SFT across every registered trainer architecture, two 4-GPU parallel layouts and a few memory features, to catch catastrophic trainer breakage before it reaches a training run. This PR is the infra. It already catches real bugs (below), which get fixed in follow-ups.

• Each case trains a seeded tiny checkpoint with uv run sft for 5 steps on one 4-GPU node and checks for a clean exit and finite loss. qwen3_moe's fsdp4 base case instead trains to step 8 with a checkpoint at step 5, resumes from it in a fresh process, and requires both runs' losses at steps 6-8 to agree within 1e-3. It trains at lr 1e-4, since at the default 1e-6 a few steps barely move the weights. Calibrated on that case, correct resumes differ by at most 2.5e-5, and a resume without the optimizer or dataloader state by 1.5e-2 or more. Save and resume is mostly arch-independent code, so one case covers it.
• Tiny models start from each arch's published HF config, cut to the fewest leading layers that cover every layer kind. MoE archs then get a smaller hidden size, smaller intermediate sizes and fewer experts. Layer kinds, vocab, head counts and dims, and the attention and indexer shapes stay published. At published widths the MoE checkpoints reached 26 GB, took up to 77 GB of host memory to build, and made GLM-5 and minimax_m2 launches about 80 s slower each. qwen4_exp's n-gram table is also shrunk from about 51B params. GLM-5 and DSv4 use index_topk 128 and 64 instead of 2048 and 512: the suite's documents are shorter than the published top-k, so the indexer kept every key. Now 84% and 68% of query tokens have more candidates than picks. 128 is the smallest value FlashMLA's SM90 sparse kernel accepts for GLM-5.
• MoE archs run fsdp4_ep4 and cp4_ep4, and dense archs run the same layouts with ep=1. The full suite is 33 launches.
• Every arch's cp4 base case also launches with TORCH_LOGS=graph_breaks,dynamo, and test_no_graph_breaks fails on any graph break in that launch, reported as its own test next to test_sft. fullgraph=True is not a substitute: it enables capture_dynamic_output_shape_ops and can hide a fullgraph=False break.
• Feature rows on qwen3_moe each change one setting: selective AC, activation offloading off, optimizer offloading off, full offloading, DeepEP dispatch at EP=4 (deepep_full_ac), and a fullgraph=True row on cp4 that fails on any graph break. Every other case keeps the defaults: compiled, full AC and both fusions.
• qwen3.5 and qwen4_exp are covered through their text configs only. llama uses the ungated unsloth/Llama-3.2-1B-Instruct tokenizer with an explicit llama-3 renderer, since SFT's llama key is the gated meta-llama repo. afmoe is left out, as a TODO in archs.py: SFT has no typed renderer for it.
• tiny_models.py holds one TinyModel per arch (tiny config, tokenizer, renderer) and the checkpoint writer, with nothing SFT-specific, so other suites can reuse it. sft/archs.py defines the matrix: the parallel layouts, the features, and one Arch per arch pairing its TinyModel with how the suite runs it. sft/harness.py renders and checks cases, and sft/README.md gives a short overview.
• --shard I/N splits the suite across 4-GPU runners. Each arch's base cases and its feature cases are separate units, packed longest first by estimated launch time, so a shard builds each checkpoint it needs once.
• A failed launch reports crash or timeout, the run dir, and the tails of trainer.log and of the rank that torchrun names as the root cause.
• Runs write under the shared output_dir (PYTEST_OUTPUT_DIR in CI), use HF_HUB_OFFLINE=1 after downloading each tokenizer once per session, and pass --no-monitors.prime, as the integration tests do: the suite only reads the file monitor's metrics.jsonl.
• The suite lives in tests/regression/distributed/, and test_parallel_invariance.py moves there too, since it is also a multi-GPU trainer test. Its path in regression.yml and the docs follows.
• scripts/ci/run_distributed_nightly.sh runs the suite. The GitHub workflow comes in a later PR, once the 4-GPU runner pool exists.

The suite caught the crash fixed in #4037. Case ids are as they were before #4134 removed ring CP, and the fp32-router FSDP unit the crash needed was later removed by #4177:

case	1e5318108 (before #4037)	4f4d7e440 (#4037)
qwen3_moe-cp2_ring_ep4-base	aten.empty_like.default ... not found in storage	pass
qwen3_moe-cp4_ulysses_ep8-base	aten.empty.memory_format ... not found in storage	pass

Known failures (follow-ups)

These are left failing on purpose, so the nightly shows them until they are fixed. A full run of the matrix on one node gave 80 passes and exactly the 5 graph-break failures below. The DeepEP row, added after that run, fails on its own run as listed. Every test_sft case passes: the AC failure that broke qwen3_5_moe and qwen4_exp and qwen3_5_dense's Ulysses CP failure, both seen on the earlier 60-case matrix, no longer reproduce on current main.

case	failure
test_sft[qwen3_moe-fsdp4_ep4-deepep_full_ac]	Autograd not support dtype: Byte in DeepEP combine (#3954)
test_no_graph_breaks[qwen3_5_dense-cp4_ep1-base]	torch.compiler.disabled call (gated_delta_net.py:88)
test_no_graph_breaks[qwen3_5_moe-cp4_ep4-base]	same gated_delta_net.py:88 break
test_no_graph_breaks[qwen4_exp-cp4_ep4-base]	same gated_delta_net.py:88 break
test_no_graph_breaks[gpt_oss-cp4_ep4-base]	dynamo skips FA4's interface (flash_attn/cute/testing.py:492)
test_no_graph_breaks[nemotron_h-cp4_ep4-base]	dynamic-shape op in FLA (fla/ops/utils/index.py:128)

• gpt_oss's first-use FA4 kernel compile once took over 600 s on a cold node, so each run gets 900 s. The slowest run in the full pass took 416 s (deepseek_v4 fsdp4).
• _from_config leaves MoE expert weights uninitialized on qwen3_moe, glm4_moe, glm_moe_dsa, minimax_m2 and laguna, so the suite's tiny-checkpoint writer initializes them itself. A shared fix in model code must not re-initialize loaded weights on the from_pretrained path.

A full pass of the matrix (32 launches before the DeepEP row) takes about 2 runner-hours on our cluster, 140 to 300 s per run with cold compile caches and imports from /home, so 4 shards finish in about 30 min plus runner setup. Per-launch overhead, not training, dominates: a 5-step launch spends about 120 s before step 1 and 35 to 75 s compiling in step 1. Two production changes cut that and come as separate PRs: a light import path for the launcher's pre_download_* (about 44 s per launch) and removing a CPU-side sum in the EP token dispatcher that makes inductor build C++ kernels (11 to 15 s of step 1).

Running on SLURM

From the repo root, on a node with at least 4 GPUs. Extra args go to pytest. -k selects cases by arch, layout or feature name, and --shard I/N runs one shard.

srun -N1 --exclusive -p <partition> --time=02:00:00 scripts/ci/run_distributed_nightly.sh -k qwen3_moe

Python imports from a venv on /home dominate per-launch time on our cluster. A node-local venv (UV_PROJECT_ENVIRONMENT=/tmp/$USER/venv uv sync) cuts a launch from about 3.2 min to about 70 s. None of that is in the code.

🤖 Generated with Claude Code

[wiz-prime-intellect]

Wiz Scan Summary

Scanner	Findings
Vulnerabilities	-
Sensitive Data	-
Secrets	-
IaC Misconfigurations	-
SAST Findings	-
Software Management Findings	-
Total	-

View scan details in Wiz

To detect these findings earlier in the dev lifecycle, try the Wiz Code extension for VS Code, JetBrains, or Visual Studio.