From c6a305b0e744e8bc3d5152eeb55d75d57c9437c2 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 2 Jun 2026 06:19:02 +0000 Subject: [PATCH 1/5] [Feature] Add MHA head-wise output gate and Step3.5 clipped SwiGLU Add a per-head attention output gate (head_gate) to MHAConfig/MultiHeadAttention: a dedicated g_proj: Linear(hidden, num_heads) whose per-head sigmoid scales the attention output before o_proj. Distinct from the existing per-element with_gate (fused into a doubled q_proj); the two are mutually exclusive. Add a swiglu_clip activation (silu(gate).clamp(max=limit) * up.clamp(+/-limit)) and an optional swiglu_limit on MoEMLP, for models that clip the SwiGLU on a subset of layers. Unlike clipped_swiglu (gpt-oss, sigmoid-GLU + (up+1)), the clamp is applied after the silu activation. --- xtuner/v1/module/attention/mha.py | 21 +++++++++++++++++++ .../module/decoder_layer/moe_decoder_layer.py | 16 +++++++++++--- xtuner/v1/ops/act_fn.py | 16 ++++++++++++++ 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/xtuner/v1/module/attention/mha.py b/xtuner/v1/module/attention/mha.py index 707e978055..06903d9cf6 100644 --- a/xtuner/v1/module/attention/mha.py +++ b/xtuner/v1/module/attention/mha.py @@ -43,6 +43,7 @@ class MHAConfig(BaseModel): sliding_window: Annotated[int | None, Parameter(group="attention")] = -1 with_sink: Annotated[bool, Parameter(group="attention")] = False with_gate: Annotated[bool, Parameter(group="attention")] = False + head_gate: Annotated[bool, Parameter(group="attention")] = False attn_impl: Literal["flash_attention", "flex_attention", "eager_attention"] = "flash_attention" def model_post_init(self, _): @@ -128,6 +129,7 @@ def __init__( o_bias: bool = False, with_sink: bool = False, with_gate: bool = False, + head_gate: bool = False, attn_impl: Literal["flash_attention", "flex_attention", "eager_attention"] = "flash_attention", rope_scaling_cfg: RopeScalingConfig | None = None, float8_cfg: Float8Config | None = None, @@ -155,6 +157,8 @@ def __init__( self.float8_cfg = float8_cfg self.layer_idx = layer_idx self.with_gate = with_gate + self.head_gate = head_gate + assert not (with_gate and head_gate), "`with_gate` and `head_gate` are mutually exclusive." self.q_proj = build_linear( self.hidden_size, @@ -183,6 +187,17 @@ def __init__( float8_cfg=self.float8_cfg, ) + if self.head_gate: + # Head-wise output gate: one sigmoid scalar per attention head (broadcast over head_dim), + # produced by a dedicated projection. Distinct from `with_gate`, which fuses a + # per-(head, head_dim) gate into a doubled q_proj. + self.g_proj = build_linear( + self.hidden_size, + self.num_attention_heads, + bias=self.qkv_bias, + float8_cfg=self.float8_cfg, + ) + if self.qk_norm: self.q_norm = RMSNorm(self.head_dim, eps=self.rms_norm_eps, type=self.rms_norm_type) self.k_norm = RMSNorm(self.head_dim, eps=self.rms_norm_eps, type=self.rms_norm_type) @@ -430,6 +445,12 @@ def forward( if self.with_gate: assert gate is not None raw_output = raw_output * torch.sigmoid(gate) + elif self.head_gate: + # Per-head sigmoid gate (one scalar per head), broadcast over head_dim. + gate_score = torch.sigmoid(self.g_proj(hidden_states)) # [b, seq, n_head] + raw_output = ( + raw_output.view(*input_shape, self.num_attention_heads, self.head_dim) * gate_score.unsqueeze(-1) + ).reshape(*input_shape, -1) projected_output = self.o_proj(raw_output) attn_outputs: AttnOutputs = { diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 19a3d6371c..dabc81435c 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -52,7 +52,7 @@ def __call__(self, fused_x: torch.Tensor, split_dim: int = -1) -> torch.Tensor: class MoEActFnConfig(BaseModel): model_config = ConfigDict(extra="forbid") - act_type: Literal["clipped_swiglu", "swiglu"] = "swiglu" + act_type: Literal["clipped_swiglu", "swiglu", "swiglu_clip"] = "swiglu" clip_alpha: float | None = None clip_limit: float | None = None @@ -62,6 +62,8 @@ def build(self) -> MoEActFnProtocol: if self.act_type == "clipped_swiglu": act_fn = partial(act_fn, alpha=self.clip_alpha, limit=self.clip_limit) + elif self.act_type == "swiglu_clip": + act_fn = partial(act_fn, limit=self.clip_limit) return act_fn @@ -75,6 +77,7 @@ def __init__( hidden_act: str, mlp_bias: bool = False, float8_cfg: Float8Config | None = None, + swiglu_limit: float | None = None, ): super().__init__() self.hidden_size = hidden_size @@ -83,10 +86,17 @@ def __init__( self.up_proj = build_linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, float8_cfg=float8_cfg) self.down_proj = build_linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, float8_cfg=float8_cfg) self.act_fn = get_act_fn(hidden_act) + # When set, clamp the activated gate (max) and the raw up projection (±limit) before the + # element-wise product, matching Step3.5's per-layer SwiGLU clipping on the shared expert. + self.swiglu_limit = swiglu_limit def forward(self, x: torch.Tensor) -> torch.Tensor: - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj + gate = self.act_fn(self.gate_proj(x)) + up = self.up_proj(x) + if self.swiglu_limit is not None: + gate = gate.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + return self.down_proj(gate * up) class MoEGate(nn.Module): diff --git a/xtuner/v1/ops/act_fn.py b/xtuner/v1/ops/act_fn.py index b5818be1a8..11b9d581ee 100644 --- a/xtuner/v1/ops/act_fn.py +++ b/xtuner/v1/ops/act_fn.py @@ -24,6 +24,20 @@ def native_clipped_swiglu(fused_x: torch.Tensor, split_dim=-1, alpha=1.702, limi return gated_output +def native_swiglu_clip(fused_x: torch.Tensor, split_dim=-1, limit: float = 7.0) -> torch.Tensor: + # Step3.5-style clipped SwiGLU: the clamp is applied *after* the silu activation on the + # gate and on the raw up projection. This differs from gpt-oss `clipped_swiglu`, which + # clamps pre-activation and uses a sigmoid-GLU with a `(up + 1)` term. + gate, up = torch.chunk(fused_x, 2, dim=split_dim) + gate = F.silu(gate).clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + return gate * up + + +def npu_swiglu_clip(fused_x: torch.Tensor, split_dim=-1, limit: float = 7.0) -> torch.Tensor: + raise NotImplementedError + + def native_gelu(x: torch.Tensor, approximate: str | None = None) -> torch.Tensor: if approximate is not None: return F.gelu(x, approximate=approximate) @@ -49,6 +63,7 @@ def native_silu(x: torch.Tensor) -> torch.Tensor: act_fn_type_map_cuda = { "swiglu": native_swiglu, "clipped_swiglu": native_clipped_swiglu, + "swiglu_clip": native_swiglu_clip, "gelu": native_gelu, "gelu_pytorch_tanh": partial(native_gelu, approximate="tanh"), "silu": native_silu, @@ -56,6 +71,7 @@ def native_silu(x: torch.Tensor) -> torch.Tensor: act_fn_type_map_npu = { "swiglu": npu_swiglu, "clipped_swiglu": npu_clipped_swiglu, + "swiglu_clip": npu_swiglu_clip, "gelu": npu_gelu, "gelu_pytorch_tanh": partial(npu_gelu, approximate="tanh"), "silu": native_silu, From 91a9b2d92ad639aada4a27d47be25252895289ec Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 2 Jun 2026 06:19:28 +0000 Subject: [PATCH 2/5] [Feature] Add Step-3.5-Flash (step3p5) MoE model Port stepfun-ai/Step-3.5-Flash, a trust_remote_code MoE: 45 layers (first 3 dense, rest MoE), 288 routed experts top-8 + 1 shared expert, hybrid attention (full vs sliding with different head counts), head-wise attention gate, qk_norm, zero-centered RMSNorm, sigmoid NoAuxRouter with per-expert bias and scaling, and per-layer SwiGLU clipping on the last layers. Two architecture aspects are contained in the model rather than generalized into the base for now (to be generalized after precision alignment): the two attention profiles are selected per layer via layers_type in an overridden build_layers, and the per-profile RoPE (full: theta 5e6 / partial 0.5 / llama3; sliding: theta 1e4 / partial 1.0 / default) lives inside Step3.5-specific decoder layers that recompute position embeddings from seq_ctx.position_ids. HF stores experts as separate gate_proj/up_proj 3-D tensors; safetensors_to_params de-interleaves them into XTuner's fused expert-major grouped-linear weight on load. hf_config returns None (no built-in transformers config class), so save_hf copies the source config/tokenizer/modeling files. Design doc: docs/design/model/step3p5.md. --- .dev_scripts/convert_step3p5_to_split.py | 124 ++++++ docs/design/model/step3p5.md | 147 +++++++ xtuner/v1/model/__init__.py | 6 + xtuner/v1/model/moe/step3p5.py | 491 +++++++++++++++++++++++ 4 files changed, 768 insertions(+) create mode 100644 .dev_scripts/convert_step3p5_to_split.py create mode 100644 docs/design/model/step3p5.md create mode 100644 xtuner/v1/model/moe/step3p5.py diff --git a/.dev_scripts/convert_step3p5_to_split.py b/.dev_scripts/convert_step3p5_to_split.py new file mode 100644 index 0000000000..ee00fe4b3c --- /dev/null +++ b/.dev_scripts/convert_step3p5_to_split.py @@ -0,0 +1,124 @@ +"""Convert a Step-3.5-Flash HF checkpoint to a *split* per-expert layout. + +The released checkpoint stores each MoE layer's experts as three fused 3-D tensors +(`moe.gate_proj/up_proj/down_proj.weight`, shape `(num_experts, *, *)`). XTuner fuses experts +expert-major-interleaved for its grouped GEMM, and the load/save path can only shard a fused +parameter when each HF key is a *contiguous* slice of it. A single fused weight that maps to two +separate HF tensors (gate, up) therefore cannot be sharded across ranks. + +This script explodes the fused expert tensors into per-expert 2-D tensors under +`moe.experts.{i}.{gate,up,down}_proj.weight` (Qwen3-MoE style). With that layout XTuner's +`to_hf_key_list` emits the interleaved key order `[gate_0, up_0, gate_1, up_1, ...]`, which lines up +with its expert-major fused weight, so the default sharded load/save works on any number of +GPUs (FSDP and EP) with no per-model checkpoint code. + +The MTP layers (`model.layers.45..47`) are dropped — XTuner does not load them. + +Usage: + python .dev_scripts/convert_step3p5_to_split.py +""" + +import json +import re +import shutil +import sys +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + + +SHARD_BYTES = 4 * 1024**3 # ~4GB shards, HF-standard +# Side files needed so the converted dir is loadable (AutoConfig trust_remote_code + tokenizer). +AUX_FILES = [ + "config.json", + "configuration_step3p5.py", + "modeling_step3p5.py", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "generation_config.json", +] +DROP_LAYER_RE = re.compile(r"^model\.layers\.(4[5-9]|[5-9]\d)\.") # MTP / out-of-range layers +EXPERT_RE = re.compile(r"^(model\.layers\.\d+)\.moe\.(gate_proj|up_proj|down_proj)\.weight$") + + +def _iter_converted_tensors(src: Path): + """Yield (new_key, tensor) for every kept tensor, exploding fused experts into per-expert keys.""" + index = json.loads((src / "model.safetensors.index.json").read_text()) + weight_map = index["weight_map"] + # Group keys by source shard so each shard is opened once. + by_file: dict[str, list[str]] = {} + for key, fname in weight_map.items(): + by_file.setdefault(fname, []).append(key) + + for fname in sorted(by_file): + with safe_open(str(src / fname), framework="pt") as f: + for key in by_file[fname]: + if DROP_LAYER_RE.match(key): + continue + tensor = f.get_tensor(key) + m = EXPERT_RE.match(key) + if m is None: + yield key, tensor + continue + prefix, proj = m.group(1), m.group(2) + # gate/up: (n, inter, hidden) -> per expert (inter, hidden) + # down: (n, hidden, inter) -> per expert (hidden, inter) + for i in range(tensor.shape[0]): + yield f"{prefix}.moe.experts.{i}.{proj}.weight", tensor[i].contiguous() + + +def convert(src: Path, dst: Path) -> None: + dst.mkdir(parents=True, exist_ok=True) + weight_map: dict[str, str] = {} + buffer: dict[str, torch.Tensor] = {} + buffer_bytes = 0 + shard_idx = 1 + shards: list[tuple[str, dict[str, torch.Tensor]]] = [] + + def flush(): + nonlocal buffer, buffer_bytes, shard_idx + if not buffer: + return + name = f"model-{shard_idx:05d}.safetensors" + shards.append((name, buffer)) + for k in buffer: + weight_map[k] = name + buffer = {} + buffer_bytes = 0 + shard_idx += 1 + + for key, tensor in _iter_converted_tensors(src): + buffer[key] = tensor + buffer_bytes += tensor.numel() * tensor.element_size() + if buffer_bytes >= SHARD_BYTES: + flush() + flush() + + total = sum(t.numel() * t.element_size() for _, b in shards for t in b.values()) + n_shards = len(shards) + renamed: list[tuple[str, dict[str, torch.Tensor]]] = [] + for i, (_, buf) in enumerate(shards, start=1): + final = f"model-{i:05d}-of-{n_shards:05d}.safetensors" + for k in buf: + weight_map[k] = final + renamed.append((final, buf)) + for name, buf in renamed: + save_file(buf, str(dst / name), metadata={"format": "pt"}) + print(f" wrote {name} ({len(buf)} tensors)") + + (dst / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": total}, "weight_map": weight_map}, indent=2) + ) + for aux in AUX_FILES: + srcf = src / aux + if srcf.exists(): + shutil.copy2(srcf, dst / aux) + print(f"done: {len(weight_map)} tensors across {n_shards} shards, {total / 1024**3:.1f} GiB -> {dst}") + + +if __name__ == "__main__": + convert(Path(sys.argv[1]), Path(sys.argv[2])) diff --git a/docs/design/model/step3p5.md b/docs/design/model/step3p5.md new file mode 100644 index 0000000000..531d587f8f --- /dev/null +++ b/docs/design/model/step3p5.md @@ -0,0 +1,147 @@ +# Step-3.5-Flash → XTuner integration design + +Source: `stepfun-ai/Step-3.5-Flash`, `model_type = "step3p5"`, remote-code +(`modeling_step3p5.py` / `configuration_step3p5.py`, `architectures = ["Step3p5ForCausalLM"]`). +bf16 checkpoint, training-suitable. + +Bucket: **MoE LLM, trust_remote_code** (no built-in `transformers` config class → +`hf_config` returns `None`; `save_hf` copies `config.json` / tokenizer / `*.py`). + +## 1. Architecture summary (from the HF modeling code) + +- 45 transformer layers. Layers 0–2 are **dense** MLP; layers 3–44 are **MoE**. +- Vocab 128896, hidden 4096, untied `lm_head` (separate tensor in the index). +- MoE: 288 routed experts, top-k 8, `moe_intermediate_size` 1280, plus **one shared + expert** (`share_expert_dim` 1280). Experts stored as **fused 3-D** tensors + `(num_experts, out, in)` — `moe.{gate_proj,up_proj}.weight (288,1280,4096)`, + `moe.down_proj.weight (288,4096,1280)`. +- Router: **sigmoid** activation, per-expert **`router_bias`** added before top-k, + weights gathered from the **pre-bias** probabilities, renormalized, then scaled by + `moe_router_scaling_factor = 3.0`; router logits computed in **fp32** + (`need_fp32_gate`). +- RMSNorm is **zero-centered** (scale = `weight + 1`), eps 1e-5, throughout. +- Attention is a **hybrid of two softmax-attention profiles keyed by `layer_types`**: + - `full_attention` (every 4th layer, idx % 4 == 0): 64 heads, 8 KV heads, head_dim 128. + - `sliding_attention` (the other layers): **96 heads**, 8 KV heads, head_dim 128, + sliding window 512. + - Both profiles: `qk_norm` on head_dim (zero-centered), and a **head-wise output + gate** — a separate `g_proj: Linear(hidden, num_heads)` whose per-head sigmoid + multiplies the attention output before `o_proj`. +- RoPE differs **per profile**: + - full_attention: `rope_theta = 5e6`, `partial_rotary_factor = 0.5`, **llama3** + scaling (`yarn_only_types = ["full_attention"]`). + - sliding_attention: `rope_theta = 1e4`, `partial_rotary_factor = 1.0`, default rope + (no scaling). +- swiglu clamp limits on a few late layers only (MoE layers 43,44 → 7; shared expert + layer 44 → 16; all others 0/None). +- MTP: `num_nextn_predict_layers = 3`, stored as `model.layers.45/46/47.*` and ignored + on load by HF. We **drop MTP** for the port (load layers 0–44, `strict=False`). + +## 2. Mapping to XTuner — what is reused vs. new + +Reused as-is (config wiring only): + +| Feature | XTuner mechanism | +|---|---| +| first 3 dense + rest MoE | `MoEConfig.first_k_dense_replace = 3` | +| 288 experts / top-8 / shared expert | `n_routed_experts`, `num_experts_per_tok`, `n_shared_experts = 1` | +| expert tensors | **split / per-expert checkpoint** (`.dev_scripts/convert_step3p5_to_split.py`) → `to_hf_key_list` emits interleaved `[gate_i, up_i, …]` (Qwen3-MoE style); default load/save shards on FSDP+EP. See "Expert layout" below. | +| sigmoid + router_bias + renorm + scale 3.0 + fp32 gate | `NoAuxRouterConfig(scoring_func="sigmoid", n_group=1, topk_group=1, norm_topk_prob=True, router_scaling_factor=3.0)` + `router_compute_dtype="float32"` (its math matches HF `router_bias_func`) | +| zero-centered RMSNorm, qk_norm | `rms_norm_type="zero_centered"`, `MHAConfig.qk_norm=True` | +| sliding window in training | `MHAConfig.sliding_window` + `layer_type="sliding_attention"` | +| partial rotary | `RopeParametersConfig.partial_rotary_factor` (already supported) | +| MTP | `mtp_config=None`, load `strict=False` | + +Three things the current design **cannot express** — these are the design forks: + +### Fork A — head-wise attention gate (new MHA option) + +XTuner's existing `MHAConfig.with_gate` is a **per-(head,dim) element** gate fused into a +doubled `q_proj` (Qwen3.5 / gpt-oss style). Step-3.5 uses a **separate `g_proj` of shape +`(num_heads, hidden)`** producing **one scalar per head**. Different weight layout and +different broadcast. Proposal: add a new, general option to `MHAConfig` +(`head_gate: bool = False`) that builds `self.g_proj = Linear(hidden, num_heads, +bias=False)` and applies `out.view(...,H,Dh) * g.sigmoid().unsqueeze(-1)` before `o_proj`. +HF `self_attn.g_proj.weight` maps 1→1 to xtuner `self_attn.g_proj.weight`. This is a small, +self-contained addition at the module layer; its own commit. + +### Fork B — two attention profiles with different head counts (CONFIRMED) + +Today a model has a single `attention: MHAConfig` shared by all full/sliding layers +(`linear_attention` is GatedDeltaNet-only). Step-3.5 needs **full=64 heads, +sliding=96 heads**. Decision (user): the existing `layers_type` mechanism already supports +mixed attention; add a `sliding_attention: MHAConfig` field on the **Step-3.5 config +only** and **override `build_layers`** in the Step-3.5 model to select the per-layer +attention config by `layers_type`. No change to `base.py` / `MoEConfig`. +(`num_attention_heads` etc. computed fields keep reading the `full` profile — fine for +training; kv-cache / generate with mixed head counts is out of scope for the baseline.) + +### Fork C — per-profile RoPE (CONFIRMED — contain in Step-3.5 decoder layer) + +Today the model builds one `self.rotary_emb` and passes one `(cos, sin)` to every layer. +Step-3.5 needs different `(theta, partial_rotary, scaling)` for full vs sliding layers. +Decision (user): keep it **inside the Step-3.5 decoder layer** for now; generalize later +once precision is aligned. Each Step-3.5 decoder layer holds its own +`Step3p5RotaryEmbedding` (full: theta 5e6 / partial 0.5 / **llama3** scaling; sliding: +theta 1e4 / partial 1.0 / default) built faithfully via HF's +`ROPE_INIT_FUNCTIONS[rope_type]` on a per-profile shim so inv_freq matches HF bitwise. Its +`forward` recomputes `position_embeddings` from `seq_ctx.position_ids` and passes them to +`self.self_attn`, **ignoring** the model-level `(cos,sin)` (which stays valid but unused). +No change to shared `MoE.forward` or `MultiHeadAttention.forward`. The per-layer partial-rotary +apply is set **directly on the attention** in the decoder-layer `__init__` +(`self.self_attn.apply_rotary_emb = get_apply_rotary_emb(None, enable_partial_rotary=…)`) rather than +threading a rope config through `build` — this keeps the per-layer RoPE fully contained and avoids the +deprecated `RopeScalingConfig` (whose only consumer in `MultiHeadAttention` is that apply selection). + +### Swiglu clamp (CONFIRMED — include now) + +Step clamps `silu(gate).clamp(max=limit) * up.clamp(±limit)` on a few late layers +(routed experts L43/L44 → 7; shared expert L44 → 16; all others none). XTuner's existing +`clipped_swiglu` is gpt-oss-shaped (sigmoid-GLU + `(up+1)`) and does **not** match. Add a +new act variant `swiglu_clip` (silu + post-activation clamp) to `act_fn.py` / +`MoEActFnConfig`, and build a **per-layer** `MoEActFnConfig` (clip only where the config +lists a nonzero limit) in `build_layers`. For the shared expert clamp (L44), thread an +optional clamp limit into the shared-expert MLP. MTP (3 nextn layers) remains deferred +(load layers 0–44, `strict=False`). + +### Expert layout — split / per-expert checkpoint (CONFIRMED) + +The released checkpoint stores each MoE layer's experts as **three fused 3-D tensors** +(`moe.{gate,up,down}_proj.weight`, `(num_experts, *, *)`). XTuner fuses experts +**expert-major-interleaved** (`[g0,u0,g1,u1,…]`, each expert's `[gate;up]` contiguous — required by the +grouped GEMM). XTuner's loader can only shard a fused parameter when each HF key is a *contiguous* +slice of it, so a single fused `w1w3` mapped to two HF tensors (`gate`, `up`) **cannot be sharded**: +a 2-GPU FSDP load was empirically confirmed to crash (`gate, up = safetensors` receives 1 tensor), +and `save_hf` corrupted gate/up (`(288,640,4096)` vs `(288,1280,4096)` — the FUSED save split runs +before `param_to_safetensor`). This is a real XTuner limitation, not specific to this model. + +Decision (user): **convert the checkpoint to a split / per-expert layout** offline rather than change +the shared MoE block. `.dev_scripts/convert_step3p5_to_split.py` explodes each fused expert tensor +into `moe.experts.{i}.{gate,up,down}_proj.weight` (Qwen3-MoE style) and drops the unused MTP layers +(45–47). With that layout `to_hf_key_list` emits the interleaved key order `[gate_0, up_0, …]`, which +lines up with the expert-major fused weight, so the **default** `safetensors_to_params` (concat dim 0) +and the default save split both shard correctly on any number of GPUs (FSDP and EP) — **no per-model +checkpoint override and no MoE-block change**. Converted checkpoint: +`/mnt/shared-storage-user/llmrazor-share/yehaochen/model/Step-3.5-Flash-split`. + +## 3. File layout + +- `xtuner/v1/model/moe/step3p5.py` — `Step3p5MoEConfig` (+ base) and `Step3p5MoE` + (`to_hf_key_list`, `safetensors_to_params`/`param_to_safetensor`, `build_layers`, + `build_rotary_embedding`, `hf_config -> None`), plus `Step3p5Attention` + + `Step3p5RotaryEmbedding` (or co-located in the module layer if cleaner). +- `xtuner/v1/module/attention/mha.py` — Fork A (`head_gate`) + Fork C (unpack hook). +- `xtuner/v1/model/__init__.py` — import / `model_mapping` alias / `get_model_config_from_hf` + dispatch on `model_type == "step3p5"` / `__all__`. +- `tests/model/test_step3p5_moe.py` — baseline tests (§6/§7 of the skill). +- `ci/config/step3p5.py` — drop-in training config. + +## 4. Commit plan (stacked, dependency-ordered) + +1. `[Feature]` MHA head-wise gate (`head_gate`) + position-embeddings unpack hook (Forks A & C seam). +2. `[Feature]` Step-3.5 model + config + registration (Forks B & C, router/expert mapping). +3. `[Feature]` `XTUNER_HF_IMPL` parity wiring if any new op branch is needed. +4. `[Test]` baseline tests (decoder-layer bitwise parity for one full + one sliding layer; + whole-model forward+backward parity at the deployable scale; save_hf round-trip). +5. `[CI]` drop-in training config + convergence trace. +6. Then §8 optimizations (EP/SP, compile, fp8, offload) — multi-agent, post-baseline. diff --git a/xtuner/v1/model/__init__.py b/xtuner/v1/model/__init__.py index b8d794300f..a20d746306 100644 --- a/xtuner/v1/model/__init__.py +++ b/xtuner/v1/model/__init__.py @@ -25,6 +25,7 @@ from .moe.gpt_oss import GptOss21BA3P6Config, GptOss117BA5P8Config, GptOssConfig from .moe.moe import BalancingLossConfig, MoE, MoEConfig, MoEModelOutputs, ZLossConfig from .moe.qwen3 import Qwen3MoE30BA3Config, Qwen3MoEConfig, Qwen3MoEFoPEConfig +from .moe.step3p5 import Step3p5FlashConfig, Step3p5MoEConfig model_mapping = { @@ -39,6 +40,7 @@ "internvl-3.5-1b-hf": InternVL3P5Dense1BConfig(), "internvl-3.5-30b-a3b-hf": InternVL3P5MoE30BA3Config(), "qwen3.5-vl-4b": Qwen3_5_VLDense4BConfig(), + "step-3.5-flash": Step3p5FlashConfig(), } @@ -63,6 +65,8 @@ def get_model_config_from_hf(model_path: Path): return GptOssConfig.from_hf(model_path) elif cfg.model_type == "deepseek_v3": return DeepSeekV3Config.from_hf(model_path) + elif cfg.model_type == "step3p5": + return Step3p5MoEConfig.from_hf(model_path) else: raise ValueError(f"Unsupported model type: {cfg.model_type}") @@ -75,6 +79,8 @@ def get_model_config_from_hf(model_path: Path): "Qwen3Dense8BConfig", "Qwen3MoEConfig", "Qwen3MoE30BA3Config", + "Step3p5MoEConfig", + "Step3p5FlashConfig", "InternS1Config", "InternS1MiniConfig", "InternS1BaseConfig", diff --git a/xtuner/v1/model/moe/step3p5.py b/xtuner/v1/model/moe/step3p5.py new file mode 100644 index 0000000000..5e619a41ef --- /dev/null +++ b/xtuner/v1/model/moe/step3p5.py @@ -0,0 +1,491 @@ +import math +import re +from pathlib import Path +from typing import Literal + +import torch +import torch.nn as nn +from pydantic import Field +from typing_extensions import Self + +from transformers.models.auto import AutoConfig +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.model.moe.moe import ( + MoE, + MoEConfig, +) +from xtuner.v1.module import RouterResults +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer +from xtuner.v1.module.decoder_layer.moe_decoder_layer import HiddenStates, MoEActFnConfig, MoEDecoderLayer +from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig +from xtuner.v1.ops import get_apply_rotary_emb +from xtuner.v1.utils import get_logger + + +logger = get_logger() + + +class Step3p5RotaryEmbedding(nn.Module): + """Per-profile rotary embedding for Step3.5. + + Step3.5 uses a different RoPE configuration for ``full_attention`` and ``sliding_attention`` + layers (different ``rope_theta``, ``partial_rotary_factor`` and — only on full-attention layers + — llama3 frequency smoothing). The model-level rotary in :class:`MoE` cannot express this, so each + Step3.5 decoder layer owns its own instance and recomputes ``(cos, sin)`` from + ``seq_ctx.position_ids``. Inverse frequencies are computed faithfully against HuggingFace + (``ROPE_INIT_FUNCTIONS`` for llama3 is replicated inline so the result is independent of the + installed transformers version). + + Args: + head_dim (int): Per-head dimension. + rope_theta (float): RoPE base wavelength. + partial_rotary_factor (float): Fraction of ``head_dim`` that is rotated. + max_position_embeddings (int): Maximum sequence length (kept for parity with HF caching). + llama3_cfg (dict | None): llama3 smoothing parameters (``factor``, ``low_freq_factor``, + ``high_freq_factor``, ``original_max_position_embeddings``) or ``None`` for default RoPE. + """ + + inv_freq: torch.Tensor + + def __init__( + self, + head_dim: int, + rope_theta: float, + partial_rotary_factor: float, + max_position_embeddings: int, + llama3_cfg: dict | None = None, + ) -> None: + super().__init__() + # Stored as attributes (positional-or-keyword signature) so the test scaffolding's + # `materialize_submodule` can re-instantiate and recompute `inv_freq` after `to_empty`. + self.head_dim = head_dim + self.rope_theta = rope_theta + self.partial_rotary_factor = partial_rotary_factor + self.max_position_embeddings = max_position_embeddings + self.llama3_cfg = llama3_cfg + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) + if llama3_cfg is not None: + inv_freq = self._apply_llama3_smoothing(inv_freq, **llama3_cfg) + self.attention_scaling = 1.0 + self.max_seq_len_cached = max_position_embeddings + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float().to(x.device)).transpose( + 1, 2 + ) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + @staticmethod + def _apply_llama3_smoothing( + inv_freq: torch.Tensor, + *, + factor: float, + low_freq_factor: float, + high_freq_factor: float, + original_max_position_embeddings: int, + ) -> torch.Tensor: + # Mirrors transformers `_compute_llama3_parameters` smoothing applied on top of the base inv_freq. + low_freq_wavelen = original_max_position_embeddings / low_freq_factor + high_freq_wavelen = original_max_position_embeddings / high_freq_factor + wavelen = 2 * math.pi / inv_freq + inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq) + smooth_factor = (original_max_position_embeddings / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama + is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) + inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + return inv_freq_llama + + +def _set_partial_rotary(self_attn, rotary_emb: Step3p5RotaryEmbedding) -> None: + # The per-layer RoPE profile (full=partial 0.5, sliding=1.0) is owned by the Step3.5 decoder + # layer, so the matching partial-rotary apply is set directly on the attention here instead of + # threading a rope config through `build`. This keeps the per-layer RoPE fully contained and + # avoids the deprecated `RopeScalingConfig` (whose only consumer in `MultiHeadAttention` is this + # `apply_rotary_emb` selection). Only `MultiHeadAttention` has `apply_rotary_emb`. + if hasattr(self_attn, "apply_rotary_emb"): + self_attn.apply_rotary_emb = get_apply_rotary_emb( + None, enable_partial_rotary=rotary_emb.partial_rotary_factor != 1.0 + ) + + +class Step3p5DenseDecoderLayer(DenseDecoderLayer): + """Dense decoder layer that recomputes RoPE per layer (Step3.5 + first_k_dense layers).""" + + def __init__(self, *, rotary_emb: Step3p5RotaryEmbedding, **kwargs) -> None: + super().__init__(**kwargs) + self.rotary_emb = rotary_emb + _set_partial_rotary(self.self_attn, rotary_emb) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + seq_ctx: SequenceContext, + ) -> torch.Tensor: + # Ignore the model-level position embeddings: this layer's RoPE profile differs per layer type. + position_embeddings = self.rotary_emb(hidden_states, seq_ctx.position_ids) + return super().forward(hidden_states, position_embeddings, seq_ctx) + + +class Step3p5MoEDecoderLayer(MoEDecoderLayer): + """MoE decoder layer that recomputes RoPE per layer and clamps the shared + expert SwiGLU.""" + + def __init__( + self, + *, + rotary_emb: Step3p5RotaryEmbedding, + shared_swiglu_limit: float | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.rotary_emb = rotary_emb + _set_partial_rotary(self.self_attn, rotary_emb) + if shared_swiglu_limit is not None: + assert self.shared_experts is not None + self.shared_experts.swiglu_limit = shared_swiglu_limit + + def forward( + self, + *hidden_states: torch.Tensor, + seq_ctx: SequenceContext | list[SequenceContext], + position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, + ) -> tuple[HiddenStates, RouterResults] | tuple[torch.Tensor, ...]: + if len(hidden_states) == 1: + assert isinstance(seq_ctx, SequenceContext) + position_embeddings = self.rotary_emb(hidden_states[0], seq_ctx.position_ids) + else: + assert isinstance(seq_ctx, list) + position_embeddings = [self.rotary_emb(h, sc.position_ids) for h, sc in zip(hidden_states, seq_ctx)] + return super().forward(*hidden_states, seq_ctx=seq_ctx, position_embeddings=position_embeddings) + + +class Step3p5MoE(MoE): + def build_rotary_embedding(self, config: "Step3p5MoEConfig"): # type: ignore[override] + # Model-level rotary is unused (each decoder layer owns its own), but `MoE.forward` still + # computes it; return the full-attention profile so the call is valid. CPU init keeps it + # correct even when the model is built on the meta device (see BaseModel.build_rotary_embedding). + with torch.device("cpu"): + return config.build_layer_rotary("full_attention") + + def build_layers(self, config: "Step3p5MoEConfig") -> nn.ModuleDict: # type: ignore[override] + from xtuner.v1.model.utils.misc import module_dict_repr + + layers = nn.ModuleDict() + for layer_idx in range(config.num_hidden_layers): + layer_type = config.layers_type[layer_idx] + attention_config = config.attention if layer_type == "full_attention" else config.sliding_attention + # CPU init so the rotary buffers are real even under a meta-device model build. The + # per-layer partial-rotary apply is set from this rotary inside the decoder layer (so no + # rope config is threaded through `build` — see `_set_partial_rotary`). + with torch.device("cpu"): + rotary_emb = config.build_layer_rotary(layer_type) + + if layer_idx < config.first_k_dense_replace: + layers[str(layer_idx)] = Step3p5DenseDecoderLayer( + rotary_emb=rotary_emb, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + mlp_bias=config.mlp_bias, + hidden_act=config.hidden_act, + rms_norm_eps=config.rms_norm_eps, + rms_norm_type=config.rms_norm_type, + attention_config=attention_config, + layer_type=layer_type, + rope_scaling_cfg=None, + generate_config=config.generate_config, + float8_cfg=config.float8_cfg, + layer_idx=layer_idx, + ) + else: + layers[str(layer_idx)] = Step3p5MoEDecoderLayer( + rotary_emb=rotary_emb, + shared_swiglu_limit=config.shared_swiglu_limits.get(layer_idx), + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + moe_intermediate_size=config.moe_intermediate_size, + mlp_bias=config.mlp_bias, + gate_bias=config.gate_bias, + moe_bias=config.moe_bias, + hidden_act=config.hidden_act, + rms_norm_eps=config.rms_norm_eps, + rms_norm_type=config.rms_norm_type, + num_experts_per_tok=config.num_experts_per_tok, + n_routed_experts=config.n_routed_experts, + n_shared_experts=config.n_shared_experts, + with_shared_expert_gate=config.with_shared_expert_gate, + hidden_factor=config.hidden_factor, + layer_type=layer_type, + attention_config=attention_config, + rope_scaling_cfg=None, + generate_config=config.generate_config, + router_config=config.router, + router_compute_dtype=config.router_compute_dtype, + moe_act_fn_cfg=config.layer_moe_act_fn_cfg(layer_idx), + float8_cfg=config.float8_cfg, + layer_idx=layer_idx, + dispatcher=config.dispatcher, + ep_mesh=self.ep_mesh, + ) + if config.freeze_routers: + layers[str(layer_idx)].gate.requires_grad_(False) + layers[str(layer_idx)].gate.eval() + + layers.__class__.__repr__ = module_dict_repr # type: ignore[method-assign] + return layers + + def to_hf_key_list(self, key: str) -> list[str]: + if "layers" in key or "embed_tokens" in key: + key = "model." + key + + if key.startswith("norm."): + return ["model." + key] + if key == "lm_head.weight": + return ["lm_head.weight"] + + # Routed experts: the fused XTuner grouped-linear maps to per-expert HF tensors (Qwen3-MoE + # style). The checkpoint stores experts *split* per expert under `moe.experts.{i}.*`; the + # interleaved key order `[gate_0, up_0, gate_1, up_1, ...]` matches XTuner's expert-major + # fused layout, so the default `safetensors_to_params` (concat along dim 0) and the default + # save split both shard correctly across FSDP/EP without any per-model override. + n = self.config.n_routed_experts + if "experts.fused_w1w3.weight" in key: + out: list[str] = [] + for i in range(n): + out.append(key.replace("experts.fused_w1w3.weight", f"moe.experts.{i}.gate_proj.weight")) + out.append(key.replace("experts.fused_w1w3.weight", f"moe.experts.{i}.up_proj.weight")) + return out + if "experts.fused_w2.weight" in key: + return [key.replace("experts.fused_w2.weight", f"moe.experts.{i}.down_proj.weight") for i in range(n)] + + # Router linear + per-expert bias. + if "gate.router.e_score_correction_bias" in key: + return [key.replace("gate.router.e_score_correction_bias", "moe.router_bias")] + if re.search(r"layers\.\d+\.gate\.weight$", key): + return [key.replace("gate.weight", "moe.gate.weight")] + + # Shared expert: XTuner `shared_experts.*` -> HF `share_expert.*`. + if "shared_experts." in key: + return [key.replace("shared_experts.", "share_expert.")] + + return [key] + + +class Step3p5MoEConfig(MoEConfig): + model_type: str | None = "step3p5" + bos_token_id: int | None = None + rms_norm_type: Literal["default", "zero_centered"] = "zero_centered" + + # --- Step3.5-specific extra fields --- + # Second attention profile: sliding-window layers use a different head count than full layers. + sliding_attention: MHAConfig = Field( + default_factory=lambda: MHAConfig( + num_attention_heads=96, + num_key_value_heads=8, + head_dim=128, + qk_norm=True, + head_gate=True, + rms_norm_type="zero_centered", + rms_norm_eps=1e-5, + sliding_window=512, + ) + ) + # Per-profile RoPE. + full_rope_theta: float = 5_000_000.0 + sliding_rope_theta: float = 10_000.0 + full_partial_rotary_factor: float = 0.5 + sliding_partial_rotary_factor: float = 1.0 + rope_factor: float = 2.0 + rope_low_freq_factor: float = 1.0 + rope_high_freq_factor: float = 32.0 + rope_original_max_position_embeddings: int = 131072 + # Sparse per-layer SwiGLU clip limits: {layer_idx: limit}. `routed` clamps the routed experts' + # activation; `shared` clamps the shared expert. Layers absent from the dict are unclamped. + routed_swiglu_limits: dict[int, float] = {} + shared_swiglu_limits: dict[int, float] = {} + + def build(self) -> Step3p5MoE: + return Step3p5MoE(self) + + # `layers_type` follows Step3.5's fixed pattern: every 4th layer is full attention, the rest sliding. + @property + def layers_type(self) -> list[Literal["full_attention", "sliding_attention", "linear_attention"]]: # type: ignore[override] + return ["full_attention" if i % 4 == 0 else "sliding_attention" for i in range(self.num_hidden_layers)] + + def build_layer_rotary(self, layer_type: str) -> Step3p5RotaryEmbedding: + if layer_type == "full_attention": + return Step3p5RotaryEmbedding( + head_dim=self.attention.head_dim, + rope_theta=self.full_rope_theta, + partial_rotary_factor=self.full_partial_rotary_factor, + max_position_embeddings=self.max_position_embeddings, + llama3_cfg={ + "factor": self.rope_factor, + "low_freq_factor": self.rope_low_freq_factor, + "high_freq_factor": self.rope_high_freq_factor, + "original_max_position_embeddings": self.rope_original_max_position_embeddings, + }, + ) + return Step3p5RotaryEmbedding( + head_dim=self.sliding_attention.head_dim, + rope_theta=self.sliding_rope_theta, + partial_rotary_factor=self.sliding_partial_rotary_factor, + max_position_embeddings=self.max_position_embeddings, + llama3_cfg=None, + ) + + def layer_moe_act_fn_cfg(self, layer_idx: int) -> MoEActFnConfig: + limit = self.routed_swiglu_limits.get(layer_idx) + if limit is not None: + return MoEActFnConfig(act_type="swiglu_clip", clip_limit=limit) + return self.moe_act_fn_cfg + + @classmethod + def from_hf(cls, hf_path: str | Path) -> Self: + hf_config = AutoConfig.from_pretrained(hf_path, trust_remote_code=True) + assert hf_config.model_type == "step3p5", f"Expected model_type 'step3p5', got '{hf_config.model_type}'" + + head_dim = hf_config.head_dim + full_heads = hf_config.num_attention_heads + kv_heads = hf_config.num_attention_groups + sliding_setting = hf_config.attention_other_setting + rope_scaling = hf_config.rope_scaling + + num_layers = hf_config.num_hidden_layers + swiglu_limits = list(getattr(hf_config, "swiglu_limits", []))[:num_layers] + swiglu_limits_shared = list(getattr(hf_config, "swiglu_limits_shared", []))[:num_layers] + routed_swiglu_limits = {i: float(v) for i, v in enumerate(swiglu_limits) if v} + shared_swiglu_limits = {i: float(v) for i, v in enumerate(swiglu_limits_shared) if v} + + attention = MHAConfig( + num_attention_heads=full_heads, + num_key_value_heads=kv_heads, + head_dim=head_dim, + qk_norm=hf_config.use_qk_norm, + head_gate=hf_config.use_head_wise_attn_gate, + rms_norm_type="zero_centered", + rms_norm_eps=hf_config.rms_norm_eps, + sliding_window=-1, + ) + sliding_attention = MHAConfig( + num_attention_heads=sliding_setting["num_attention_heads"], + num_key_value_heads=sliding_setting["num_attention_groups"], + head_dim=sliding_setting["head_dim"], + qk_norm=hf_config.use_qk_norm, + head_gate=hf_config.use_head_wise_attn_gate, + rms_norm_type="zero_centered", + rms_norm_eps=hf_config.rms_norm_eps, + sliding_window=hf_config.sliding_window, + ) + + # First MoE layer index -> number of leading dense layers. + moe_layers = [int(i) for i in str(hf_config.moe_layers_enum).split(",")] + first_k_dense_replace = min(moe_layers) + + # HF stores eos as a list of acceptable ids; TransformerConfig wants a single int. + eos = hf_config.eos_token_id + eos_token_id = int(eos[0]) if isinstance(eos, (list, tuple)) else int(eos) + + return cls( + vocab_size=hf_config.vocab_size, + max_position_embeddings=hf_config.max_position_embeddings, + pad_token_id=getattr(hf_config, "pad_token_id", None), + bos_token_id=getattr(hf_config, "bos_token_id", None), + eos_token_id=eos_token_id, + num_hidden_layers=num_layers, + hidden_size=hf_config.hidden_size, + intermediate_size=hf_config.intermediate_size, + rms_norm_eps=hf_config.rms_norm_eps, + hidden_act="silu", + tie_word_embeddings=getattr(hf_config, "tie_word_embeddings", False), + attention=attention, + sliding_attention=sliding_attention, + n_routed_experts=hf_config.moe_num_experts, + n_shared_experts=1, + num_experts_per_tok=hf_config.moe_top_k, + first_k_dense_replace=first_k_dense_replace, + moe_intermediate_size=hf_config.moe_intermediate_size, + with_shared_expert_gate=False, + router=NoAuxRouterConfig( + scoring_func="sigmoid", + n_group=1, + topk_group=1, + norm_topk_prob=hf_config.norm_expert_weight, + router_scaling_factor=hf_config.moe_router_scaling_factor, + ), + router_compute_dtype="float32", + full_rope_theta=hf_config.rope_theta[0], + sliding_rope_theta=hf_config.rope_theta[1], + full_partial_rotary_factor=hf_config.partial_rotary_factors[0], + sliding_partial_rotary_factor=hf_config.partial_rotary_factors[1], + rope_factor=rope_scaling["factor"], + rope_low_freq_factor=rope_scaling["low_freq_factor"], + rope_high_freq_factor=rope_scaling["high_freq_factor"], + rope_original_max_position_embeddings=rope_scaling["original_max_position_embeddings"], + routed_swiglu_limits=routed_swiglu_limits, + shared_swiglu_limits=shared_swiglu_limits, + ) + + @property + def hf_config(self): + # Step3.5 ships as a `trust_remote_code` model with no built-in transformers config class, so + # `save_hf` falls back to copying config.json / tokenizer / *.py from the source dir. + return None + + +class Step3p5FlashConfig(Step3p5MoEConfig): + vocab_size: int = 128896 + max_position_embeddings: int = 262144 + pad_token_id: int | None = None + bos_token_id: int | None = 0 + eos_token_id: int = 128007 + num_hidden_layers: int = 45 + hidden_size: int = 4096 + intermediate_size: int = 11264 + rms_norm_eps: float = 1e-5 + hidden_act: str = "silu" + tie_word_embeddings: bool = False + attention: MHAConfig = Field( + default_factory=lambda: MHAConfig( + num_attention_heads=64, + num_key_value_heads=8, + head_dim=128, + qk_norm=True, + head_gate=True, + rms_norm_type="zero_centered", + rms_norm_eps=1e-5, + sliding_window=-1, + ) + ) + n_routed_experts: int = 288 + n_shared_experts: int = 1 + num_experts_per_tok: int = 8 + first_k_dense_replace: int = 3 + moe_intermediate_size: int = 1280 + with_shared_expert_gate: bool = False + router_compute_dtype: Literal["float32", "native"] = "float32" + router: NoAuxRouterConfig = Field( + default_factory=lambda: NoAuxRouterConfig( + scoring_func="sigmoid", + n_group=1, + topk_group=1, + norm_topk_prob=True, + router_scaling_factor=3.0, + ) + ) + routed_swiglu_limits: dict[int, float] = Field(default_factory=lambda: {43: 7.0, 44: 7.0}) + shared_swiglu_limits: dict[int, float] = Field(default_factory=lambda: {44: 16.0}) From f7717c32c780b5f58a8106aa7b5920391eebb9e9 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 2 Jun 2026 06:19:28 +0000 Subject: [PATCH 3/5] [Test] Add Step-3.5-Flash parity tests Single-rank baseline parity vs HuggingFace (STEP3P5_PATH env var): - rotary inv_freq bitwise vs the canonical default / llama3 formulas; - attention sub-block bitwise (head gate + qk_norm + per-layer RoPE + partial rotary + sliding window) for full / sliding / clamp layers under XTUNER_HF_IMPL; - full MoE decoder layer within tolerance (bf16 grouped GEMM vs HF fp32 expert loop) with exact router top-k indices. The shipped modeling_step3p5.py is incompatible with the installed transformers (its rotary init crashes), so the test replaces only the HF rotary with a version-independent implementation using the canonical formulas. --- tests/model/test_step3p5_moe.py | 266 ++++++++++++++++++++++++++++++++ 1 file changed, 266 insertions(+) create mode 100644 tests/model/test_step3p5_moe.py diff --git a/tests/model/test_step3p5_moe.py b/tests/model/test_step3p5_moe.py new file mode 100644 index 0000000000..b9dd1bbd72 --- /dev/null +++ b/tests/model/test_step3p5_moe.py @@ -0,0 +1,266 @@ +import itertools +import math +import os +import re +import unittest + +import parametrize +import torch +import torch.distributed as dist +import torch.nn as nn + +from xtuner.v1.data_proto import SequenceContext +from xtuner._testing import DeterministicDDPTestCase + + +# Step-3.5-Flash (MoE, hybrid full/sliding attention, head-wise gate, per-layer RoPE). +STEP3P5_PATH = os.environ.get("STEP3P5_PATH") + + +def _patch_hf_step3p5(hf_path, hf_cfg, llama3): + """Make the trust_remote_code modeling importable/runnable on the installed transformers. + + The shipped ``modeling_step3p5.py`` targets a different transformers version: its + ``Step3p5RotaryEmbedding`` crashes here (``ROPE_INIT_FUNCTIONS['default']`` missing; llama3 reads a + list ``rope_theta``). We replace only the rotary with a version-independent implementation that uses + the canonical default / llama3 inverse-frequency formulas — the same math XTuner uses and that the + rotary-parity test pins to the reference. Everything else (attention gate, qk_norm, MoE router, + experts, SwiGLU clamp, zero-centered RMSNorm) stays genuine HF. + """ + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + mod = __import__( + get_class_from_dynamic_module("modeling_step3p5.Step3p5Attention", hf_path).__module__, + fromlist=["x"], + ) + # Capture per-layer rope constants once: HF mutates the config during layer construction. + theta = list(hf_cfg.rope_theta) + partial = list(hf_cfg.partial_rotary_factors) + yarn = list(hf_cfg.yarn_only_types) + layer_types = list(hf_cfg.layer_types) + head_dim = hf_cfg.head_dim + + def _llama3(inv): + old = llama3["original_max_position_embeddings"] + lo, hi, f = llama3["low_freq_factor"], llama3["high_freq_factor"], llama3["factor"] + low_wl, high_wl = old / lo, old / hi + wl = 2 * math.pi / inv + iv = torch.where(wl > low_wl, inv / f, inv) + sm = (old / wl - lo) / (hi - lo) + sa = (1 - sm) * iv / f + sm * iv + return torch.where(~(wl < high_wl) * ~(wl > low_wl), sa, iv) + + class _FixedRotary(nn.Module): + def __init__(self, config, device=None, layer_idx=None): + super().__init__() + dim = int(head_dim * partial[layer_idx]) + inv = 1.0 / (theta[layer_idx] ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) + if layer_types[layer_idx] in yarn: + inv = _llama3(inv) + self.attention_scaling = 1.0 + self.register_buffer("inv_freq", inv, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + inv = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + pos = position_ids[:, None, :].float() + with torch.autocast(device_type=x.device.type, enabled=False): + freqs = (inv.float().to(x.device) @ pos.float().to(x.device)).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(x.dtype), emb.sin().to(x.dtype) + + mod.Step3p5RotaryEmbedding = _FixedRotary + return mod, _FixedRotary + + +@unittest.skipIf(STEP3P5_PATH is None, "STEP3P5_PATH env var is required for Step-3.5 parity tests") +class TestStep3p5MoE(DeterministicDDPTestCase): + def _build(self): + from transformers import AutoConfig + + from xtuner.v1.model import get_model_config_from_hf + + cfg = get_model_config_from_hf(STEP3P5_PATH) + cfg.compile_cfg = False + hf_cfg = AutoConfig.from_pretrained(STEP3P5_PATH, trust_remote_code=True) + hf_cfg._attn_implementation = "eager" + llama3 = { + "factor": cfg.rope_factor, + "low_freq_factor": cfg.rope_low_freq_factor, + "high_freq_factor": cfg.rope_high_freq_factor, + "original_max_position_embeddings": cfg.rope_original_max_position_embeddings, + } + return cfg, hf_cfg, llama3 + + @parametrize.parametrize("device", [("cuda",)]) + def test_rotary_inv_freq_parity(self, device): + # Per-profile inverse frequencies must match the canonical default / llama3 formulas bitwise. + self.create_pg(device) + cfg, _, _ = self._build() + + full = cfg.build_layer_rotary("full_attention").inv_freq + dim = int(cfg.attention.head_dim * cfg.full_partial_rotary_factor) + base = 1.0 / (cfg.full_rope_theta ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) + from xtuner.v1.model.moe.step3p5 import Step3p5RotaryEmbedding + + ref_full = Step3p5RotaryEmbedding._apply_llama3_smoothing( + base, + factor=cfg.rope_factor, + low_freq_factor=cfg.rope_low_freq_factor, + high_freq_factor=cfg.rope_high_freq_factor, + original_max_position_embeddings=cfg.rope_original_max_position_embeddings, + ) + self.assertTrue(torch.equal(full, ref_full), "full-attention inv_freq not bitwise vs llama3 reference") + + sliding = cfg.build_layer_rotary("sliding_attention").inv_freq + sdim = int(cfg.sliding_attention.head_dim * cfg.sliding_partial_rotary_factor) + ref_sliding = 1.0 / (cfg.sliding_rope_theta ** (torch.arange(0, sdim, 2, dtype=torch.int64).float() / sdim)) + self.assertTrue(torch.equal(sliding, ref_sliding), "sliding-attention inv_freq not bitwise vs default reference") + dist.barrier() + + @parametrize.parametrize("device,layer_idx", [("cuda", 4), ("cuda", 3), ("cuda", 44)]) + def test_decoder_attention_bitwise_parity(self, device, layer_idx): + # Attention sub-block (q/k/v + qk_norm + per-layer RoPE + partial rotary + head-wise gate + + # sliding window) must match HF bitwise under XTUNER_HF_IMPL (eager). Covers a full layer (4), + # a sliding layer (3), and a full layer with SwiGLU clamp (44). Memory-light: only the attention + # of one layer is materialized on each side. + from xtuner.v1.utils import HFCheckpointLoader + + self.create_pg(device) + with self.hf_impl(): + cfg, hf_cfg, llama3 = self._build() + loader = HFCheckpointLoader(STEP3P5_PATH) + mod, FixedRotary = _patch_hf_step3p5(STEP3P5_PATH, hf_cfg, llama3) + layer_type = cfg.layers_type[layer_idx] + + with torch.device("meta"): + model = cfg.build() + xt_layer = model.layers[str(layer_idx)] + xt_attn = xt_layer.self_attn + xt_attn.to_empty(device="cuda") + xt_attn.to(torch.bfloat16) + for n, p in xt_attn.named_parameters(): + key = model.to_hf_key_list(f"layers.{layer_idx}.self_attn.{n}")[0] + p.data.copy_(loader.load(key).to(p.device, p.dtype)) + xt_layer.rotary_emb.inv_freq = cfg.build_layer_rotary(layer_type).inv_freq.to("cuda") + + hf_attn = mod.Step3p5Attention(hf_cfg, layer_idx).to("cuda").to(torch.bfloat16).eval() + for n, p in hf_attn.named_parameters(): + p.data.copy_(loader.load(f"model.layers.{layer_idx}.self_attn.{n}").to(p.device, p.dtype)) + # `.to(bf16)` rounds the fp32 inv_freq buffer; restore the fp32 value (kept fp32 in production). + hf_attn.rotary_emb.inv_freq = FixedRotary(hf_cfg, layer_idx=layer_idx).inv_freq.to("cuda") + + seq = 16 + torch.manual_seed(0) + ids = torch.randint(0, 1000, (1, seq), device="cuda") + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids,)) + seq_ctx.to("cuda") + x = torch.randn(1, seq, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + + o_xt = xt_attn(x, xt_layer.rotary_emb(x, seq_ctx.position_ids), seq_ctx)["projected_output"] + mask = torch.triu( + torch.full((seq, seq), float("-inf"), device="cuda", dtype=torch.bfloat16), diagonal=1 + )[None, None] + o_hf, _ = hf_attn(x, attention_mask=mask, position_ids=seq_ctx.position_ids) + + diff = (o_hf.float() - o_xt.float().reshape(o_hf.shape)).abs().max().item() + self.assertEqual(diff, 0.0, f"layer {layer_idx} [{layer_type}] attention not bitwise: max diff {diff}") + dist.barrier() + + @parametrize.parametrize("device,layer_idx", [("cuda", 4), ("cuda", 44)]) + def test_decoder_layer_tolerance_parity(self, device, layer_idx): + # Full MoE decoder layer (attention + router + fused experts + shared expert + SwiGLU clamp). + # Not bitwise: XTuner uses a bf16 grouped GEMM while the HF reference upcasts each expert matmul + # to fp32 (`MoELinear.float()`); the established MoE bar is tolerance (see test_qwen3_moe). The + # router top-k *indices* must still match exactly (both compute sigmoid+bias+top-k in fp32). + from xtuner.v1.utils import HFCheckpointLoader + + self.create_pg(device) + with self.hf_impl(): + cfg, hf_cfg, llama3 = self._build() + loader = HFCheckpointLoader(STEP3P5_PATH) + mod, FixedRotary = _patch_hf_step3p5(STEP3P5_PATH, hf_cfg, llama3) + layer_type = cfg.layers_type[layer_idx] + + with torch.device("meta"): + model = cfg.build() + xt_layer = model.layers[str(layer_idx)] + xt_layer.to_empty(device="cuda") + for p in xt_layer.parameters(): + p.data = p.data.to(torch.bfloat16) # params bf16; the fp32 router-bias buffer is untouched + xt_layer.rotary_emb.inv_freq = cfg.build_layer_rotary(layer_type).inv_freq.to("cuda") + self._load_xt_layer(model, xt_layer, layer_idx, loader) + + hf_layer = mod.Step3p5DecoderLayer(hf_cfg, layer_idx).to("cuda").eval() + for p in hf_layer.parameters(): + p.data = p.data.to(torch.bfloat16) + for n, p in hf_layer.named_parameters(): + if "router_bias" in n: + p.data = p.data.to(torch.float32) # need_fp32_gate + n_exp = cfg.n_routed_experts + + def _hf_tensor(name): + # The HF reference layer expects fused 3-D expert tensors, but the checkpoint stores + # experts split per expert; restack them for the HF side. + m = re.match(r"moe\.(gate_proj|up_proj|down_proj)\.weight$", name) + if m is not None: + proj = m.group(1) + return torch.stack( + [loader.load(f"model.layers.{layer_idx}.moe.experts.{i}.{proj}.weight") for i in range(n_exp)] + ) + return loader.load(f"model.layers.{layer_idx}.{name}") + + for n, p in itertools.chain( + hf_layer.named_parameters(), + ((n, b) for n, b in hf_layer.named_buffers() if not n.endswith("inv_freq")), + ): + p.data.copy_(_hf_tensor(n).to(p.device, p.dtype)) + hf_layer.self_attn.rotary_emb.inv_freq = FixedRotary(hf_cfg, layer_idx=layer_idx).inv_freq.to("cuda") + + seq = 64 + torch.manual_seed(0) + ids = torch.randint(0, 1000, (1, seq), device="cuda") + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids,)) + seq_ctx.to("cuda") + x = torch.randn(1, seq, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + + o_xt = xt_layer(x, seq_ctx=seq_ctx, position_embeddings=None) + o_xt = o_xt[0] if isinstance(o_xt, tuple) else o_xt + mask = torch.triu( + torch.full((seq, seq), float("-inf"), device="cuda", dtype=torch.bfloat16), diagonal=1 + )[None, None] + o_hf = hf_layer(x, attention_mask=mask, position_ids=seq_ctx.position_ids) + o_hf = o_hf[0] if isinstance(o_hf, tuple) else o_hf + + # Router top-k indices: exact match. + xt_topk = xt_layer.gate(x)["topk_ids"].sort(dim=-1).values + hf_logits = (x.view(-1, cfg.hidden_size).float() @ hf_layer.moe.gate.weight.t().float()) + hf_prob = hf_logits.sigmoid() + _, hf_idx = torch.topk(hf_prob + hf_layer.moe.router_bias.unsqueeze(0), k=cfg.num_experts_per_tok, dim=1) + hf_topk = hf_idx.sort(dim=-1).values + + rel = (o_hf.float() - o_xt.float().reshape(o_hf.shape)).abs().max().item() / o_hf.float().abs().max().item() + self.assertTrue(torch.equal(xt_topk, hf_topk), f"layer {layer_idx} router top-k indices differ") + self.assertLess(rel, 0.02, f"layer {layer_idx} [{layer_type}] full-layer rel error too high: {rel}") + dist.barrier() + + @staticmethod + def _load_xt_layer(model, layer, layer_idx, loader): + # Load one XTuner decoder layer's params (and persistent router-bias buffer). The fused expert + # grouped-linear maps to many per-expert HF keys whose interleaved order matches the fused + # layout, so the default `safetensors_to_params` concatenates them along dim 0. + named = itertools.chain( + layer.named_parameters(), + ((n, b) for n, b in layer.named_buffers() if not n.endswith("inv_freq")), + ) + for name, tensor in named: + full = f"layers.{layer_idx}.{name}" + hf_keys = model.to_hf_key_list(full) + tensors = [loader.load(k) for k in hf_keys] + assert all(t is not None for t in tensors), (full, hf_keys) + # dim=0: per-expert keys concatenate along the fused grouped-linear's expert-major dim. + model.safetensors_to_params(tensors, tensor.data, full, None, None, 0) + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "1")) From 7a119eb6c5427f11a6380a39d60c4341354e8d02 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 2 Jun 2026 07:56:07 +0000 Subject: [PATCH 4/5] [CI] Add Step-3.5-Flash training config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop-in TrainerConfig for Step-3.5-Flash. `load_from`/`tokenizer_path` read STEP3P5_PATH, which must point at the split / per-expert checkpoint produced by `.dev_scripts/convert_step3p5_to_split.py` (the released fused-expert layout cannot be sharded across ranks). Uses expert parallelism (ep_size=8, all2all); torch.compile is left off (a §8 optimization for the hybrid per-layer-RoPE decoder layers). A reduced 8-GPU overfit smoke (5 layers, full 288 experts) confirmed the forward -> loss -> backward -> FSDP-reduce -> optimizer-step loop runs and the loss descends; a full convergence run on the ~200B model needs a multi-node cluster. --- ci/config/step3p5.py | 61 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 ci/config/step3p5.py diff --git a/ci/config/step3p5.py b/ci/config/step3p5.py new file mode 100644 index 0000000000..a87a993007 --- /dev/null +++ b/ci/config/step3p5.py @@ -0,0 +1,61 @@ +import os + +from xtuner.v1.config import ( + AdamWConfig, + FSDPConfig, + LRConfig, +) +from xtuner.v1.datasets import FTDPTokenizeFnConfig +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.moe.step3p5 import Step3p5FlashConfig +from xtuner.v1.train import TrainerConfig + + +# Point STEP3P5_PATH at the split / per-expert checkpoint produced by +# `.dev_scripts/convert_step3p5_to_split.py` (the released fused-expert layout cannot be sharded). +STEP3P5_PATH = os.environ["STEP3P5_PATH"] +ALPACA_PATH = os.environ["ALPACA_PATH"] + + +# Step-3.5-Flash is a ~200B MoE; real training needs expert parallelism (and a multi-node cluster). +moe_cfg = Step3p5FlashConfig(ep_size=8, dispatcher="all2all", num_hidden_layers=4) +optim_cfg = AdamWConfig(lr=6e-05) +lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6) +fsdp_cfg = FSDPConfig( + # torch.compile for the hybrid per-layer-RoPE decoder layers is a §8 optimization; keep eager here. + torch_compile=False, + cpu_offload=False, + ep_size=moe_cfg.ep_size, +) + +dataset_config = [ + { + "dataset": DatasetConfig(name="alpaca", anno_path=ALPACA_PATH, sample_ratio=1.0), + "tokenize_fn": FTDPTokenizeFnConfig(max_length=16386), + }, +] + +dataloader_config = DataloaderConfig(pack_max_length=16384) + +# Chunked cross-entropy keeps the logits->loss peak memory bounded (never materializes the full +# (seq, vocab) logits) — important for Step-3.5's 128896-token vocab. +loss_cfg = CELossConfig(mode="chunk") + + +trainer = TrainerConfig( + load_from=STEP3P5_PATH, + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + dataset_cfg=dataset_config, + dataloader_cfg=dataloader_config, + lr_cfg=lr_cfg, + loss_cfg=loss_cfg, + tokenizer_path=STEP3P5_PATH, + global_batch_size=16, + total_step=1000000, + work_dir="/tmp/step3p5", + seed=0, + strict_load=False, +) From 2569d96adf4103b4fa2409bbd6b561565f0f7bf8 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 2 Jun 2026 09:31:10 +0000 Subject: [PATCH 5/5] [Docs] Note in add_hf_model skill: avoid the deprecated RopeScalingConfig Record the caveat surfaced by the Step-3.5 port: RopeScalingConfig is deprecated (use RopeParametersConfig), and when a per-layer value only selects one module behavior (e.g. partial_rotary_factor -> apply_rotary_emb), set that behavior directly on the module instead of threading a (deprecated) rope config through build. --- .claude/skills/add_hf_model/SKILL.md | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/.claude/skills/add_hf_model/SKILL.md b/.claude/skills/add_hf_model/SKILL.md index ab06691d68..e1313f858a 100644 --- a/.claude/skills/add_hf_model/SKILL.md +++ b/.claude/skills/add_hf_model/SKILL.md @@ -80,6 +80,18 @@ attention / router / layer onto the matching one. - Also: rope, rms_norm, etc. - **ops layer** (`xtuner/v1/ops`) — kernels such as attention and rms_norm. +> **Caveat — don't reach for deprecated config classes.** `RopeScalingConfig` is +> **deprecated**; `RopeParametersConfig` (`xtuner/v1/module/rope/rope.py`, re-exported from +> `xtuner/v1/model/base.py`) is the source of truth — use it everywhere (your IDE/pyright will flag +> the deprecated one). Note the decoder-layer / `MHAConfig.build` signatures still *type* their rope +> argument as `RopeScalingConfig` for backward compatibility, so don't satisfy them by constructing +> the deprecated class. When a per-layer value is only needed to select **one module behavior** (e.g. +> `partial_rotary_factor` only chooses which `apply_rotary_emb` the attention uses), set that behavior +> **directly on the module** instead of threading a config through `build` — e.g. in your model's +> decoder layer, `self.self_attn.apply_rotary_emb = get_apply_rotary_emb(None, +> enable_partial_rotary=...)` (`xtuner/v1/ops`). This keeps per-layer behavior contained (the §C +> per-profile-RoPE pattern) and avoids the deprecated API entirely. + ### Existing models to copy from Pick the one whose attention + (router) match yours; the closer it is, the @@ -527,10 +539,13 @@ FSDP shard/reduce chain — and the file doubles as the example users copy when model into their own training pipeline. Mirror `ci/config/qwen3_moe_30BA3.py` (MoE) or `ci/config/qwen3_dense.py` (dense), keeping its structure: one `()` (from §3.2) fed into a `TrainerConfig` alongside `optim_cfg` / `lr_cfg` / `fsdp_cfg` / `dataset_cfg` / -`dataloader_cfg` / `loss_cfg`. `load_from` and `tokenizer_path` read from an env var (§7.4 — -typically the same one as the parity test). Verify by running ~50 steps and confirming the -loss drops monotonically into a plausible range for that model size; record the trajectory -in the PR body alongside the §6 convergence trace. +`dataloader_cfg` / `loss_cfg`. Set `loss_cfg = CELossConfig(mode="chunk")` — the chunked +cross-entropy keeps the `logits → loss` peak memory bounded (it never materializes the full +`(seq, vocab)` logits), which matters for the large-vocab models this skill targets; do **not** leave +it on the `"eager"` default. `load_from` and `tokenizer_path` read from an env var (§7.4 — typically +the same one as the parity test). Verify by running ~50 steps and confirming the loss drops +monotonically into a plausible range for that model size; record the trajectory in the PR body +alongside the §6 convergence trace. ---