Add New Model

Add a New Model to PithTrain

End-to-end workflow for integrating a new MoE language model. This file is the entry point: it tells you the phase order, the gates between phases, and which reference/*.md to load before each phase. Do not try to do everything from memory - load the reference file for the phase you're in.

Input

One of the following:

• HuggingFace model ID (e.g. "mistralai/Mixtral-8x7B-v0.1"). Used as the --hf-id for snapshot_download, and as the from_pretrained source for config and tokenizer.
• Local snapshot path (a directory containing config.json, model*.safetensors, tokenizer* etc.). Treat it exactly like the HF ID case - AutoConfig.from_pretrained(path) works on both. The snapshot itself came from HF, so online reference material (HF's modeling_<model>.py, TorchTitan, OpenAI release repo, upstream papers) is still fair game and should be consulted.

Optionally a model short name (filename stem, e.g. mixtral_8x7b). If the user doesn't give one, derive it from the HF ID by lowercasing and replacing / and - with _. Confirm with the user before using.

Hard rules (apply in every phase)

These are non-negotiable. Violating any of them will cost time later - in most cases that's exactly how past bugs landed.

1. Mirror HuggingFace, not our existing models. Class names, attribute names, and tensor structure (fused vs split) must match HF's modeling_<model>.py. Do not base the new model on Qwen3 / DeepSeek-V2 / GPT-OSS and rename - that path produced the GPT-OSS gate/router mismatch. See reference/conventions.md.
2. fullgraph=True for the compiled hot regions. forward_stage1_compute and forward_stage5 must each carry @torch.compile(fullgraph=True) (plus the attention forward itself when the model uses flex_attention, as GPT-OSS does). Never reach for fullgraph=False. See reference/compile.md.
3. Shared experts live in forward_stage1, not forward_stage5. If the model has shared experts (e.g. DeepSeek-V2), fold them into the residual at the end of forward_stage1_compute. See reference/protocol.md.
4. Check nvidia-smi before every GPU command. This is a shared cluster. Free GPUs can change between commands; don't reuse indices. nvidia-smi --query-gpu=index,memory.used,memory.free --format=csv and pick indices with memory.used < 1000 MiB. Do this once per invocation, not once at the start of the whole skill.
5. Test timeouts stay short (120-180s). Do not set a 10-minute timeout and walk away - if a test hasn't progressed in 3 minutes, it's hanging (likely torch.compile retrace). Kill and diagnose.
6. When tests fail, diagnose before relaxing anything. Print actual magnitudes. A 6-order-of-magnitude gradient discrepancy is not bf16 noise. Never loosen thresholds or add name-based skips as a first move.
7. Reject unused process-group arguments explicitly. If a model's __init__ accepts a group it does not actually implement (e.g. cp_group on a model with no ring-attention path), silently ignoring it will produce wrong results when a real group is eventually passed. Raise NotImplementedError when group is not None and group.size() > 1. See reference/protocol.md (init-requirements).
8. Thread config values through __init__; don't hardcode them. Any value that appears in HF's config.json is a per-checkpoint knob - read it from the config even if every released checkpoint currently ships the same value. Only true architectural constants (paper coefficients, spec-defined magic numbers) stay as module-level literals, and they get a one-line comment naming the source. See reference/conventions.md (thread-config).
9. Never stage .agents/, .claude/, AGENTS.md, CLAUDE.md, or docs/ in commits.

Phase overview

Phases 1-4 are modeling + training correctness. Phases 5-6 are real-weight inference (only needed if the user wants to generate from trained / released weights). Phase 5 is skippable when the user only cares about training from scratch.

Do not skip ahead. Each phase is a gate: if phase N fails, do not move to phase N+1. If you're tempted to, stop and read reference/pitfalls.md.

Phase 0 - Analyze the HF reference

Before writing any code, inventory what you need to match.

Load reference/conventions.md before starting this phase. It has the diagnostic commands (grep class names, grep attribute names, safetensors shape dump) under the Quick diagnostic commands section.

Work through these sources in order:

1. modeling_<model>.py - class names, attribute names, fused vs split projections, special-case features (shared experts, sinks, sliding window, YaRN RoPE, clamped SwiGLU, attention biases).
2. A safetensors shard - actual expert-weight shapes and dtypes, not comments.
3. TorchTitan / Megatron-LM reference for this model - the training-framework-consensus expert layout ([E, out, in] vs [E, in, out]).
4. configuration_<model>.py - every default in <Model>Config.__init__. When model-specific defaults disagree with a generic fallback path, match the model-specific default (see reference/conventions.md (example-config)).
5. HF's MLP / activation forward - read the math directly. Don't trust config.hidden_act to tell you the whole activation. See reference/conventions.md (activation-math).

Record in a scratch doc (not a committed file): class names, attribute names, expert tensor layout, fused/split projections, per-checkpoint knobs (thread through __init__) vs architectural constants (module literals with a source comment), process groups the model accepts but doesn't implement (reject via NotImplementedError; see reference/protocol.md (init-requirements)), and any special-case features that map to entries in reference/pitfalls.md.

Gate: you can articulate exactly which class names, attribute names, tensor shapes, and config knobs you will wire. If any item is a guess, go back and print() it from the actual data.

Phase 1 - Write pithtrain/models/<model>.py

Load reference/protocol.md and reference/compile.md before starting.

1. Start from templates/model_skeleton.py. It is a structural outline (NOT a copy of Qwen3). Fill in the TODO placeholders with the HF- derived names and shapes from phase 0.
2. Implement in this order:
   • RotaryEmbedding (mirror HF's)
   • Attention (mirror HF's kernel choice - flash_attn for standard MHA/GQA, flex_attention for sinks/sliding)
   • Experts module
   • Router / Gate
   • MLP (the MoE block that wires router + experts)
   • DecoderLayer (the 5-stage split)
   • Model (forward via model_forward, posemb / prolog / epilog, reference_forward)
3. Checklist for the decoder layer:
   • self.idx = layer_id and self.mlp assigned (satisfies LayerProtocol in pithtrain/models/interface.py)
   • EP sizing read from the distributed context (distributed.ep_size / distributed.ep_group), with the local expert count on self.mlp.experts_per_rank
   • @torch.compile(fullgraph=True) on forward_stage1_compute and forward_stage5
   • Shared experts (if any) fold into residual at the end of forward_stage1_compute, before the return
   • reference_forward runs eager (no compile) and is numerically equivalent to forward_stage1 -> forward_stage3 -> forward_stage5
   • forward_stage3 truncates expert input by sum(ks) if the expert block has biases or elementwise post-ops (prevents 0*NaN=NaN in backward)
   • forward_stage3 uses padded_index_gather (not raw indexing) for both expand and reverse shuffle
4. Checklist for the model class:
   • Layers built via layer_partition(config.num_hidden_layers, stage_count, stage_index) from pithtrain/pipeline/dualpipev.py
   • forward delegates to model_forward(self, hidden_states, self.chunk_record, cu_seqlens); the engine records each stage into self.chunk_record (a ChunkRecord from pithtrain/pipeline/execution.py) for the pipeline backward, which the engine drives via model_backward.
   • forward_posemb, forward_prolog, forward_epilog, and reference_forward implemented per ModelProtocol.

Gate: file imports cleanly (python -c "from pithtrain.models.<model> import <Model>").

Phase 2 - Wire into the training framework

No new reference file needed - the changes are small and mechanical.

1. pithtrain/modules/training.py:
   • Import the new <Model>Model class.
   • Add a branch in setup_model:

     elif module_config.model_type == "<model_type>":
         ModelClass = <Model>Model
         model_kwargs = {"cp_group": cp_group}  # or {} if no CP support

   • Add the new class to the apply_fsdp isinstance assertion tuple.
   • Add the HF ID to the TrainingCfg.model Literal[...] union (if the user wants the HF ID to be an accepted value; config-path usage doesn't require this).
2. tests/test_dualpipev.py:
   • Import the new model + router/gate class (+ Experts class if it stores raw nn.Parameter expert weights - see reference/pitfalls.md).
   • Add the new class to the apply_fsdp isinstance assertion tuple.
   • Add a branch in the config.model_type switch in main. Slice num_hidden_layers down to 8 (and any parallel arrays like layer_types) to keep the test fast.
   • Add a fill_weights branch if:
     • The expert module stores raw nn.Parameter (not GroupedLinear). Without this, expert weights default to zero and the MoE subtree silently produces all-zero outputs - see reference/pitfalls.md.
     • The router/gate has new Parameters beyond weight (e.g. a per-expert bias).
   • Verify shard_experts can detect the experts module. If using raw nn.Parameter, the fallback gate on gate_up_proj already handles it. If the Parameter name is different, extend the fallback
     • gate on the distinctive weight name, not on num_experts alone (the router has num_experts too and must not be sharded).
3. Add the model config and HF ID to the models list at the bottom of tests/test_dualpipev.py.
4. Write examples/pretrain_lm/<model>/config.json. Mirror upstream HF's config.json field-by-field - including every nested block (rope_scaling, quantization_config, etc.). See reference/conventions.md (example-config) for the diff command and the three layered defaults you need to reconcile.

Gate: python -c "import tests.test_dualpipev" imports cleanly AND the example-config diff is either empty or has a documented reason for each remaining difference.