add-hf-model

name: add_hf_model description: > Add support for a HuggingFace model in XTuner so it can be trained with the full set of parallelism strategies and optimization switches. The user can provide a HuggingFace Hub repo id, a local model directory, or a model family already supported by transformers. This skill walks through the full path: locating the reference implementation, classifying the model, splitting new code across the model / module / ops layers, implementing the XTuner model class and config, registering the entry point, validating bitwise numerical parity with HuggingFace, writing regression tests, and enabling the training optimizations (EP/SP, micro-batch, torch.compile, fp8, activation offload).

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:

   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.

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

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):

# 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.

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 normalizes case and -/_, so "my-model-7b" and "my_model_7b" are the same key.
3. Dispatch — add elif cfg.model_type == "<model_type>": to get_model_config_from_hf returning <NewConfigClass>.from_hf(model_path). The catch-all raise ValueError(...) must stay last.
4. __all__ — re-export the new public config names.

Never key dispatch on parameter names, file paths, or architecture strings. Only model_type.

5. Compose / VLM specifics

Compose models live under xtuner/v1/model/compose/<name>/ and follow a three-tower pattern. See xtuner/v1/model/compose/base.py and xtuner/v1/model/compose/qwen3_vl/qwen3_vl_config.py.

5.1 Config shape

BaseComposeConfig (compose/base.py) carries three sub-configs:

vision_config:    XTunerBaseModelConfig     # e.g. Qwen3VLVisionConfig
projector_config: XTunerBaseModelConfig     # e.g. Qwen3VLProjectorConfig
text_config:      XTunerBaseModelConfig     # any Dense/MoEConfig subclass
freeze_vision:    bool = False
freeze_projector: bool = False
freeze_language:  bool = False

Reuse existing text-tower configs where possible. Qwen3VLDense* delegates to the Dense path via xtuner/v1/model/dense/qwen3vl_text.py. Do not copy-paste the entire dense/MoE config — compose over them.

Reused text-tower key prefixes — use hf_key_mapping, not to_hf_key_list. A text tower reused inside a VLM keeps its standalone to_hf_key_list (it emits model.<...> as if it were a plain LLM). The compose checkpoint nests it one level deeper (model.language_model.<...>). Do not teach the text tower's to_hf_key_list about that prefix — that leaks the compose context into a tower that must also work standalone. Set the remap on the text config instead: hf_key_mapping = {r"^model\.": "model.language_model."} (applied in _init_load_spec). Keep structural renames in to_hf_key_list; the deployment-dependent prefix region belongs in hf_key_mapping.

5.2 Model shape

BaseComposeModel (compose/base.py) builds three sub-models: vision_tower, multi_modal_projector, language_model.

• from_hf(hf_path) — delegates to each sub-model with strict=False (vision keys are missing from the language-model state dict and vice versa) and unions the missing-keys sets.
• save_hf(hf_dir) — saves each tower with a distinct prefix ("model-language", "model-vision", "model-projector") and merges the three weight_maps into one model.safetensors.index.json.

5.3 hf_config for VLMs

Many VLMs ship without a stable HF top-level config class; hf_config returns None (compose/qwen3_vl/qwen3_vl_config.py). In that case, save_hf falls back to copying source files from self._hf_path (see §3.1).

5.4 Vision-side quirks to watch for

• attn_impl selectable per checkpoint: "flash_attention" | "flex_attention" | "eager_attention".
• deepstack_visual_indexes (e.g. [8, 16, 24]) — auxiliary supervision depths on the vision tower; must round-trip.
• Specialized rope: rope_type="qwen3_vl", mrope_section=[24, 20, 20] — delegated to RopeParametersConfig.from_hf_config.
• The projector’s torch.compile path is enabled on torch ≥ 2.9.1 (see commit f6d74efb). Tests must validate both compiled and eager outputs match.

6. Baseline parity — no complex parallelism yet

First make the model correct on the simplest execution path: single rank, no EP/SP, no torch.compile, no fp8, no activation offload. Align precision from the inside out — ops → decoder layer → full model — reusing XTuner's own modules at every level. Only once this baseline holds do you add the parallel and optimization features (§8).

1. Reuse, don't reimplement. Use the existing XTuner modules — MHA, MLA, GatedDeltaNet, GroupLinear, etc. Do not re-implement them with naive torch.

2. Route op differences through XTUNER_HF_IMPL. This env var already exists (xtuner/v1/utils/misc.py) and selects HF-exact ops via per-op selectors: xtuner/v1/ops/attn_imp.py::get_attn_impl_fn (eager attention), xtuner/v1/ops/rms_norm/__init__.py::get_rms_norm_fn (native torch rms_norm), and xtuner/v1/ops/gated_deltanet/__init__.py::{get_chunk_gated_delta_rule_fn, get_causal_conv1d_fn} (return the canonical fla / causal_conv1d wrappers HF calls instead of XTuner's compile-friendly torch.library.custom_op wraps — same forward output, but the backward op graph now matches HF). Selectors read the env var live so a test can toggle it per model instance. If your model introduces a new op whose fast path is not bitwise against HF (typical for fla-backed linear-attention variants), add its HF-exact branch in the corresponding op selector the same way. You may only patch at the ops level — never take the shortcut of patching the entire XTunerDecoderLayer.forward.

   Pair XTUNER_HF_IMPL with XTUNER_DETERMINISTIC=true. The latter gates an autotune-pin block in xtuner/v1/__init__.py that monkey-patches triton.autotune to lock every kernel to its first config and disable the result cache. Without the pin, fla's @triton.autotune picks tiling / reduction order per kernel and the choice drifts between runs, producing 1 ULP per linear-attention layer that the chain rule amplifies to ~1.0 at the model boundary (§12). tests/conftest.py sets the env var and runs import xtuner.v1 so the patch installs before fla is imported. The patch lives on the xtuner.v1.* import chain (not in the conftest) because multiprocessing.spawn children of MultiProcessTestCase don't go through pytest's conftest — they re-import the test class top-level (which pulls xtuner.v1.model.*) and the patch must run in those processes too. Production training doesn't set the env var, so the block is inert there.

3. Decoder-layer bitwise parity — forward AND backward, both required. Build one XTuner DecoderLayer (one per distinct layer type — e.g. a linear and a full layer for a hybrid model) and the matching HF layer, set XTUNER_HF_IMPL and XTUNER_DETERMINISTIC, feed identical inputs, and align bitwise — not a tolerance. Backward is not optional: a single layer → final norm → lm_head → CE loss → loss.backward() test asserts the layer output, the loss, and the input gradient dL/dx all bitwise-equal to HF — test one full and one linear layer separately (reference: test_decoder_layer_bitwise_parity). A bitwise forward does not imply a bitwise backward — the forward values can round identically while the backward graph differs (see the fp32-softmax note in §12). This is the primary path, required at all scales, and the only feasible one for large models (it never loads the full model — see "load only what you use" below). When you hand-build a single layer's inputs, the two sides' causal mask / seq_ctx must match exactly — see §12. Load only what you use so the test fits any model size: build the XTuner tower on meta and materialize (to_empty + selective load) only the tested layer + norm + lm_head via HFCheckpointLoader.load(key) (it reads just the needed safetensors shard); build a standalone HF DecoderLayer + norm + lm_head and load the same keys. Never from_pretrained the whole model for a single-layer test. (build_rotary_embedding builds the rotary on CPU with real buffers even under meta, so model.rotary_emb is usable without materializing the layers.) For < 40B you may also read per-layer hidden states from one full-model forward (HF output_hidden_states + XTuner return_hidden_states); mind the off-by-one in §12.

4. Model config matches HF hyperparameters — see the §3.1 built-in vs. remote-code contract for how from_hf / hf_config differ between the two.

5. Full-model forward parity for models < 40B, bitwise against HF, on the single-rank path (FSDP on or off; no EP / compile / fp8 yet).

6. save_hf / from_hf correctness (round-trip; see §7).

7. Loss-convergence trace + engine test. Once model forward and decoder-layer forward/backward are aligned, record a loss-convergence trajectory and add an engine test case.

Commit the baseline — do not stop to ask. Committing is part of this workflow: invoking this skill authorizes the in-workflow commits. Once parity holds on the simple path and lint/tests pass, commit immediately so this correct baseline is preserved before any parallel/optimization work begins (the §8 sub-agents branch off this commit). Commit each logical step as you finish it (§10) rather than batching everything to the end or waiting for a separate go-ahead.

Report back to the user: whether bitwise parity was achieved, and the residual error (decoder-layer level and model level) once XTuner's internal components are used.

7. Baseline tests — add these first

Before any parallel feature, add the most basic regression tests so the baseline (§6) is locked and the §8 work has a fixed reference to measure against. Add tests/model/test_<name>_<bucket>.py, mirroring tests/model/test_qwen3_moe.py (MoE) or tests/model/test_qwen3_dense.py (Dense). All tests extend xtuner._testing.DeterministicDDPTestCase for deterministic distributed setup.

Reference test cases:

• decoder-layer bitwise parity + save_hf round-trip: tests/model/test_qwen3_5_dense.py (hybrid linear/full dense VLM; toggles XTUNER_HF_IMPL per case)
• model forward & save & load parity: tests/model/test_qwen3_moe.py
• engine training test: tests/engine/test_moe_train_engine.py

7.1 Baseline test matrix

All single execution path — no EP / dispatcher / compile / fp8 here; those are §8 tests, written against this baseline.

7.2 Parity bar by model scale

Two independent axes:

Per-layer (op) parity — required at all scales, bitwise. Test one full and one linear decoder layer separately; for each assert layer output + loss + dL/dx bitwise vs HF under XTUNER_HF_IMPL. Keep it memory-light (XTuner meta-build + materialize only the tested layer + norm + lm_head; standalone HF layer; load via HFCheckpointLoader) so it runs even for models too large to forward end-to-end. This is the bitwise guarantee for both forward and backward ops.

Whole-model (integration) parity — on the top-level model, graded by what fits one GPU. Run the real model end-to-end: for a VLM that's the compose model on an image prompt (one test covers vision + projector + text — don't add a separate text-tower model test); for a plain LLM, the text/LM model.

• Small (e.g. dense 4B): both forward (logits/loss) AND backward (dL/d(pixel_values) for a VLM, dL/d(inputs_embeds) for text) bitwise. Prerequisites: XTUNER_HF_IMPL=true so every op selector returns the HF-canonical callable (§6.2) AND XTUNER_DETERMINISTIC=true so Triton autotune is pinned (§6.2). With both on, even hybrid linear/full models match HF byte-for-byte through the whole compose chain. Don't check weight gradients — they accumulate across positions + tied lm_head and aren't the right granularity for parity.
• Medium: forward only (logits/loss bitwise).
• Large (cannot forward end-to-end on the test GPUs): skip the whole-model test; leave it to integration tests. The per-layer test still gives op-level forward+backward bitwise.

7.3 Required test idioms

tests/model/test_qwen3_5_dense.py is the reference template; extend DeterministicDDPTestCase (xtuner._testing), which provides the shared parity scaffolding so the test file only holds model-specific logic:

• with self.hf_impl(): — sets/restores XTUNER_HF_IMPL=true; build the XTuner model inside the block so its attention picks the eager op at construction.
• self.materialize_xtuner_submodule(model, submodule, prefix, loader) — for a model built on meta, materialize only that submodule on CUDA and load its weights (so a single-layer test fits any model size).
• self.load_params_from_hf(module, key_for, loader) / self.xtuner_ckpt_key(model, name) — selective per-param load by checkpoint key (HFCheckpointLoader reads only the needed shard); xtuner_ckpt_key applies to_hf_key_list + the config's hf_key_mapping.

# Construct on meta, materialize on device — single-rank baseline, eager, no EP.
with torch.device("meta"):
    cfg = get_model_config_from_hf(hf_model_path)
    cfg.compile_cfg = False       # baseline tests run eager
    model = cfg.build()._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True)

# Optionally shard with FSDP (data-parallel only), then load.
model.fully_shard(fsdp_config=FSDPConfig(cpu_offload=False))
model.from_hf(hf_model_path)

Older patch_hf_rms_norm / patch_hf_rope helpers also live in xtuner._testing. Check they actually fit your model before using them — they were written for earlier models and match HF modules by class-name substring and attribute name (e.g. variance_epsilon). On a newer model they can silently mis-patch (wrong eps attribute, or replacing a zero_centered RMSNorm with a default one). Often the native XTuner ops already match HF bitwise (see §12), so no patching is needed; prefer XTUNER_HF_IMPL over these helpers for parity. The §8 parallel-feature tests reuse this same idiom with cfg.dispatcher / cfg.ep_size (and FSDPConfig(ep_size=...)) set.

7.4 Checkpoint paths

Read from env vars; never hard-code. Example: QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"] at tests/model/test_qwen3_moe.py. Document new env vars in the PR body.

7.5 Drop-in training config

Ship a runnable training config at ci/config/<name>.py so a developer (and CI) can verify the port end-to-end with one command: load the HF checkpoint, run a few real training steps, watch the loss curve descend as expected. This is the smoke test that catches anything the unit tests miss — the full from_hf → model.forward → loss → backward → optimizer step → FSDP shard/reduce chain — and the file doubles as the example users copy when wiring the 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 <NewSizeConfig>() (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.

8. Training optimizations

Only after the §6 baseline is committed and the §7 tests are green: verify each switch runs correctly and keeps precision within budget against that baseline.

1. EP and SP can be enabled normally.
2. With ep_size > 1, intra_layer_micro_batch enables correctly and stays precise.
3. torch.compile runs correctly.
4. fp8 can be enabled.
5. XTUNER_ACTIVATION_OFFLOAD can be enabled.

Each switch needs a test that compares its output to the §6 baseline — e.g. the dispatcher × ep_size parity matrix (at minimum {(None,1), ("all2all",4), ("all2all",8)}, add "deepep" if supported), compiled-vs-eager, or fp8-vs-bf16 within tolerance.

These items are well-suited to the multi-agent workflow in §11: the main agent analyzes and creates tasks (with test cases), sub-agents develop in parallel, and the main agent merges. Report the precision error introduced by each switch.

9. Reference cheat-sheet

10. Commit discipline

0. Commit as you go — this is authorized, do not pause for a separate go-ahead. The user invoking this skill authorizes the commits this workflow produces. Finish a logical step → lint/test → commit, then move on. Do not leave the whole port uncommitted and end with "should I commit?"; that strands the work. (Hard-to-reverse or outward actions like git push / opening a PR still need explicit confirmation — committing locally does not.)
1. Follow the stacked-PR convention. Plan upfront how many commits the model port needs — one per logical step (e.g. shared-path fixes; model class + config + registration; the XTUNER_HF_IMPL switch; baseline tests; then one per parallel/optimization feature). Order them by dependency (a shared-path fix the baseline relies on goes first). Land every later change as git commit --fixup=<sha> + git rebase --autosquash into the commit it belongs to. Never grow the history with endless patch commits just to keep fixing things up.
2. Every commit must pass lint:

   pre-commit run --files $(find xtuner/v1)
   

11. Multi-agent workflow

This applies only to the §8 parallel/optimization features — and only once the §6 baseline and §7 baseline tests are done and committed (§10). The baseline itself is built sequentially, not multi-agent; the §8 features are independent of each other, which is what makes them safe to develop in parallel.

1. The main agent splits the §8 features into tasks and, per §8, writes each feature's comparison test against the committed baseline and commits the tests.
2. Sub-agents branch off the baseline/test commit and develop one feature each (e.g. one switch / one tower) in parallel, making its test pass.
3. The two sides adjust each other — a shared-interface change is made once by the main agent (it updates the test/baseline; sub-agents rebase), and a failing test pushes the sub-agent to fix its feature, not to weaken the test.
4. The main agent merges, resolves conflicts, and reports the precision error each feature introduces relative to the baseline.

12. Parity debugging pitfalls

Hard-won notes from aligning hybrid (GatedDeltaNet + gated-MHA) models bitwise. The meta-rule: measure on real GPU and bisect to localize the divergence before concluding anything — most "bugs" here are measurement artifacts, not model errors. A clean way to bisect is to feed one side's intermediate tensors into both modules and compare step by step (norm → q/k/v → qk_norm → rope → attention core → o_proj).

• Single-layer parity: the two sides' inputs must match exactly — especially the causal mask. When you hand-build one layer's inputs, passing attention_mask=None to an HF attention/layer makes HF apply no causal mask while XTuner stays causal, producing a large false diff. Build a causal mask equivalent to XTuner's seq_ctx, or drive the layer through the model so the mask is constructed internally.

• flash ≠ eager, and HF eager upcasts softmax to fp32. XTuner defaults to flash attention; HF reference eager does softmax in fp32. They are numerically close but not bitwise equal, so bitwise parity requires forcing eager on the XTuner side (that is exactly what XTUNER_HF_IMPL does). XTuner's eager_attention matches HF's eager_attention_forward bitwise.

• A bitwise forward does not imply a bitwise backward. The eager softmax must run in fp32 (softmax(..., dtype=torch.float32).to(x.dtype), like HF), not bf16. With a bf16 softmax the forward logits can still round identically to HF (bitwise), but the backward differs (~1e-5 on dL/dx) because the gradient accumulates at bf16 precision. So always check backward bitwise separately (the loss.backward() test), not just forward — and align op dtypes (softmax, norms) in the backward-sensitive direction, fp32 where HF uses fp32.

• Whole-model backward is bitwise — but only with the full env stack. Per-layer bitwise backward does not on its own imply whole-model bitwise backward; the missing pieces are XTUNER_HF_IMPL=true AND XTUNER_DETERMINISTIC=true AND the conftest must import xtuner.v1 so the autotune-pin block in xtuner/v1/__init__.py runs before any test reaches fla. Without those, two distinct ULP-level sources compound through the LM chain rule and look indistinguishable from "fundamental bf16 noise":

  1. Triton autotune drift in fla. fla's @triton.autotune picks tiling / num_warps / reduction-order per kernel. Across runs the choice can change, producing a ~1-ULP backward diff per linear-attention layer that the chain rule amplifies — a 32-layer LM stack turned a per-layer ~2e-3 abs diff at grad magnitude ~30 into a ~1.0 diff on the massive-activation channels of dL/d(pixel_values). The pin (configs held to the first one, cache disabled) is in xtuner/v1/__init__.py under an XTUNER_DETERMINISTIC gate. It must run before any module imports fla — fla's @triton.autotune decorators evaluate at import time, so patching after the fact has no effect on kernels already constructed. That's why the patch is on the xtuner.v1.* import chain rather than in the conftest: multiprocessing.spawn children of MultiProcessTestCase re-import the test class top-level (which pulls xtuner.v1.model.*) without running conftest, so the conftest alone can't install the patch in those processes.
  2. torch.library.custom_op wraps for the linear-attention path (GatedDeltaNet, causal_conv1d) produce a different backward autograd graph from the canonical fla / causal_conv1d call HF uses, even though the underlying kernel and forward output are identical. XTUNER_HF_IMPL routes through get_chunk_gated_delta_rule_fn / get_causal_conv1d_fn (§6.2) so the parity path matches HF byte-for-byte.

  Production train doesn't set either env var; the custom_op wraps + autotune still operate. These switches are parity-test only. If the whole-model backward is not bitwise but the per-layer backward is, suspect one of the two before suspecting depth-accumulated bf16 noise.

• flash(XTuner varlen) vs flash_attention_2(HF) is bitwise only for short, single-tile sequences. They are the same kernel, so at a short seqlen (e.g.

  32. the logits are bitwise (0.0); at longer sequences the tiling / accumulation order differs and they drift ~1e-2. So a short-seq text test can assert flash-vs-FA2 bitwise, but anything realistic-length (or a VLM whose image tokens make the sequence long) needs the eager path for seqlen-independent bitwise. Don't conclude "flash matches FA2" from one short input.

• Shared kernels are already bitwise; focus on what isn't. GatedDeltaNet / linear attention go through the same fla / causal_conv1d kernels as HF, and native rms_norm / rope match HF bitwise — these come out at 0.0 with no work. Spend your effort on the path that actually differs (usually full attention).

• output_hidden_states off-by-one (full-forward route, < 40B only). HF's tuple is [emb, L0_out, …, L{n-2}_out, post_norm]: index i+1 is layer i's pre-norm output only for i ≤ n-2; the last entry is post-norm, not the last layer's pre-norm output. Compare the last layer via model.norm(...) or you will see the final layer "explode".

• A big max-diff is not automatically a bug. Late layers have massive activations (a few channels in the tens/hundreds); a tiny relative error there shows up as a large absolute max-diff. Look at the mean and the trend, not just the max. Conversely, feeding randn inputs to a single attention layer is ill-conditioned and misleads — use real (or model-produced) hidden states.

• A single-model failure is often a general-path bug. E.g. Dense.build_layers not forwarding rms_norm_type to DenseDecoderLayer broke any zero_centered dense model, not just this one. Fix the shared invariant, not the symptom.

• Loading a standalone sub-tower uses strict=False. A text tower loaded directly from a full VLM checkpoint will see vision (and deferred mtp.*) keys as unexpected; load with strict=False so they are ignored.

• Test the whole VLM, not just the text tower. Text-tower parity leaves the vision tower and projector completely untested. A compose port needs a full-model forward-parity test with real image inputs (§7.1). The text path of the compose forward still runs a dummy vision forward to keep the tower in the autograd graph, so vision bugs surface even on text prompts.

• Vision attention is non-causal — give eager_attention a causal=False path. XTuner's eager_attention defaulted to a causal block-diagonal mask; that is wrong for a vision tower (bidirectional within each image), which is why forcing vision to the causal eager op produced large diffs. The op now takes causal=True|False (False → non-causal block-diagonal mask), the vision module already calls it with causal=False, and the vision tower routes through get_attn_impl_fn so XTUNER_HF_IMPL selects eager. With that, the whole VLM is bitwise in one run, everything eager: XTUNER_HF_IMPL on the XTuner side (text causal-eager + vision non-causal-eager) and HF loaded with a single attn_implementation="eager". Prefer eager for parity — it is seqlen-independent, unlike flash/FA2 (see the flash seqlen bullet above).

• Vision pos-embed dtype when reusing the Qwen3-VL tower. HF's fast_pos_embed_interpolate returns fp32 (bilinear weights are fp32) and HF's vision forward casts it (pos_embeds.to(hidden_states.dtype)); XTuner's vision forward adds it without a cast. If you patch in the HF method for parity, cast the result back to the param dtype, or you get an fp32/bf16 LayerNorm mismatch (expected scalar type Float but found BFloat16).

• Hybrid-model torch.compile: choose fullgraph per layer — verify, then fall back; don't gate whether to compile. When a decoder layer's forward is wrapped into a higher-order op (HOP) — e.g. by an activation-checkpoint wrapper, which torch.compile lowers to a HOP whose body must be functional so it can be recomputed in backward — and that layer is a linear-attention layer (GatedDeltaNet), the layer mutates an outer-scope variable inside the HOP (GatedDeltaNet writes back into seq_ctx). Under fullgraph=True that raises HigherOrderOperator ... Mutating a variable not in the current scope (SideEffects). This is not a blanket "linear attention ⇒ fullgraph=False" rule — whether it triggers depends on whether the layer is actually HOP-wrapped and whether that linear impl really writes outer state. So for such a HOP + linear combination, first test whether fullgraph=True runs; keep it if it does (it compiles more of the layer), and only fall back to fullgraph=False if it doesn't — fullgraph=False survives by graph-breaking at the mutation and running that side-effecting write in eager. Make this a per-layer choice in Dense.fully_shard (decide each layer's fullgraph), not a decision about whether a layer is compiled at all.