Add Moe Model

Add HF MoE Model to Transformers Backend

Add support for a new HuggingFace MoE model in the torchtitan/experiments/transformers_modeling_backend/ experiment.

Scope of changes

ONLY modify code under torchtitan/experiments/transformers_modeling_backend/. If supporting the new model appears to require a change outside this folder (e.g. in torchtitan/models/, torchtitan/distributed/, torchtitan/components/, or any other core path), STOP and escalate to the user — describe the change you believe is needed and why, and wait for their decision instead of editing core yourself.

The one exception is the temporary, must-be-reverted diagnostic patches in Phase 3c (numerical-gap analysis): those edit core only to measure where titan diverges from HF, are reverted before the change is done, and are never committed. Any permanent change outside the experiment folder still requires escalation.

DO modify experiment code to support the new model. The experiment folder (torchtitan/experiments/transformers_modeling_backend/) is designed to be extended for new models. Expected changes include:

• state_dict_adapter.py — add regex patterns for new key naming conventions, handle transposed weight layouts, fix round-trip mappings
• moe_replacement.py — add detection logic for new scoring functions, router types, expert layouts, or shared expert patterns in _probe_hf_moe_block and its helper functions
• hf_sharding.py — give every parameter/buffer-bearing module the model introduces a sharding config. This is not limited to attention projections: set_hf_sharding_configs covers the root-level tok_embeddings, norm, lm_head, and rotary_emb (and their buffers), and _set_layer_sharding_configs covers per-layer norms, the attention block, and the dense MLP. A module with params/buffers but no config will mix a plain tensor with a DTensor and crash under TP (the _assert_all_states_sharded backstop now makes that a loud setup-time error for any param/buffer-bearing dense-path module — root modules and every decoder layer, excluding the Titan-swapped MoE). For a genuinely new module, shard it (colwise_config/rowwise_config) if it has a clear TP layout, otherwise default to _replicate_config(module) (replicate weights, no TP). If a new module type doesn't fit any existing helper, that's a signal to stop and think (or escalate) rather than leave it unconfigured.
• model.py — handle alternative submodule names: HF models name the same component differently (feed-forward mlp vs feed_forward, final norm norm vs final_layernorm). Add a new model's variant to the relevant lookup (_get_moe_attr_name, the HFTransformerModel accessor properties).
• scripts/numerical_equivalence.py — add the model's synthetic test config

When blockers are found in Phase 1, fix them in the experiment code before proceeding to Phase 2. Do NOT stop and ask — implement the fixes, then continue the workflow.

Input

The user provides a HuggingFace model ID (e.g. Qwen/Qwen3-30B-A3B). Optionally, a local path to downloaded weights for pretrained testing.

Before starting, ask the user (if not already specified):

1. Which branch to work on — should changes go on the current branch, or a new branch? If running in a worktree, clarify whether the commit should target the main feature branch or stay on the worktree branch.
2. Whether to commit — after all tests pass, should the changes be committed automatically, or left as uncommitted changes for the user to review first?

Phase 1: Compatibility Probe

Load the model config and a single MoE layer on meta device to verify titan compatibility. No file changes in this phase — but record all issues found. After the probe, fix any blockers in the experiment code before moving to Phase 2.

1a. Load and inspect config

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)

Record: model_type, hidden_size, intermediate_size, moe_intermediate_size, num_attention_heads, num_key_value_heads, num_hidden_layers.

1b. Instantiate on meta device

with torch.device("meta"):
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)

1c. Find MoE layers

Scan model.model.layers for layers where hasattr(mlp, "gate") or hasattr(mlp, "router") AND hasattr(mlp, "experts").

Record:

• Which layer indices are MoE vs dense
• The mechanism: first_k_dense_replace, decoder_sparse_step, or mlp_layer_types
• Minimum num_hidden_layers needed to have at least one MoE layer

1d. Probe MoE block

Call _probe_hf_moe_block(moe_block, config) from torchtitan/experiments/transformers_modeling_backend/moe_replacement.py.

Check each field succeeds without error:

Issues to detect and fix in experiment code:

• Expert weight layout transposed (e.g. (E, H, 2*I) instead of (E, 2*I, H)) — fix state_dict_adapter.py to transpose during conversion, and moe_replacement.py to detect the layout
• Experts stored as nn.ModuleList of individual MLPs instead of fused 3D tensors — this is a true blocker (needs HF upstream change), report to user and stop
• Scoring function not softmax/sigmoid but detectable from the model code (e.g. torch.sigmoid call without e_score_correction_bias) — fix _resolve_score_func() in moe_replacement.py
• Router uses nn.Linear directly (key is router.weight not gate.weight) — add pattern to state_dict_adapter.py
• MoE layer attribute is not mlp (e.g. feed_forward) — fix model.py and moe_replacement.py to handle the alternative name
• Shared expert round-trip fails (singular/plural mismatch) — fix reverse patterns in state_dict_adapter.py
• Model doesn't support grouped_mm expert implementation — adjust the numerical test to skip the grouped_mm alignment for this model
• Layer-level MoE: router and experts are siblings of the dense MLP at the decoder layer level, not nested inside layer.mlp (e.g. Gemma4 has layer.mlp + layer.router + layer.experts). In this case the dense MLP is functionally a shared expert — it runs on every token and its output is summed with the routed expert output. The correct mapping is to treat the dense MLP as shared_experts inside Titan's MoE module and replace layer.mlp with the full Titan MoE (router + experts + shared_experts). This way the entire block has a single ShardingConfig and TP/EP work correctly. Fix in moe_replacement.py: when probing, detect that the dense MLP is a shared expert, include it in shared_expert_info, and when swapping, replace layer.mlp with the Titan MoE that includes the dense MLP as its shared expert. Wire the router and expert weights from the layer-level attributes into the Titan MoE.

1e. Probe attention for ShardingConfig

Check which sharding configs _set_layer_sharding_configs (in hf_sharding.py) needs:

Q and KV are detected independently, so any mix of full-rank/low-rank works (e.g. DeepSeek-V2-Lite has full-rank q_proj + low-rank KV):

The _assert_all_states_sharded backstop will fail loud at setup if any param/buffer-bearing attention (or other dense-path) module is left without a config, so an unhandled projection surfaces immediately. If found, report: "attention has unsupported projection <name>, needs _set_layer_sharding_configs update in hf_sharding.py."

Also check: does the layer have post_attention_layernorm? (Some models omit it — already handled.) Does the embedding have extra buffers (e.g. Gemma4's embed_scale)? (Already handled — buffers are enumerated dynamically.)

1f. State dict adapter round-trip

from torchtitan.experiments.transformers_modeling_backend.state_dict_adapter import (
    hf_to_titan_moe_state_dict, titan_to_hf_moe_state_dict,
)

hf_sd = model.state_dict()
titan_sd = hf_to_titan_moe_state_dict(hf_sd)
roundtrip_sd = titan_to_hf_moe_state_dict(titan_sd)

Verify: all original keys present, no extra keys, values match. If keys are missing, the adapter needs new patterns for this model's key naming convention.

1g. Report

Print a compatibility summary table:

Compatibility report for <model_id>
  model_type:          <type>
  hidden_size:         <H>
  moe_intermediate:    <I>
  num_experts:         <E>
  top_k:               <K>
  score_func:          softmax|sigmoid
  route_norm:          True|False
  shared_experts:      none|additive|sigmoid-gated
  attention:           gqa|mla|mla+dsa
  MoE layers:          <indices>
  min_layers_needed:   <N>
  state_dict_roundtrip: PASS|FAIL
  issues:              none | <list>

1h. Fix issues

If issues were found, fix them in the experiment code now. Consult the "Issue Resolution Reference" at the bottom of this document for where each fix goes. The only true blocker is nn.ModuleList experts — for everything else, implement the fix and re-run the probe to verify.

After all issues are resolved and the probe passes cleanly, proceed to Phase 2.

Phase 2: Add Synthetic Test Config

Add an entry to _MODEL_CONFIGS in torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py.

Config rules

Use production-scale hidden_size and moe_intermediate_size from the real config. Override everything else for a small single-GPU test:

"<model_type>": dict(
    model_type="<model_type>",
    hidden_size=<real>,           # from pretrained config
    intermediate_size=<real>,     # from pretrained config
    moe_intermediate_size=<real>, # from pretrained config
    num_local_experts=8,          # reduced from real
    num_experts_per_tok=2,        # reduced from real
    num_hidden_layers=<min>,      # minimum for MoE layer to exist
    num_attention_heads=<fits>,   # must divide hidden_size evenly
    num_key_value_heads=<fits>,   # must divide num_attention_heads evenly
    vocab_size=256,
    max_position_embeddings=64,
    attn_implementation="sdpa",
    use_cache=False,
    # Include model-specific fields:
    # first_k_dense_replace, decoder_sparse_step, mlp_layer_types,
    # n_group, topk_group, n_shared_experts, norm_topk_prob,
    # routed_scaling_factor, q_lora_rank, kv_lora_rank, etc.
)

Attention dimension rules:

• num_attention_heads must divide hidden_size evenly
• num_key_value_heads must divide num_attention_heads evenly
• For MLA models: use small q_lora_rank (16), kv_lora_rank (16), qk_rope_head_dim (8), qk_nope_head_dim (8), v_head_dim (16)
• For DSA models: use small index_n_heads (2), index_head_dim (8), index_topk (8)

MoE layer selection rules:

• If first_k_dense_replace > 0: set it to 0 (force first layer to be MoE)
• If model uses decoder_sparse_step: set to 1
• If model uses mlp_layer_types: set to ["sparse"]

Expert group routing:

• If n_group is set, num_local_experts must be divisible by n_group and num_local_experts / n_group >= 2

Phase 3: Run Numerical Equivalence

HF is the gold standard. The numerical equivalence test compares Titan's MoE output against HF's output. The test must run HF's forward pass unmodified — do NOT patch, monkey-patch, or alter HF's computation to make the test pass. The only permitted modifications to close accuracy gaps are temporary changes to titan core code, which must be reverted after verification.

Specifically:

• Do NOT patch the HF model's activation function to match titan's
• Do NOT replace HF's routing with titan's routing
• Do NOT modify the test's comparison logic or thresholds
• Do NOT skip or disable parts of the comparison
• Do NOT cast HF tensors to different dtypes to hide precision diffs (the existing float32 gate cast is an exception — it aligns titan's autocast behavior, not the other way around)
• If titan produces different results from HF, the gap analysis (Phase 3c) must identify WHY by temporarily changing titan to match HF, not by changing HF to match titan

Do NOT modify the pass/warn/fail thresholds. The thresholds (KL < 1e-6 PASS, KL < 1e-3 WARN, KL >= 1e-3 FAIL) are calibrated for the known dispatcher precision gap. If a new model hits WARN, verify via Phase 3c that the diff is from known causes. If it hits FAIL, fix the integration, do not loosen the thresholds.

3a. Synthetic test

CUDA_VISIBLE_DEVICES=<gpu> python \
    torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py \
    --models <model_type>

Expected: PASS with KL < 1e-6 and round-trip max_diff = 0.00.

If it fails, diagnose:

• ValueError from _probe_hf_moe_block — missing config attribute, fix the synthetic config
• Shape mismatch in load_state_dict — state dict adapter needs update
• RuntimeError from grouped_mm — check expert weight shapes
• Numerical FAIL — investigate routing or weight transfer issue

3b. Pretrained test (if weights available)

CUDA_VISIBLE_DEVICES=<gpu> python \
    torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence_pretrained.py \
    --model_dir <path>

Expected: PASS with KL < 1e-3 (known dispatcher precision gap).

3c. Verify numerical gaps via temporary core alignment (MANDATORY)

This step is NOT optional. Execute it fully and automatically whenever any test shows WARN or higher. Do NOT ask the user before proceeding — apply patches, run tests, record results, report, revert. The entire sequence runs without user interaction.

The goal: bring max_diff to 0.00 by temporarily making titan's implementation match HF's exactly. Every remaining diff must be experimentally explained — never claim a root cause without running code that proves it.

Rules:

• Never give up on a patch because it crashes. If a patch causes an error (e.g. dtype mismatch in scatter), fix the error by extending the patch. The patch is temporary — make it work, don't declare it "not applicable."
• Never claim a root cause without verification. If you think "the diff is from bmm vs grouped_mm", prove it: make both sides use the same kernel and show the diff drops. If you can't change one side (e.g. HF rejects grouped_mm), change the other side (e.g. make titan use bmm temporarily, or manually compute the expert forward with the same kernel as HF using titan's weights).
• Adapt patches to the model's code path. The known patches target specific code paths (e.g. f32 accumulation in LocalTokenDispatcher.combine()). If the model uses a different path, write an equivalent patch for that path. The principle is the same -- align titan's operations with HF's.
• Keep going until max_diff=0.00. If the known patches don't get there, trace the forward step by step to find additional differences.

Procedure:

1. Record baseline — note current max_diff and KL before any patches.

2. Apply patches incrementally — apply each applicable patch one at a time. After each, re-run the test and record the new max_diff. If a patch crashes, fix the crash (extend the patch as needed) and re-run. If a patch has no effect (e.g. top_k=1 makes reshape+sum trivial), record "no effect (top_k=1)" and move on.

3. Handle kernel mismatches — if the model doesn't support grouped_mm on the HF side, the test can't force the same kernel. In that case, verify experimentally:

   • Extract titan's expert weights after weight transfer
   • Run the expert forward manually using HF's kernel (e.g. torch.bmm) with titan's weights and the same input
   • Compare that output against HF's expert output
   • If they match (0.00), the remaining diff is confirmed as kernel difference, not a bug
   • Record this in the report as a verified kernel difference

4. Investigate remaining diffs — if max_diff is still non-zero after all patches and kernel alignment, trace the forward path step by step: router logits → routing weights → per-expert gate/up/down → score application → combine. Compare intermediate values between HF and titan to find where they first diverge. Make additional minimal patches to test hypotheses — if a patch doesn't help, revert it and keep digging.

5. Report — present a table to the user showing each patch and its measured impact:

Numerical gap analysis for <model_name>
| Patch | max_diff | KL | Description |
|-------|----------|----|-------------|
| baseline | X.XXe-XX | X.XXe-XX | no patches |
| +patch 1 | X.XXe-XX | X.XXe-XX | f32 accumulation in combine |
| +patch 2 | X.XXe-XX | X.XXe-XX | reshape+sum (no effect, top_k=1) |
| +kernel align | X.XXe-XX | X.XXe-XX | titan experts via bmm |
| final | 0.00e+00 | X.XXe-XX | all diffs explained |

6. Revert ALL patches — restore all core files:

   git checkout torchtitan/models/common/token_dispatcher.py torchtitan/models/common/moe.py
   

   Verify with git diff that the revert is clean. The temporary patches are NOT part of the commit.

Known sources of numerical difference:

Do NOT copy-paste patches from this document — the titan core code may have changed since these were written. Instead, read the current code in torchtitan/models/common/token_dispatcher.py and torchtitan/models/common/moe.py, understand the principle behind each fix, and write a patch that applies cleanly to the current code.

The three known differences and how to fix them:

1. bf16 accumulation before scatter_add — LocalTokenDispatcher.combine() multiplies expert output by routing scores in f32 but casts back to bf16 before deterministic_scatter_add, so the accumulation happens in bf16. Fix principle: remove the cast to bf16, keep everything in f32 through the scatter_add, allocate an f32 output buffer, cast to the input dtype only at the very end.

Example (illustrative, may not match current code):

# principle: remove .to(routed_output.dtype), use f32 buffer
routed_output = routed_output.to(torch.float32) * scores
out = scatter_add(f32_buffer, ..., routed_output)
return out.to(x.dtype)

2. scatter_add vs reshape+sum accumulation order — scatter_add accumulates top_k contributions in expert-sorted order. HF uses reshape(N, K, D).sum(dim=1) which accumulates in token order (consecutive top_k entries per token). Different f32 summation order produces different results. Fix principle: store the argsort permutation from dispatch, in combine unsort back to original token order, then reshape and sum over the top_k dimension. For top_k=1 this has no effect (only one element to sum).

3. topk sorted flag — titan uses torch.topk(sorted=False), HF defaults to sorted=True. With sorted=True topk values come back in descending order; the normalization sum is computed over differently ordered values, producing a small f32 rounding difference. Fix principle: change sorted=False to sorted=True in the router's topk call. For top_k=1 this has no effect (only one element).

With all three fixed, all existing models produce max_diff=0.00. If the new model still shows non-zero diff after all applicable fixes, the remaining difference is model-specific. Investigate as described in step 4 above.

Phase 4: Parallelism Integration Tests

Run the new model through the same parallelism configurations used for existing models. The integration test infrastructure is in torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py.

4a. Determine available GPUs

nvidia-smi -L | wc -l

4b. Run parallelism configs

Test the new model with at least these configurations (adapt GPU counts to what is available). Use --training.steps 2 to keep runs short — the goal is to verify the model initializes, parallelizes, and runs forward/backward without errors, not to converge.

# FSDP only (minimum viable test)
torchrun --nproc-per-node <NGPU> --module transformers_modeling_backend \
    --config transformers_modeling_backend_debugmodel_moe \
    --hf_model <model_id> \
    --training.steps 2

# FSDP + EP
torchrun --nproc-per-node <NGPU> --module transformers_modeling_backend \
    --config transformers_modeling_backend_debugmodel_moe \
    --hf_model <model_id> \
    --parallelism.data_parallel_shard_degree -1 \
    --parallelism.expert_parallel_degree <EP> \
    --training.steps 2

# FSDP + TP + EP (if NGPU >= 4)
torchrun --nproc-per-node <NGPU> --module transformers_modeling_backend \
    --config transformers_modeling_backend_debugmodel_moe \
    --hf_model <model_id> \
    --parallelism.data_parallel_shard_degree -1 \
    --parallelism.tensor_parallel_degree 2 \
    --parallelism.expert_parallel_degree <EP> \
    --training.steps 2

Every config must be run. There are only three — run all of them. Do NOT skip, do NOT mark as "NOT TESTED", do NOT mark as "SKIPPED". Every config in the final report must show either PASS or FAIL with a diagnosed error message and root cause.

If a config fails:

1. Read the error traceback
2. Diagnose the root cause
3. If fixable in experiment code, fix it and re-run
4. If it's a fundamental limitation (e.g. DTensor mixing), report FAIL with the specific error and why it can't be fixed in experiment code

Redundant compute is a blocker. If the only way to make TP work for a module is to run the full computation redundantly on every TP rank (e.g. all-gather input, run full forward, slice output), that is NOT an acceptable fix — it defeats the purpose of TP. If a model component (e.g. a custom attention mechanism) cannot be properly sharded, report it as a FAIL with the reason "no proper TP sharding available, would require redundant compute". Do NOT implement redundant-compute workarounds with pre/post hooks that all-gather and re-shard.

Use ShardingConfig, not hooks. The HF backend sets _sharding_config on every sub-module so model.parallelize() handles all DTensor distribution and forward wrapping. Available helpers from torchtitan/models/common/decoder_sharding.py and torchtitan/experiments/transformers_modeling_backend/hf_sharding.py:

• colwise_config() — for linear projections (weight Shard(0), output Shard(-1))
• rowwise_config(output_sp=True) — for output projections (weight Shard(1), output Shard(1) for SP)
• _hf_norm_config(enable_sp=True) — for norms (weight Replicate, activations Shard(1) when SP enabled)
• _replicate_config(module) — for modules that should have Replicate params/buffers without TP sharding (dynamically enumerates params and buffers to avoid _shard_states errors)
• ShardingConfig(in_src_shardings=..., in_dst_shardings=...) — for input redistribution at module boundaries (e.g. attention, MLP gather from Shard(1) to Replicate)

Do NOT use register_forward_pre_hook / register_forward_hook to manually convert DTensor ↔ local tensors. Do NOT use parallelize_module from torch.distributed.tensor.parallel. Use _sharding_config exclusively.

The --hf_model <model_id> flag overrides the default model in the base config. If the base config has incompatible defaults for the new model (e.g. wrong tokenizer path), override those too with additional -- flags. Do NOT give up because "the config is incompatible" — fix the incompatibility with overrides.

Common failure modes and fixes:

• TP failures from unsupported attention projections — fix in hf_sharding.py by adding ShardingConfig in _set_layer_sharding_configs
• EP failures from expert weight sharding — check FSDP shard_placement_fn
• Shape mismatches from model-specific layer structure — fix in experiment code
• Config incompatibility — add more --override.flags to the command
• Tokenizer mismatch — override --hf_assets_path to point to the model
• Custom attention with no standard Q/K/V/O projections — FAIL, not fixable without redundant compute
• Mixed DTensor / plain Tensor — a module has buffers not declared in its ShardingConfig's state_shardings. Fix by using _replicate_config(module) which dynamically enumerates all params and buffers

Note: If no GPUs are available, skip this phase and note it in the report. But if GPUs are available, all configs must be run and reported.

Phase 5: Update Model Compatibility Doc

Update torchtitan/experiments/transformers_modeling_backend/MODEL_COMPATIBILITY.md with the new model:

1. Add a row to the Numerical Equivalence Summary table with the model's status, KL, max_diff, cos_sim, and round-trip result.

2. Add a section under Architectural Differences per Model describing:

   • Router type and features (softmax/sigmoid, normalization, bias, scaling)
   • Expert format (standard/transposed, activation function)
   • Shared experts (none/additive/sigmoid-gated/layer-level)
   • Attention type (GQA/MLA/DSA/custom)
   • Any unsupported features and their impact
   • Summary of differences from Titan's MoE

3. Add a row to the Parallelism Support table (FSDP and FSDP+EP).

4. If the model has unsupported features requiring core changes, add or update entries under Core Changes Needed for Full Support.

Phase 6: Lint and Commit

1. Run pre-commit on all changed files:

   SKIP=pyrefly-check pre-commit run --files <changed_files>
   

   Fix any failures and re-run until clean.

2. Verify no core files were modified:

   git diff --stat  # only files in experiments/transformers_modeling_backend/
   

3. Create commit:

   Add <ModelName> to transformers MoE backend
   
   Adds synthetic test config with production-scale dimensions
   (H=<hidden_size>, I=<moe_intermediate_size>, E=8). Verified via
   numerical equivalence test (KL=<value>, max_diff=<value>).
   

Issue Resolution Reference

Most issues are fixable in the experiment folder. Only nn.ModuleList experts is a true blocker requiring upstream changes.

Files touched (experiment folder only)