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

> Expand AReaL's Archon engine by registering ModelSpecs to support new HuggingFace models. Define model class, parser, and strategies seamlessly.

- Repository: [inclusionAI/areal](https://github.com/inclusionai/areal)
- Tags: how-to-guide
- Published: 2026-03-04

---

**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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/config.json) fields. This file typically resides at `areal/experimental/models/archon/{model_name}/model/args.py`.

```python

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

```python

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

```python

# 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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/spec.py):

```python

# 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`](https://github.com/inclusionai/areal/blob/main/areal/experimental/models/archon/__init__.py) or relying on import side-effects.

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

```python
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`](https://github.com/inclusionai/areal/blob/main/config.json), extracts the `model_type` field, and calls `get_model_spec(model_type)` from [`areal/experimental/models/archon/model_spec.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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.