Add Archon Model

Add support for a new HuggingFace model architecture in the Archon training engine.

When to Use

This skill is triggered when:

• User asks "how do I add a model to Archon?"
• User wants to support a new model family (e.g., Llama, Mistral, DeepSeek) in

ArchonEngine

• User mentions adding a new ModelSpec or model type for Archon

Prerequisites

Before starting, ensure:

• The target model is available on HuggingFace (has config.json with model_type)
• You know the HuggingFace model ID (e.g., meta-llama/Llama-3-8B)
• The model uses a standard transformer architecture (decoder-only)

Step-by-Step Guide

Step 1: Analyze the Target Model Architecture

Read the HuggingFace model's source code to extract key architecture information.

Action: Fetch and analyze the model's HuggingFace configuration and modeling files.

1. Read the model's config.json (via AutoConfig.from_pretrained) to identify:

- modeltype string (this is the key used for registry lookup) - All architecture hyperparameters (hiddensize, numlayers, etc.) - Any model-specific fields (e.g., qknorm, attention_bias, MoE fields)

1. Read the HuggingFace modeling_*.py source to identify:

- Attention variant: Does it have Q/K norm? Attention bias? Sliding window? Multi-latent attention? - FFN variant: SwiGLU (gateproj + upproj + downproj)? GeGLU? Standard MLP? - MoE support: Does it have MoE layers? What router type? Shared experts? - RoPE variant: Standard RoPE? YaRN? NTK-aware scaling? What is the invfreq formula? - Normalization: RMSNorm or LayerNorm? Pre-norm or post-norm? Elementwise affine? - Weight tying: Does tiewordembeddings appear in config? - State dict key names: What are the HF weight key naming conventions?

1. Summarize findings in a checklist like:

Target model: <name>
HF model_type: "<model_type>" (and variants like "<model_type>_moe" if applicable)
Attention: [standard GQA / with QK norm / with bias / sliding window / ...]
FFN: [SwiGLU / GeGLU / standard MLP / ...]
MoE: [no / yes - num_experts, top_k, shared_experts]
RoPE: [standard / YaRN / NTK-aware / ...]
Norm: [RMSNorm / LayerNorm] with [pre-norm / post-norm]
Weight tying: [yes / no]

Step 2: Select the Reference Model

Choose the closest existing implementation as a starting point:

Action: Copy the reference model directory as the starting point:

areal/experimental/models/archon/<model>/
  __init__.py
  spec.py
  model/
    args.py
    model.py
    rope.py
    state_dict_adapter.py
  infra/
    parallelize.py

Step 3: Implement args.py

Adapt <Model>ModelArgs to match the target model's HuggingFace config fields.

Key changes from reference:

1. Update the @dataclass fields to match the target model's hyperparameters:

- Field names should use Archon conventions (dim, nlayers, nheads, nkvheads, vocabsize, headdim, hiddendim, normeps, ropetheta, etc.) - Default values should match the smallest variant of the target model - Add model-specific fields (e.g., attentionbias, qknorm, slidingwindow)

1. Update fromhfconfig() to correctly map HuggingFace config attributes:

- Use getattr(hfconfig, "fieldname", default) for optional fields - Handle variant-specific fields (e.g., MoE fields only present in MoE variants) - The method must return an instance of the model args class

Critical: Verify every field mapping against the HF model's config.json. Incorrect mappings here cause silent errors downstream.

Base class contract (BaseModelArgs):

@dataclass
class <Model>ModelArgs(BaseModelArgs):
    # ... model-specific fields ...

    @classmethod
    def from_hf_config(
        cls,
        hf_config: PretrainedConfig,
        is_critic: bool = False,
        **kwargs,
    ) -> <Model>ModelArgs:
        # Map HF config fields to Archon model args
        ...

Step 4: Implement model.py

Adapt the model architecture to match the target model.

Key components to adapt:

1. Normalization (RMSNorm or similar):

- Check if elementwise_affine is configurable - Check the epsilon default value - If the model uses LayerNorm, implement accordingly

1. Attention module:

- Q/K/V projection: Check bias presence (nn.Linear(..., bias=True/False)) - QK norm: Add qnorm/knorm if the model has them, remove if it doesn't - GQA: nkvheads \< nheads for grouped-query attention - Ulysses SP: Keep the setcpgroup / sp_enabled pattern from the reference - Output projection: Check bias presence

1. FeedForward module:

- SwiGLU: w2(silu(w1(x)) * w3(x)) -- most common for modern LLMs - Check bias in linear layers - For MoE models: MoE module replaces FeedForward on designated layers

1. TransformerBlock: Pre-norm (most modern LLMs) vs post-norm

- MoE layer detection via ismoe_layer() if applicable

1. Top-level Model (<Model>Model(BaseArchonModel)):

- tokembeddings, layers (as ModuleDict), norm, output/score - initweights(): Match initialization scheme from HF - initbuffers(): RoPE cache + MoE buffers - forward(): Must follow BaseArchonModel signature: (tokens, positions, cuseqlens, maxseqlen, treeattn_meta=None) -> Tensor

Base class contract (BaseArchonModel):

class <Model>Model(BaseArchonModel):
    def forward(self, tokens, positions, cu_seqlens, max_seqlen, tree_attn_meta=None) -> torch.Tensor: ...
    def init_weights(self) -> None: ...
    def init_buffers(self, buffer_device) -> None: ...

Step 5: Implement rope.py

Handle the rotary position embedding variant.

Options:

1. Standard RoPE (same as qwen2/qwen3): Re-export from qwen2:

``python from areal.experimental.models.archon.qwen2.model.rope import ( applyrotaryemb, precomputeropecache, repeatkv, reshapeforbroadcast, rotatehalf, ) ``

1. Custom RoPE (YaRN, NTK-aware, etc.): Implement custom precomputeropecache()

and applyrotaryemb() functions. The key difference is usually in how inv_freq is computed (scaling factors, interpolation, etc.).

Step 6: Implement statedictadapter.py

Map between HuggingFace and Archon weight key names.

This is the most error-prone step. The adapter must correctly handle:

1. Key name mapping (fromhfmap dict):

- Embedding: model.embedtokens.weight -> tokembeddings.weight - Attention: model.layers.{}.selfattn.qproj.weight -> layers.{}.attention.wq.weight - FFN: model.layers.{}.mlp.gateproj.weight -> layers.{}.feedforward.w1.weight - Norms: model.layers.{}.inputlayernorm.weight -> layers.{}.attentionnorm.weight - Output: lmhead.weight -> output.weight - Skip keys (set to None): rotaryemb.inv_freq (computed at runtime) - Model-specific keys: bias terms, QK norm weights, etc.

1. Reverse mapping (tohfmap): Auto-generated from fromhfmap

1. MoE expert weights (if applicable): 3D\<->2D conversion for expert weights. Copy

the MoE handling from qwen3 if the model has MoE.

1. Weight tying: Skip output.weight during tohf() if tieword_embeddings=True

Verification approach: After implementation, the adapter should satisfy:

# Roundtrip: archon -> hf -> archon preserves all keys
hf_sd = adapter.to_hf(archon_sd)
roundtrip_sd = adapter.from_hf(hf_sd)
assert set(roundtrip_sd.keys()) == set(archon_sd.keys())

Base class contract (BaseStateDictAdapter):

class <Model>StateDictAdapter(BaseStateDictAdapter):
    def from_hf(self, hf_state_dict) -> dict[str, Any]: ...
    def to_hf(self, archon_state_dict) -> dict[str, Any]: ...
    def convert_single_to_hf(self, name, tensor) -> list[tuple[str, torch.Tensor]]: ...

Step 7: Implement parallelize.py

Define the parallelization strategy for the model.

The parallelize function applies parallelism in this order:

1. TP (Tensor Parallelism) -- shard attention/FFN across devices
2. EP (Expert Parallelism) -- for MoE models only
3. CP (Context Parallelism / Ulysses SP) -- sequence parallelism
4. AC (Activation Checkpointing) -- memory optimization
5. torch.compile -- compilation optimization
6. FSDP (Fully Sharded Data Parallelism) -- data parallelism

Key adaptations by model architecture:

• Attention with QK norm: wq/wk use uselocaloutput=False (DTensor output for

norm), add SequenceParallel(sequencedim=2) for qnorm/k_norm

• Attention without QK norm: wq/wk/wv all use uselocaloutput=True
• Attention with bias: Bias terms follow the same parallel plan as their weights
• MoE layers: Separate TP plan for MoE input/output, router gate, and expert

weights. Copy from qwen3's applymoeeptp() and applynonmoetp()

• Dense-only models: Simpler plan without MoE handling. Copy from qwen2

Function signature (must match ParallelizeFn protocol):

def parallelize_<model>(
    model: nn.Module,
    parallel_dims: ArchonParallelDims,
    param_dtype: torch.dtype = torch.bfloat16,
    reduce_dtype: torch.dtype = torch.float32,
    loss_parallel: bool = True,
    cpu_offload: bool = False,
    reshard_after_forward_policy: str = "default",
    ac_config: ActivationCheckpointConfig | None = None,
    enable_compile: bool = True,
) -> nn.Module:

Step 8: Create spec.py and Register

Assemble the ModelSpec and register it.

from areal.experimental.models.archon.model_spec import ModelSpec, register_model_spec
from areal.experimental.models.archon.pipeline_parallel import pipeline_llm
from areal.experimental.models.archon.<model>.infra.parallelize import parallelize_<model>
from areal.experimental.models.archon.<model>.model.args import <Model>ModelArgs
from areal.experimental.models.archon.<model>.model.model import <Model>Model
from areal.experimental.models.archon.<model>.model.state_dict_adapter import (
    <Model>StateDictAdapter,
)

<MODEL>_SPEC = ModelSpec(
    name="<Model>",
    model_class=<Model>Model,
    model_args_class=<Model>ModelArgs,
    state_dict_adapter_class=<Model>StateDictAdapter,
    parallelize_fn=parallelize_<model>,
    supported_model_types=frozenset({"<model_type>"}),  # From HF config.json
    pipelining_fn=pipeline_llm,
)

# Auto-register when module is imported
register_model_spec(<MODEL>_SPEC)

__all__ = ["<MODEL>_SPEC"]

Note: supportedmodeltypes should include all HF modeltype strings that this implementation handles (e.g., {"qwen3", "qwen3moe"} for Qwen3).

Step 9: Register in init.py

Add the import to areal/experimental/models/archon/init.py:

from areal.experimental.models.archon.<model> import spec as <model>_spec  # noqa: F401

This triggers auto-registration when the module is imported.

Step 10: Verify and Test

Verification should be done in stages, adapting based on available hardware and the test patterns in tests/experimental/archon/.

Before writing tests, examine the existing test files to understand current patterns:

tests/experimental/archon/
  conftest.py             -- Pytest configuration (version checks)
  utils.py                -- Shared utilities (model loading, comparison)
  test_qwen3_args.py      -- Args unit tests (CPU-only)
  test_state_dict_adapter.py  -- State dict roundtrip tests
  test_weight_sync.py     -- Weight completeness tests (meta device)
  test_forward.py         -- Forward precision comparison (single GPU)
  ...

Test stages (write tests appropriate for the model's complexity):

Stage 1: Args Tests (CPU-only, always write these)

Test fromhfconfig() with mock HuggingFace configs:

# Pattern: Create mock PretrainedConfig, verify args mapping
from unittest.mock import MagicMock

def test_args_from_hf_config():
    hf_config = MagicMock()
    hf_config.hidden_size = 4096
    hf_config.num_hidden_layers = 32
    # ... set all required fields
    args = <Model>ModelArgs.from_hf_config(hf_config)
    assert args.dim == 4096
    assert args.n_layers == 32

Stage 2: State Dict Adapter Tests (CPU-only)

Test key mapping roundtrip:

def test_state_dict_roundtrip():
    # Create adapter with mock config
    adapter = <Model>StateDictAdapter(mock_config)
    # Create fake archon state dict with expected keys
    archon_sd = {"tok_embeddings.weight": torch.randn(vocab, dim), ...}
    # Roundtrip
    hf_sd = adapter.to_hf(archon_sd)
    roundtrip = adapter.from_hf(hf_sd)
    assert set(roundtrip.keys()) == set(archon_sd.keys())

Stage 3: Weight Completeness (meta device, CPU-only)

Verify all model parameters have HF mappings:

def test_weight_completeness():
    # Create model on meta device
    with torch.device("meta"):
        model = <Model>Model(args)
    adapter = <Model>StateDictAdapter(hf_config)
    # Check every archon param has a HF mapping
    for name, _ in model.named_parameters():
        hf_pairs = adapter.convert_single_to_hf(name, torch.empty(0))
        assert len(hf_pairs) > 0, f"No HF mapping for {name}"

Stage 4: Forward Precision (single GPU, if available)

Compare Archon model output against HuggingFace reference:

@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
def test_forward_matches_hf():
    # Load both HF and Archon models
    # Run forward on same input
    # Compare logits within tolerance

Important: Do NOT hardcode the test categories. Inspect the existing test files in tests/experimental/archon/ and follow the same patterns, fixtures, and markers. Adapt test scope to the model's specific features (e.g., add MoE-specific tests only if the model has MoE).

Reference Implementations

Architecture Decision Map

Common Mistakes

• Not mapping all HF keys in statedictadapter.py (causes silent weight drops)
• Wrong fromhfconfig() field mapping (uses wrong HF config attribute name)
• Forgetting to handle None keys in fromhfmap (keys to skip like

rotaryemb.invfreq)

• Missing MoE expert weight 3D\<->2D conversion when model has MoE
• Wrong TP plan for attention with/without QK norm (uselocaloutput must match)
• Forgetting to add import line in areal/experimental/models/archon/init.py
• Not including all modeltype variants in supportedmodel_types frozenset
• Using print instead of areal.utils.logging.getLogger()

File Checklist

After completion, verify all files exist and are consistent:

• areal/experimental/models/archon/<model>/init.py
• areal/experimental/models/archon/<model>/spec.py -- ModelSpec + register
• areal/experimental/models/archon/<model>/model/args.py -- ModelArgs +

fromhfconfig

• areal/experimental/models/archon/<model>/model/model.py -- Model + Attention +

FFN

• areal/experimental/models/archon/<model>/model/rope.py -- RoPE (or re-export)
• areal/experimental/models/archon/<model>/model/statedictadapter.py -- Key

mapping

• areal/experimental/models/archon/<model>/infra/parallelize.py -- Parallel

strategy

• areal/experimental/models/archon/init.py -- Import line added
• tests/experimental/archon/test<model>*.py -- Tests

<!-- ================================================================================ MAINTAINER GUIDE ================================================================================

Location: .claude/skills/add-archon-model/SKILL.md Invocation: /add-archon-model <model_name>

Purpose

Semi-automated guide for adding new model architectures to the Archon training engine. Unlike simpler skills (add-reward, add-dataset), this skill actively guides Claude to:

1. Analyze HuggingFace source code to extract architecture details
2. Select the closest reference implementation (qwen2 or qwen3)
3. Generate code skeletons adapted to the target architecture
4. Create appropriate tests based on existing test patterns

How to Update

When New Reference Models Are Added

1. Add to "Reference Implementations" table
2. Update "Architecture Decision Map" with new feature columns
3. Update Step 2 (reference selection) with new options

When Base Classes Change

1. Update contract signatures in Steps 3, 4, 6, 7
2. Update file checklist if new files are required

When ModelSpec Changes

1. Update Step 8 with new ModelSpec fields
2. Update spec.py template

When Test Patterns Change

1. Update Step 10 with new test patterns
2. Do NOT hardcode test categories -- keep it flexible

Important Design Decisions

• This skill is SEMI-AUTOMATED: Claude should read HF source and generate code,

not just provide templates for the user to fill in manually

• The skill references existing test files rather than hardcoding test categories,

ensuring it stays current as the test suite evolves

• Reference model selection (qwen2 vs qwen3) is based on MoE and QK norm presence

================================================================================ -->