Adding Support for New HuggingFace Models to the Archon Engine in AReaL

To add support for a new HuggingFace model in AReaL's Archon engine, you must register a ModelSpec that defines the model class, HuggingFace configuration parser, parallelization strategy, and state-dict adapter.

The Archon engine is AReaL's torch-native training backend designed for distributed large language model training with tensor, context, pipeline, and expert parallelism. Adding support for new HuggingFace models to the Archon engine requires implementing a structured registration pattern that connects HuggingFace config.json schemas to AReaL's distributed training capabilities.

Understanding the Archon Engine and ModelSpec

The Archon engine supports tensor-parallel (TP), context-parallel (CP / Ulysses SP), pipeline-parallel (PP), and expert-parallel (EP) strategies. To integrate a new HuggingFace model, you create a ModelSpec dataclass that instructs the engine how to:

1. Instantiate the model on a meta device without memory allocation.
2. Parse HuggingFace configs into a BaseModelArgs dataclass.
3. Apply parallelism via a custom parallelize_fn.
4. Split into pipeline stages via an optional pipeline_fn.
5. Load and save weights through a BaseStateDictAdapter.

The registration mechanism lives in areal/experimental/models/archon/model_spec.py. When a module containing a spec is imported, register_model_spec adds the spec to the global _MODEL_SPECS dictionary. The engine validates model types via ArchonEngine._validate_model_type in areal/experimental/engine/archon_engine.py, which calls get_model_spec(model_type) to retrieve the registered specification.

Core Components of a ModelSpec

Model Arguments and Configuration

Create a dataclass inheriting from BaseModelArgs that mirrors the HuggingFace config.json fields. This file typically resides at areal/experimental/models/archon/{model_name}/model/args.py.

# my_gpt/model/args.py

from __future__ import annotations
from dataclasses import dataclass
from areal.experimental.models.archon.base import BaseModelArgs

@dataclass
class MyGPTModelArgs(BaseModelArgs):
    n_layers: int
    dim: int
    n_heads: int
    vocab_size: int
    max_seq_len: int = 8192
    qk_norm: bool = False
    moe_enabled: bool = False

Model Implementation

Implement the model class inheriting from BaseArchonModel, typically in areal/experimental/models/archon/{model_name}/model/model.py. This class defines the forward pass, embedding layers, transformer blocks, and normalization.

# my_gpt/model/model.py

from __future__ import annotations
import torch.nn as nn
from areal.experimental.models.archon.base import BaseArchonModel
from .args import MyGPTModelArgs

class MyGPTModel(BaseArchonModel):
    def __init__(self, model_args: MyGPTModelArgs):
        super().__init__()
        self.model_args = model_args
        self.tok_embeddings = nn.Embedding(
            model_args.vocab_size, 
            model_args.dim
        )
        self.layers = nn.ModuleDict({
            str(i): nn.TransformerEncoderLayer(
                d_model=model_args.dim,
                nhead=model_args.n_heads,
                dim_feedforward=model_args.dim * 4,
                batch_first=True,
            )
            for i in range(model_args.n_layers)
        })
        self.norm = nn.LayerNorm(model_args.dim)

    def forward(self, tokens, positions, cu_seqlens, max_seqlen, tree_attn_meta=None):
        h = self.tok_embeddings(tokens)
        for layer in self.layers.values():
            h = layer(h)
        return self.norm(h)

State-Dict Adapter

Create an adapter inheriting from BaseStateDictAdapter to map HuggingFace checkpoint keys to Archon parameter names. This resides in areal/experimental/models/archon/{model_name}/model/state_dict_adapter.py.

Parallelization Function

Implement a parallelize_fn that applies tensor, context, and expert parallelism. This typically calls archon_parallelize from areal/experimental/engine/archon_utils.

# my_gpt/infra/parallelize.py

import torch
from areal.experimental.engine.archon_utils import parallelize_archon
from .model import MyGPTModel

def parallelize_my_gpt(model: MyGPTModel, parallel_dims, **kwargs):
    return parallelize_archon(
        model=model,
        parallel_dims=parallel_dims,
        param_dtype=kwargs.get("param_dtype", torch.bfloat16),
        loss_parallel=True,
        cpu_offload=False,
        reshard_after_forward_policy="default",
        ac_config=None,
        enable_compile=kwargs.get("enable_compile", True),
    )

Pipeline Parallel Function (Optional)

If your model supports pipeline parallelism, set the pipelining_fn in your ModelSpec to a function like pipeline_llm from areal/experimental/models/archon/pipeline_parallel.py. This function computes virtual stages and splits the model into PipelineStage objects.

Step-by-Step Implementation Guide

Follow these steps to register a new HuggingFace model (e.g., "my-gpt") with the Archon engine:

1. Define ModelArgs: Create MyGPTModelArgs inheriting from BaseModelArgs to parse HuggingFace config fields.

2. Implement the Model: Write MyGPTModel inheriting from BaseArchonModel with the forward pass and layer definitions.

3. Create State-Dict Adapter: Implement MyGPTStateDictAdapter to map HF checkpoint keys to Archon parameters.

4. Write Parallelize Function: Define parallelize_my_gpt that calls archon_parallelize for TP/CP/EP support.

5. Register the Spec: Create a ModelSpec instance and call register_model_spec in your package's spec.py:

# my_gpt/spec.py

from areal.experimental.models.archon.model_spec import ModelSpec, register_model_spec
from areal.experimental.models.archon.pipeline_parallel import pipeline_llm
from .model.args import MyGPTModelArgs
from .model.model import MyGPTModel
from .model.state_dict_adapter import MyGPTStateDictAdapter
from .infra.parallelize import parallelize_my_gpt

MYGPT_SPEC = ModelSpec(
    name="MyGPT",
    model_class=MyGPTModel,
    model_args_class=MyGPTModelArgs,
    state_dict_adapter_class=MyGPTStateDictAdapter,
    parallelize_fn=parallelize_my_gpt,
    supported_model_types=frozenset({"my_gpt"}),
    pipelining_fn=pipeline_llm,  # Set to None if not using PP

)

register_model_spec(MYGPT_SPEC)
__all__ = ["MYGPT_SPEC"]

6. Expose the Model: Ensure your spec module is importable by adding it to areal/experimental/models/archon/__init__.py or relying on import side-effects.

7. Use the Model: Instantiate via the AReaL CLI or Python API:

from areal.api.cli_args import TrainEngineConfig, ArchonEngineConfig
from areal.experimental.engine.archon_engine import ArchonLMEngine

cfg = TrainEngineConfig(
    path="my_gpt_hf_repo",
    dtype="bfloat16",
    archon=ArchonEngineConfig(
        pp_schedule="1F1B",
        pp_layers_per_stage=12,
    ),
)

engine = ArchonLMEngine(cfg)
engine.create_process_group()
engine.initialize(addr=None, ft_spec=None)

Integration with the Archon Engine

When the engine initializes, ArchonEngine.__init__ reads the HuggingFace config.json, extracts the model_type field, and calls get_model_spec(model_type) from areal/experimental/models/archon/model_spec.py. If your spec is registered with a matching supported_model_types entry, the engine retrieves your ModelSpec and uses it to:

• Instantiate the model on a meta device via model_class
• Parse configuration via model_args_class
• Apply tensor, context, and expert parallelism via parallelize_fn
• Split the model into pipeline stages via pipelining_fn (if provided)
• Load and save checkpoints via state_dict_adapter_class

If the model_type is not found in the global registry, ArchonEngine._validate_model_type raises a clear error listing supported types via get_supported_model_types().

Summary

• Register a ModelSpec via areal/experimental/models/archon/model_spec.py to teach Archon how to handle a new HuggingFace model.
• Implement five core components: BaseModelArgs for config parsing, BaseArchonModel for the architecture, BaseStateDictAdapter for checkpoint mapping, a parallelize_fn for distributed training, and an optional pipelining_fn for pipeline parallelism.
• Use register_model_spec in your model's spec.py file to add the model to the global registry with a specific model_type identifier.
• Reference existing implementations like areal/experimental/models/archon/qwen3/spec.py for production-ready patterns.
• Validate integration by instantiating ArchonLMEngine with your HuggingFace model path; the engine automatically resolves the correct spec via get_model_spec.

Frequently Asked Questions

What is the Archon engine in AReaL?

The Archon engine is AReaL's native PyTorch training backend that supports large language model training with tensor-parallel (TP), context-parallel (CP), pipeline-parallel (PP), and expert-parallel (EP) strategies. It operates as the execution layer that takes HuggingFace model configurations and distributes them across GPU clusters using the specifications defined in a ModelSpec.

Do I need to implement pipeline parallelism for every new model?

No, pipeline parallelism is optional. When creating your ModelSpec in areal/experimental/models/archon/model_spec.py, you can set pipelining_fn=None if your model does not require pipeline parallelism. However, if you need PP support, you should set this to a function like pipeline_llm from areal/experimental/models/archon/pipeline_parallel.py, which handles stage generation and splitting.

How does the state-dict adapter handle HuggingFace checkpoint formats?

The BaseStateDictAdapter implementation maps HuggingFace checkpoint key names to the parameter names expected by your Archon model class. When the engine loads or saves checkpoints, it uses the adapter to translate between the HuggingFace state dict format (e.g., model.layers.0.self_attn.q_proj.weight) and your internal parameter naming convention. This ensures compatibility with pre-trained HuggingFace models while allowing flexible internal architectures.

Can I use existing parallelization utilities for custom models?

Yes, you should reuse the generic parallelization utilities provided by Archon rather than implementing distributed logic from scratch. In your parallelize_fn, call parallelize_archon from areal/experimental/engine/archon_utils, which automatically applies tensor-parallel sharding, FSDP wrapping, and context-parallel group registration. This ensures your custom model benefits from AReaL's optimized distributed training implementations without additional complexity.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}