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

docs/user-guide/training-backend.md

@@ -557,8 +557,10 @@ the config a torchtitan user would write by hand.
 --seq-length 16384
 ```
 
-`--titan-model-name` is a package under `torchtitan/models/`, and `--titan-model-flavor` one
-of the sizes that package registers. `--hf-checkpoint` supplies the weights and tokenizer;
+`--titan-model-name` names a model package and `--titan-model-flavor` one of the sizes it
+registers. The name is looked up first under `miles/backends/torchtitan_utils/models/`, then
+under `torchtitan/models/`, so a model torchtitan does not ship can be added on the miles side
+without touching torchtitan (see [below](#going-deeper-a-model-torchtitan-does-not-ship)). `--hf-checkpoint` supplies the weights and tokenizer;
 torchtitan's own state-dict adapter converts them, so there is no offline conversion step.
 
 `--seq-length` sizes the rotary tables and the buffers pipeline stages exchange, so it has
@@ -619,8 +621,9 @@ because that check compares the engine against the original HF checkpoint.
 
 <Warning>
 
-**The architecture must be one torchtitan implements.** There is no per-architecture spec to
-write, which also means there is no way to bring up a model torchtitan does not have.
+**The architecture must be one torchtitan implements or one registered under
+`miles/backends/torchtitan_utils/models/`.** There is no per-architecture spec to write; a
+new model is a Python package that assembles torchtitan's own blocks.
 
 **No LoRA, no optimizer CPU offload, no disk offload, no on-policy distillation, and
 `--ref-update-interval` is rejected** rather than silently ignored.
@@ -631,6 +634,34 @@ MoE works, including expert parallelism and `--use-rollout-routing-replay` (R3).
 more than it looks on MoE: without it, training and rollout routing drift apart as the policy
 moves, and the disagreement compounds over rollouts rather than staying flat.
 
+### Going deeper: a model torchtitan does not ship
+
+torchtitan finds a model by importing `torchtitan.models.<name>` and calling its
+`model_registry(flavor, attn_backend)`, which returns a `ModelSpec`: the model config, the
+parallelize and pipelining functions, and the state-dict adapter that maps HF names to
+torchtitan's. Miles looks in `miles/backends/torchtitan_utils/models/<name>/` first and expects
+exactly the same function, so a miles-side package is a drop-in peer of a torchtitan one.
+
+Most new architectures are a recombination of blocks torchtitan already has — MLA or GQA
+attention, the sigmoid or softmax token-choice router, grouped experts, shared experts — and
+then the package is a flavor table plus a `model_registry` that points at an existing
+`parallelize_fn` and adapter. `models/glm4_moe_lite/` is the worked example: GLM-4.7-Flash is
+DeepSeek-V3's layer with different dimensions and the same HF parameter names, so its package
+builds the layer list with torchtitan's DeepSeek-V3 helpers, sets the GLM sizes, and reuses
+`DeepSeekV3StateDictAdapter` unchanged.
+
+```bash
+--train-backend torchtitan \
+--titan-model-name glm4_moe_lite \
+--titan-model-flavor 30B-A3B \
+--hf-checkpoint /root/models/GLM-4.7-Flash
+```
+
+When the architecture genuinely differs, the package grows in this order: a `model.py` with
+the new block (subclassing torchtitan's `TransformerBlock` / `BaseAttention`), then a
+`state_dict_adapter.py` if the HF names differ, and only then a `parallelize.py` if the
+existing sharding plans do not apply. Two checks belong with every package: a fast test that pins the flavor's dimensions to the values recorded from the model's `config.json` and checks the adapter's map covers every key pattern recorded from `model.safetensors.index.json`, and an e2e case under `tests/e2e/torchtitan/` — SGLang has to implement the architecture too, since it is what consumes the streamed weights.
+
 ### Try it
 
 ```bash

miles/backends/torchtitan_utils/config.py

@@ -16,18 +16,28 @@
 logger = logging.getLogger(__name__)
 
 
+_MODEL_PACKAGE_ROOTS = ("miles.backends.torchtitan_utils.models", "torchtitan.models")
+
+
+def _model_package(name: str):
+    for root in _MODEL_PACKAGE_ROOTS:
+        module_name = f"{root}.{name}"
+        try:
+            return importlib.import_module(module_name)
+        except ModuleNotFoundError as e:
+            if e.name != module_name:
+                raise
+    raise ValueError(
+        f"--titan-model-name {name!r} is neither a miles model package nor a torchtitan one "
+        f"(looked for {', '.join(f'{root}.{name}' for root in _MODEL_PACKAGE_ROOTS)})"
+    )
+
+
 def resolve_model_spec(args: Namespace):
-    module_name = f"torchtitan.models.{args.titan_model_name}"
-    try:
-        module = importlib.import_module(module_name)
-    except ModuleNotFoundError as e:
-        raise ValueError(
-            f"--titan-model-name {args.titan_model_name!r} does not resolve to a torchtitan "
-            f"model package ({module_name}). Check the pinned torchtitan checkout."
-        ) from e
+    module = _model_package(args.titan_model_name)
     registry = getattr(module, "model_registry", None)
     if registry is None:
-        raise ValueError(f"{module_name} exposes no model_registry(); cannot build a ModelSpec")
+        raise ValueError(f"{module.__name__} exposes no model_registry(); cannot build a ModelSpec")
     return registry(args.titan_model_flavor, attn_backend="flex")
 
 

miles/backends/torchtitan_utils/hf_weight_iterator.py

@@ -3,6 +3,7 @@
 import json
 import os
 from collections.abc import Iterator
+from dataclasses import dataclass
 
 import safetensors
 import torch
@@ -14,6 +15,7 @@
 from miles.backends.training_utils.weight_update.hf_weight_iterator.checkpoint_towers import (
     iter_checkpoint_tower_units,
 )
+from miles.utils.hf_utils.config import load_hf_config
 
 
 class TitanHfWeightIterator(HfWeightIteratorBase):
@@ -22,6 +24,7 @@ class TitanHfWeightIterator(HfWeightIteratorBase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._engine_dtypes = _checkpoint_dtypes(self.args.hf_checkpoint)
+        self._q_lora_rank = getattr(load_hf_config(self.args.hf_checkpoint), "q_lora_rank", None) or None
 
     def _iter_hf_param_units(self, weights, *, materialize):
         for name, tensor in hf_weights(self.model, complete_across_pp=self.placement.gather_pp):
@@ -36,8 +39,7 @@ def _to_engine_dtype(self, name: str, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(target)
 
     def _hf_atomic_update_groups(self):
-        q_lora_rank = getattr(self.model.model_config, "q_lora_rank", None) or None
-        return get_hf_atomic_update_groups(self.model_name, q_lora_rank=q_lora_rank)
+        return get_hf_atomic_update_groups(self.model_name, q_lora_rank=self._q_lora_rank)
 
     def _iter_hf_adapter_units(self, lora_name, adapter, *, materialize):
         raise NotImplementedError("the torchtitan backend has no LoRA")
@@ -62,58 +64,88 @@ def _stage_process_groups(stage_ranks: tuple[tuple[int, ...], ...]) -> tuple:
     return tuple(dist.new_group(list(ranks)) for ranks in stage_ranks)
 
 
+_GROUPED_EXPERTS = "moe.routed_experts.inner_experts"
+
+
 def _hf_weights_on_device(trainer, *, complete_across_pp: bool) -> Iterator[tuple[str, torch.Tensor]]:
     sd_adapter = getattr(trainer.checkpointer, "sd_adapter", None)
     if sd_adapter is None:
         sd_adapter = trainer.config.model_spec.state_dict_adapter(trainer.model_config, trainer.config.hf_assets_path)
-    local = sd_adapter.to_hf({k: v for part in trainer.model_parts for k, v in part.state_dict().items()})
-
-    world = dist.get_world_size()
-    local_meta = {name: (tuple(t.shape), str(t.dtype)) for name, t in local.items()}
-    gathered: list = [None] * world
-    dist.all_gather_object(gathered, local_meta)
-
-    if all(meta.keys() == local_meta.keys() for meta in gathered):
-        for name in sorted(local):
-            yield name, gather_full_param(local[name])
-        return
-
-    owners: dict[str, list[int]] = {}
-    specs: dict[str, tuple] = {}
-    for rank, meta in enumerate(gathered):
-        for name, (shape, dtype) in meta.items():
-            owners.setdefault(name, []).append(rank)
-            if specs.setdefault(name, (shape, dtype)) != (shape, dtype):
-                raise RuntimeError(f"ranks disagree on the shape/dtype of {name}")
-
-    my_rank = dist.get_rank()
-    stage_of: list = [None] * world
-    pp_mesh = trainer.parallel_dims.get_optional_mesh("pp")
-    my_stage_id = dist.get_rank(group=pp_mesh.get_group()) if pp_mesh is not None else 0
-    dist.all_gather_object(stage_of, my_stage_id)
-    stage_groups: dict[int, list[int]] = {}
-    for rank, stage in enumerate(stage_of):
-        stage_groups.setdefault(stage, []).append(rank)
-    my_stage = stage_of[my_rank]
-
-    if complete_across_pp:
-        audience, broadcast_group = list(range(world)), None
-    else:
-        groups = _stage_process_groups(tuple(tuple(ranks) for _, ranks in sorted(stage_groups.items())))
-        audience, broadcast_group = stage_groups[my_stage], groups[sorted(stage_groups).index(my_stage)]
-
-    audience_set = set(audience)
-    names = [name for name in sorted(owners) if audience_set.intersection(owners[name])]
-    for name in names:
-        shape, dtype = specs[name]
-        holders = [rank for rank in owners[name] if rank in audience_set]
-        if my_rank in holders:
-            tensor = gather_full_param(local[name]).contiguous()
+    state = {k: v for part in trainer.model_parts for k, v in part.state_dict().items()}
+    layout = _StageLayout.from_trainer(trainer, complete_across_pp=complete_across_pp)
+
+    dense = {k: v for k, v in state.items() if _GROUPED_EXPERTS not in k}
+    yield from layout.stream(sd_adapter.to_hf(dense))
+
+    mine = sorted(k for k in state if _GROUPED_EXPERTS in k)
+    everyone: list = [None] * layout.world
+    dist.all_gather_object(everyone, mine)
+    for key in sorted(set().union(*everyone)):
+        yield from layout.stream(sd_adapter.to_hf({key: gather_full_param(state[key])}) if key in state else {})
+
+
+@dataclass(frozen=True)
+class _StageLayout:
+    world: int
+    my_rank: int
+    device: torch.device
+    audience: frozenset[int]
+    broadcast_group: dist.ProcessGroup | None
+
+    @classmethod
+    def from_trainer(cls, trainer, *, complete_across_pp: bool) -> "_StageLayout":
+        world = dist.get_world_size()
+        my_rank = dist.get_rank()
+        stage_of: list = [None] * world
+        pp_mesh = trainer.parallel_dims.get_optional_mesh("pp")
+        my_stage_id = dist.get_rank(group=pp_mesh.get_group()) if pp_mesh is not None else 0
+        dist.all_gather_object(stage_of, my_stage_id)
+        stage_groups: dict[int, list[int]] = {}
+        for rank, stage in enumerate(stage_of):
+            stage_groups.setdefault(stage, []).append(rank)
+        my_stage = stage_of[my_rank]
+        if complete_across_pp:
+            audience, broadcast_group = list(range(world)), None
         else:
-            tensor = torch.empty(shape, dtype=getattr(torch, dtype.split(".")[-1]), device=trainer.device)
-        dist.broadcast(tensor, src=holders[0], group=broadcast_group)
-        yield name, tensor
-        del tensor
+            groups = _stage_process_groups(tuple(tuple(ranks) for _, ranks in sorted(stage_groups.items())))
+            audience, broadcast_group = stage_groups[my_stage], groups[sorted(stage_groups).index(my_stage)]
+        return cls(
+            world=world,
+            my_rank=my_rank,
+            device=trainer.device,
+            audience=frozenset(audience),
+            broadcast_group=broadcast_group,
+        )
+
+    def stream(self, local: dict[str, torch.Tensor]) -> Iterator[tuple[str, torch.Tensor]]:
+        local_meta = {name: (tuple(t.shape), str(t.dtype)) for name, t in local.items()}
+        gathered: list = [None] * self.world
+        dist.all_gather_object(gathered, local_meta)
+
+        if all(meta.keys() == local_meta.keys() for meta in gathered):
+            for name in sorted(local):
+                yield name, gather_full_param(local[name])
+            return
+
+        owners: dict[str, list[int]] = {}
+        specs: dict[str, tuple] = {}
+        for rank, meta in enumerate(gathered):
+            for name, (shape, dtype) in meta.items():
+                owners.setdefault(name, []).append(rank)
+                if specs.setdefault(name, (shape, dtype)) != (shape, dtype):
+                    raise RuntimeError(f"ranks disagree on the shape/dtype of {name}")
+
+        names = [name for name in sorted(owners) if self.audience.intersection(owners[name])]
+        for name in names:
+            shape, dtype = specs[name]
+            holders = [rank for rank in owners[name] if rank in self.audience]
+            if self.my_rank in holders:
+                tensor = gather_full_param(local[name]).contiguous()
+            else:
+                tensor = torch.empty(shape, dtype=getattr(torch, dtype.split(".")[-1]), device=self.device)
+            dist.broadcast(tensor, src=holders[0], group=self.broadcast_group)
+            yield name, tensor
+            del tensor
 
 
 _SAFETENSORS_DTYPES = {

miles/backends/torchtitan_utils/models/glm4_moe_lite/__init__.py

@@ -0,0 +1,70 @@
+from torchtitan.components.optimizer import register_moe_load_balancing_hook
+from torchtitan.distributed.pipeline_parallel import pipeline_llm
+from torchtitan.models.common import ComplexRoPE, Embedding, Linear
+from torchtitan.models.common.nn_modules import RMSNorm
+from torchtitan.models.deepseek_v3 import (
+    _EMBEDDING_INIT,
+    _NORM_INIT,
+    DeepSeekV3Model,
+    _build_dsv3_layers,
+    _output_linear_init,
+    parallelize_deepseekv3,
+)
+from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter
+from torchtitan.protocols.model_spec import ModelSpec
+
+__all__ = ["glm4_moe_lite_configs", "model_registry"]
+
+
+def _30b_a3b(attn_backend: str) -> DeepSeekV3Model.Config:
+    dim = 2048
+    vocab_size = 154880
+    rope_dim = 64
+    layers = _build_dsv3_layers(
+        n_layers=47,
+        n_dense_layers=1,
+        dim=dim,
+        n_heads=20,
+        q_lora_rank=768,
+        kv_lora_rank=512,
+        qk_nope_head_dim=192,
+        qk_rope_head_dim=rope_dim,
+        v_head_dim=256,
+        mscale=1.0,
+        dense_hidden_dim=10240,
+        moe_hidden_dim=1536,
+        num_experts=64,
+        num_shared_experts=1,
+        router_top_k=4,
+        router_score_func="sigmoid",
+        router_route_scale=1.8,
+        router_route_norm=True,
+        attn_backend=attn_backend,
+        moe_comm_backend="standard",
+        non_blocking_capacity_factor=None,
+        rope=ComplexRoPE.Config(dim=rope_dim, max_seq_len=202752, theta=1_000_000.0),
+    )
+    return DeepSeekV3Model.Config(
+        vocab_size=vocab_size,
+        dim=dim,
+        tok_embeddings=Embedding.Config(num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT),
+        norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+        lm_head=Linear.Config(in_features=dim, out_features=vocab_size, param_init=_output_linear_init(dim)),
+        layers=layers,
+        mtp_layers=[],
+    )
+
+
+glm4_moe_lite_configs = {"30B-A3B": _30b_a3b}
+
+
+def model_registry(flavor: str, attn_backend: str = "flex") -> ModelSpec:
+    return ModelSpec(
+        name="glm4_moe_lite",
+        flavor=flavor,
+        model=glm4_moe_lite_configs[flavor](attn_backend=attn_backend),
+        parallelize_fn=parallelize_deepseekv3,
+        pipelining_fn=pipeline_llm,
+        post_optimizer_build_fn=register_moe_load_balancing_hook,
+        state_dict_adapter=DeepSeekV3StateDictAdapter,
+    )

tests/e2e/torchtitan/_common.py

@@ -38,6 +38,7 @@ class CaseConfig:
     with_ref: bool = False
     colocate: bool = True
     rollout_num_gpus: int | None = None
+    rollout_num_gpus_per_engine: int | None = None
     fully_async: bool = False
     mem_fraction_static: float = 0.7
     transfer_mode: str | None = None
@@ -68,6 +69,8 @@ def __post_init__(self):
 
     @property
     def engine_size(self) -> int:
+        if self.rollout_num_gpus_per_engine is not None:
+            return self.rollout_num_gpus_per_engine
         return self.num_gpus if self.colocate else self.rollout_num_gpus
 
     @property

tests/e2e/torchtitan/test_glm47_flash_tp_ep_r3.py

@@ -0,0 +1,34 @@
+from tests.ci.ci_register import register_cuda_ci
+from tests.e2e.torchtitan._common import CaseConfig, execute, prepare
+
+register_cuda_ci(est_time=2400, suite="stage-c-8-gpu-h200", labels=["torchtitan", "replay"], hardware=["hopper"])
+
+# The first model torchtitan itself does not ship: the flavor lives under
+# miles/backends/torchtitan_utils/models/ and reuses torchtitan's DeepSeek-V3
+# blocks (MLA attention, sigmoid router with the correction bias, shared
+# expert), so this case covers the miles-side registry end to end -- HF load
+# through the DeepSeek adapter, tensor and expert parallelism, routing replay, and the
+# weight stream back into SGLang's own glm4_moe_lite implementation. The engines
+# take four GPUs each because the model's 20 attention heads do not divide by eight. SGLang runs
+# attention on triton so the case is the same on Hopper and Blackwell: flashinfer's
+# SM100 prefill kernel has no head_dim 256 instantiation, which is what MLA's
+# 192 + 64 query head pads to.
+CASE = CaseConfig(
+    model_repo="zai-org/GLM-4.7-Flash",
+    titan_model_name="glm4_moe_lite",
+    titan_model_flavor="30B-A3B",
+    num_gpus=8,
+    tp_size=2,
+    ep_size=4,
+    rollout_num_gpus_per_engine=4,
+    seq_len=16384,
+    max_response_len=8192,
+    use_r3=True,
+    mem_fraction_static=0.5,
+    extra_args="--sglang-attention-backend triton --sglang-moe-runner-backend triton ",
+)
+
+
+if __name__ == "__main__":
+    prepare(CASE)
+    execute(CASE, wandb_file=__file__)