add-hf-model

# Add a HuggingFace model to XTuner A complete handbook for porting any HuggingFace causal LM into XTuner with correct training behavior. Covers four buckets: **Dense LLM**, **MoE LLM**, **VLM / compose model**, and **`trust_remote_code` model**. Every step below is grounded in concrete `file:line` references in this repository — read the cited code before writing yours. > **Anti-spaghetti reminder** — XTuner separates *what the model is* (`*Config`), > *what it computes* (`Dense` / `MoE` / `BaseComposeModel` subclass), and *how its > weights map to HF safetensors* (`to_hf_key_list`, optional `safetensors_to_params` > / `param_to_safetensor`). Keep new code inside those three seams. **Never** add > `if model_name == "..."` branches to the base classes. --- ## 0. Locate the source code The only goal of this step is to find the **HuggingFace reference implementation you will port from**, and to note the `model_type` you will register under. Do not pick attention/router/layer here — that happens naturally in §1 as you read the code. 1. If the user gives a **remote-code** model, the modeling code ships in the repo as `modeling_*.py` — read it directly. 2. If the user gives **weights only**, confirm the target Python environment with the user, then find the matching implementation inside that env's `transformers` install. If no code exists for that `model_type`, stop and ask the user. 3. Confirm the weights are **training-suitable** — at least bf16 — so training precision can be aligned later. Flag fp8/int4-only checkpoints to the user. 4. Note two strings from the HF config — you will need both downstream: ```python from transformers import AutoConfig cfg = AutoConfig.from_pretrained(path, trust_remote_code=True) cfg.model_type # registration key (§4) cfg.architectures[0] # the modeling class you are mirroring ``` --- ## 1. How to do the migration XTuner splits a model into three layers. You port by reusing what already exists at each layer and only adding what is genuinely new — so the work is "read the HF modeling code, then match each of its pieces to the XTuner component below". The locations are listed plainly; you do not need extra guidance to map your model's attention / router / layer onto the matching one. - **model layer** (`xtuner/v1/model`) — the model body. - Dense body: `xtuner/v1/model/dense/dense.py` - MoE body: `xtuner/v1/model/moe/moe.py` - Multi-modal (compose) bodies: `xtuner/v1/model/compose/` — reuse the language model and add only the ViT / projector pieces. - **module layer** (`xtuner/v1/module`): - Decoder layers: `DenseDecoderLayer`, `MoEDecoderLayer` — `xtuner/v1/module/decoder_layer`. Small changes add a switch on the existing layer; a genuinely novel architecture may add `{model_name}_decoder_layer`. - Attention configs: `MHAConfig` (`xtuner/v1/module/attention/mha.py`), `MLAConfig` (`.../mla.py`, DeepSeek-style latent), `GatedDeltaNetConfig` (`.../gated_deltanet.py`, linear). - Router configs (MoE): `GreedyRouterConfig` (`xtuner/v1/module/router/greedy.py`, softmax top-k — Qwen3 MoE, GptOss), `NoAuxRouterConfig` (`xtuner/v1/module/router/noaux_router.py`, grouped sigmoid aux-loss-free — DeepSeek V3). - Also: rope, rms_norm, etc. - **ops layer** (`xtuner/v1/ops`) — kernels such as attention and rms_norm. ### Existing models to copy from Pick the one whose attention + (router) match yours; the closer it is, the smaller your diff. | Your model is… | Copy from | |--------------------------------------|--------------------------------------| | Dense LLM | `xtuner/v1/model/dense/qwen3.py` | | MoE LLM | `xtuner/v1/model/moe/qwen3.py` | | MoE with MLA + aux-loss-free router | `xtuner/v1/model/moe/deepseek_v3.py` | | MoE needing weight-layout transform | `xtuner/v1/model/moe/gpt_oss.py` | | Compose / VLM | `xtuner/v1/model/compose/qwen3_vl/` | After you have a plan, **write an integration design doc into `docs/design/model`** and confirm with the user that the decomposition is reasonable before implementing. ### File layout for the new model | Bucket | Model file | Public re-export | Test file | |----------------------|-------------------------------------------------------------|-----------------------------------------------|--------------------------------------------| | Dense LLM | `xtuner/v1/model/dense/<name>.py` | `xtuner/v1/model/__init__.py` | `tests/model/test_<name>_dense.py` | | MoE LLM | `xtuner/v1/model/moe/<name>.py` | `xtuner/v1/model/__init__.py` | `tests/model/test_<name>_moe.py` | | Compose / VLM | `xtuner/v1/model/compose/<name>/...` | `xtuner/v1/model/__init__.py` | `tests/model/test_<name>.py` | --- ## 2. Implement the model class Subclass `Dense`, `MoE`, or `BaseComposeModel`. You typically only need to override one method. ### 2.1 `to_hf_key_list(self, key: str) -> list[str]` — **mandatory, abstract on `BaseModel`** Declared at `xtuner/v1/model/base.py`. Translates an **XTuner-side** parameter name into one or more **HF-side** safetensors keys. Return-list cardinality: - **1 → 1** for plain weights (most params). - **1 → N** for fused MoE experts (one fused param explodes into N per-expert HF keys; see `xtuner/v1/model/moe/qwen3.py`). Patterns you almost always need (study `xtuner/v1/model/dense/qwen3.py` and `xtuner/v1/model/moe/qwen3.py`): ```python # 1. Tied embeddings: redirect lm_head → embed_tokens before adding the model. prefix. if self.config.tie_word_embeddings and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") # 2. Add the "model." prefix that HF wraps everything in. if "layers" in key or "embed_tokens" in key: key = "model." + key # 3. HF MoE nests experts/gate under .mlp.; XTuner does not. if "layers" in key: key = re.sub(r"layers\.(\d+)\.(experts|gate)", r"layers.\1.mlp.\2", key) # 4. Top-level norm. if key.startswith("norm."): return [key.replace("norm.", "model.norm.")] # 5. MoE expert fusion: one fused param → N HF keys. if "fused_w1w3.weight" in key: out = [] for i in range(self.config.n_routed_experts): out.append(key.replace("fused_w1w3.weight", f"{i}.gate_proj.weight")) out.append(key.replace("fused_w1w3.weight", f"{i}.up_proj.weight")) return out # 6. Model-specific buffers (e.g. FoPE: rotary_emb.sin_coef / cos_coef are persistent buffers). if key.startswith("rotary_emb."): return [key.replace("rotary_emb.", "model.rotary_emb.")] ``` A prefix-region remap that differs by *deployment* (e.g. a text tower nested under a VLM as `model.language_model.`) does **not** belong here — put it in the config's `hf_key_mapping`. See §5.1. ### 2.2 `safetensors_to_params` / `param_to_safetensor` — **optional, only for layout mismatches** If the HF storage layout cannot be expressed as “same tensor, different name”, you must transform the bytes. Canonical example: GptOss stores expert weights as `(num_experts, hidden_size, expert_dim * 2)` but XTuner fuses them as `(num_experts * 2 * expert_dim, hidden_size)` — see overrides at `xtuner/v1/model/moe/gpt_oss.py`. Use this hook only when transposing/reshaping is unavoidable. Renames go in `to_hf_key_list`; numerics go here. ### 2.3 Do **not** override - `build_layers`, `build_embeddings`, `_init_weights` — the base handles them via the config. Override only if the new model adds a layer type the base cannot express (rare). - `forward` — same reasoning. If you find yourself touching it, re-evaluate the bucket choice; the problem is probably in the config (`attention`, `router`, `first_k_dense_replace`). --- ## 3. Implement the config Two layers: a **base config** that reads HF and a **size-specific subclass** per released checkpoint. ### 3.1 Base config — three required members Mirror `Qwen3MoEConfig` (`xtuner/v1/model/moe/qwen3.py`) or `Qwen3DenseConfig` (`xtuner/v1/model/dense/qwen3.py`). 1. **`build(self) -> <ModelClass>`** — return `<ModelClass>(self)`. 2. **`@classmethod from_hf(cls, hf_path) -> Self`** — read the HF config and map its fields one-to-one onto your config. The field set is whatever the live config class declares (`TransformerConfig` in `xtuner/v1/model/base.py`, `MoEConfig` in `xtuner/v1/model/moe/moe.py`); read it, don't work from a hard-coded list. A few non-obvious traps: - Use `RopeParametersConfig.from_hf_config(hf_config)` from `xtuner/v1/module/rope/rope.py`. It already handles both the `rope_parameters` dict (HF ≥ 5.2.0) and the legacy `rope_scaling` dict (HF 4.57.x), plus YARN / FoPE / mrope special fields. **Do not** parse rope yourself. - For optional HF fields, use `getattr(hf_config, "<field>", <default>)` rather than `hf_config.<field>` — older checkpoints drop fields. 3. **`@property hf_config(self) -> <HFConfigClass> | None`** — the inverse of `from_hf`: every field `from_hf` reads must be re-emitted so `save_hf` round-trips. **Built-in vs. `trust_remote_code` — the only axis that changes `from_hf` / `hf_config`:** a `trust_remote_code` model is structurally still Dense / MoE / VLM; the *only* difference is whether `transformers` ships a built-in `<Name>Config` class. Decide which case you are in and follow the matching column — this is the single place that contract is defined. | | Built-in config (e.g. Qwen3) | `trust_remote_code` (no built-in config) | |---|---|---| | `from_hf` reads via | `<HFConfigClass>.from_pretrained(hf_path)` | `AutoConfig.from_pretrained(hf_path, trust_remote_code=True)`, then `assert hf_config.model_type == "<expected>"` | | `hf_config` returns | a populated `<HFConfigClass>` re-emitting every field | `None` | | `save_hf` weight-map | written from `hf_config` | same | | `save_hf` side files | derived from `hf_config` | falls back to copying `config.json` / tokenizer / `*.py` from `self._hf_path` so the dir stays loadable with `trust_remote_code=True`; module-cache test in the matrix below | Compare `Qwen3MoEConfig.hf_config` (built-in) and `Qwen3MoEFoPEConfig.hf_config` (returns `None`), both in `xtuner/v1/model/moe/qwen3.py`. **Do not** invent a built-in HF config class for a remote-code model — return `None`. **Invariants:** - `from_hf` ↔ `hf_config` must round-trip. The `test_save_hf` test (§7) is the enforcement mechanism — every key in the original index must appear in the saved index with byte-equal tensors. - The config must be **pickleable and side-effect free**. No env vars, no file I/O, no `torch.cuda.*` at construction time. ### 3.2 Size-specific subclasses Add one subclass per released checkpoint that hard-codes the published dimensions. Mirror `Qwen3MoE30BA3Config` / `Qwen3MoE235BA22Config` in `xtuner/v1/model/moe/qwen3.py`. Even if only one size ships, define at least one — it is the entry point for `get_model_config` and for tests. --- ## 4. Register the entry point Open `xtuner/v1/model/__init__.py` — all four edits are short: 1. **Import** — bring in the new config classes. 2. **Alias** — add `"<alias>": <NewSizeConfig>()` to `model_mapping`. Lookup