torchtitan: miles-side model packages, with GLM-4.7-Flash as the first

[Zhichenzzz]

ci-image-tag: pr-3171

Stacked on #3248 (which is stacked on #3171). Review only the top four commits.

What

A place for models torchtitan does not ship, and the first one: GLM-4.7-Flash.

--titan-model-name <name> now resolves miles.backends.torchtitan_utils.models.<name> before torchtitan.models.<name>. Both roots expose the same model_registry(flavor, attn_backend) -> ModelSpec, so nothing downstream of resolve_model_spec distinguishes a miles package from a torchtitan one.

Why this shape

torchtitan's own extension point is the ModelSpec: model config, parallelize_fn, pipelining_fn, state_dict_adapter. Most new architectures are a recombination of blocks torchtitan already has (MLA/GQA attention, sigmoid or softmax token-choice router, grouped experts, shared experts), so a miles-side package is usually a flavor table plus a model_registry pointing at an existing parallelize_fn and adapter. When the architecture genuinely differs, the package grows a model.py, then a state_dict_adapter.py, then a parallelize.py, in that order.

GLM-4.7-Flash

HF glm4_moe_lite is DeepSeek-V3's layer with different sizes: MLA (q_lora 768 / kv_lora 512 / nope 192 / rope 64 / v 256), sigmoid top-4 router with the correction bias, route scale 1.8 with normalisation, one shared expert, a dense first layer, RoPE theta 1e6 without YaRN, 47 layers, vocab 154880. Its checkpoint uses DeepSeek-V3's parameter names verbatim, so models/glm4_moe_lite/ builds the layers with torchtitan's DeepSeek-V3 helpers and reuses parallelize_deepseekv3 and DeepSeekV3StateDictAdapter unchanged.

The checkpoint's MTP layer (layers.47) is not built. The HF load requests only the keys the model's to_hf() produces, so those tensors are never read, and the weight stream to SGLang does not carry them.

A bug in #3171 this surfaced

SGLang's DeepSeek family loads q_a_proj and kv_a_proj_with_mqa into one fused parameter and writes it only when both halves arrive in the same load call; miles keeps the pair in one bucket through get_hf_atomic_update_groups(). The torchtitan bridge asked the torchtitan model config for q_lora_rank, which DeepSeek-V3-shaped configs keep on layers[0].attention, so it read None, registered no group, and whichever layer's pair straddled a bucket boundary was silently never updated in the engine. The first GLM run failed the post-update weight check on exactly one tensor (layers.10...fused_qkv_a_proj_with_mqa). The bridge now reads q_lora_rank from the HF config. Separate commit.

Tests

• fast: the flavor's dimensions are pinned to GLM-4.7-Flash's config.json, and the adapter's map must cover every key pattern in the checkpoint's model.safetensors.index.json; miles-first resolution and the two-root error message.
• e2e: tests/e2e/torchtitan/test_glm47_flash_tp_ep_r3.py, colocated on 8 GPUs, TP2 EP4 with 4-GPU engines, routing replay, SGLang's own glm4_moe_lite. SGLang attention is pinned to triton: flashinfer's SM100 prefill kernel has no head_dim 256 instantiation (MLA's 192 + 64 query head pads to it).

Validated on 4x GB300 (run: https://wandb.ai/miles_training/backend_unify/runs/1m82oxpk): weight check clean on both syncs, 2 rollouts, 2 training steps. That pod's 626 GiB cgroup cannot hold the 4-rank CPU offload of a 30B-A3B trainer, so the run there used --no-offload-train with --sglang-mem-fraction-static 0.35; the committed case keeps the default colocate offload, which fits the 2 TB H200 CI nodes.

Docs

docs/user-guide/training-backend.md: the model lookup order, the Limits warning, and a "Going deeper: a model torchtitan does not ship" section.

Megatron vs torchtitan on GLM-4.7-Flash, same setting

4x GB300, colocated, TP2 EP2, routing replay, 8k context (4096 prompt + 4096 response), 32 samples per step, --use-dynamic-batch-size --max-tokens-per-gpu 8192 on both, trainer kept on GPU, the same SGLang configuration (triton attention and MoE runner, mem_fraction_static 0.35). Steady-state steps 1-5, 6 rollouts each. Rollout time is the same on both (42-48 s), so the table is the training side.

per step	torchtitan	Megatron
actor_train_time (s)	21-38	13-24
log_probs_time (s)	5.7-10.7	3.8-7.4
update_weights_time (s)	4.8-5.0	3.0
train tok/s	3.5k-4.9k	5.3k-8.0k
train/rollout logprob abs diff	0.032-0.041	0.032-0.039

torchtitan's dynamic batching needed no code: the shared get_batch already packs samples into one row with per-sample position ids, and torchtitan's post_dataloading_process builds the packed-document mask from those positions. The remaining gap is kernels (TE fused attention and grouped GEMM against flex attention and torch grouped_mm) and 2 s of weight sync.

Runs: torchtitan, Megatron.

The same comparison on 8x H200

TP2 EP4, colocated, two 4-GPU SGLang engines (20 attention heads do not divide by 8), otherwise the same flags as above. Steady-state steps 1-5.

per step	torchtitan	Megatron
actor_train_time (s)	10.1-15.5	5.6-9.4
log_probs_time (s)	3.2-4.5	1.8-3.0
update_weights_time (s)	3.0-3.2	2.3-2.5
train tok/s	8.5k-11.8k	13.7k-18.9k
train/rollout logprob abs diff	0.031-0.039	0.032-0.039

Runs: torchtitan group cmp_glm47_h200_titan_tp2ep4_r3_8k, Megatron group cmp_glm47_h200_megatron_tp2ep4_r3_8k.

What was tried to close the gap, and what stuck

experiment	result
compile.enable=True (per-block torch.compile, R3 router kept eager)	no change: train 26.6/20.9/38.3 s vs 24.2/20.7/38.0 s; the time is in the kernels, not eager overhead
flex attention Flash backend (flex_flash)	refused on SM100: head_dim=256 does not support block sparsity
varlen attention (torch 2.13 routes to FA4)	crashes inside torch's FA4 glue on this shape
gather each grouped expert tensor once, split locally	kept (last commit): update_weights 4.8-5.0 s -> 4.1-4.2 s

On Blackwell, MLA's 192 + 64 query head is the wall for every flex/FA path torchtitan exposes; Megatron's advantage on this model is TE fused attention and grouped GEMM, which support head_dim 256.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[Rockdu]

LGTM