feat: load HuggingFace checkpoints via key/config conversion

- Add astrai.serialization.hf_adapter mapping LLaMA-style HF keys to AstrAI names (input_layernorm, gate_proj, MoE experts/shared_experts) with config aliases for dense and MoE (Mixtral/DeepSeek-V3) layouts; reject biased projections, mismatched head_dim and MLA
- Give AutoModel.from_pretrained weights_format=auto|astrai|hf with auto-detection; read sharded safetensors via model.safetensors.index.json
- Adapt preloaded weights/config in train_context and benchmark CLI

astrai/model/automodel.py

@@ -10,7 +10,15 @@ import torch.nn as nn
 
 from astrai.config.model_config import BaseModelConfig, ConfigFactory
 from astrai.factory import BaseFactory
-from astrai.serialization import load_model_config, load_model_weights, save_model
+from astrai.serialization import (
+    HF_MODEL_TYPES,
+    adapt_config,
+    convert_hf_weights,
+    load_model_config,
+    load_model_weights,
+    looks_like_hf_state_dict,
+    save_model,
+)
 
 
 @contextmanager
@@ -57,7 +65,25 @@ class AutoModel(nn.Module):
         path: Union[str, Path],
         disable_random_init: bool = True,
         strict: bool = True,
+        weights_format: str = "auto",
     ) -> nn.Module:
+        """Load a model directory.
+
+        Args:
+            path: Directory containing ``config.json`` and optionally
+                ``model.safetensors``.
+            disable_random_init: Replace parameter initializers with no-ops
+                while building the model.
+            strict: Passed to ``load_state_dict``.
+            weights_format: ``"auto"`` detects HuggingFace checkpoints
+                (LLaMA-style keys and ``model_type``) and converts them;
+                ``"astrai"`` skips conversion; ``"hf"`` forces it.
+        """
+        if weights_format not in ("auto", "astrai", "hf"):
+            raise ValueError(
+                f"weights_format must be one of 'auto', 'astrai', 'hf', "
+                f"got {weights_format!r}"
+            )
 
         model_path = Path(path)
 
@@ -66,6 +92,12 @@ class AutoModel(nn.Module):
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
         raw = load_model_config(str(model_path))
+        is_hf_config = weights_format == "hf" or (
+            weights_format == "auto" and raw.get("model_type") in HF_MODEL_TYPES
+        )
+        if is_hf_config:
+            raw = adapt_config(raw)
+
         config = ConfigFactory.load(raw)
         model_type = config.model_type or "autoregressive_lm"
 
@@ -75,8 +107,14 @@ class AutoModel(nn.Module):
             model = actual_cls(config)
 
         weights_path = model_path / "model.safetensors"
-        if weights_path.exists():
+        index_path = model_path / "model.safetensors.index.json"
+        if weights_path.exists() or index_path.exists():
             state_dict = load_model_weights(str(model_path))
+            is_hf_weights = is_hf_config or (
+                weights_format == "auto" and looks_like_hf_state_dict(state_dict)
+            )
+            if is_hf_weights:
+                state_dict = convert_hf_weights(state_dict, config)
             model.load_state_dict(state_dict, strict=strict)
 
         return model

astrai/serialization/__init__.py

@@ -22,9 +22,21 @@ from astrai.serialization.dataset import (
     load_bin_offsets,
     save_bin,
 )
+from astrai.serialization.hf_adapter import (
+    HF_MODEL_TYPES,
+    adapt_config,
+    convert_hf_config,
+    convert_hf_weights,
+    looks_like_hf_state_dict,
+)
 
 __all__ = [
     "Checkpoint",
+    "HF_MODEL_TYPES",
+    "adapt_config",
+    "convert_hf_config",
+    "convert_hf_weights",
+    "looks_like_hf_state_dict",
     "load_json",
     "load_model_config",
     "load_model_weights",

astrai/serialization/checkpoint.py

@@ -91,7 +91,21 @@ def load_model_config(save_directory: str) -> dict:
 
 
 def load_model_weights(save_directory: str) -> dict:
-    return load_state_dict(Path(save_directory) / _WEIGHTS_FILE)
+    save_path = Path(save_directory)
+    weights_file = save_path / _WEIGHTS_FILE
+    if weights_file.exists():
+        return load_state_dict(weights_file)
+
+    index_path = save_path / "model.safetensors.index.json"
+    if index_path.exists():
+        index = load_json(index_path)
+        weight_map = index.get("weight_map", {})
+        state_dict = {}
+        for shard in sorted(set(weight_map.values())):
+            state_dict.update(load_state_dict(save_path / shard))
+        return state_dict
+
+    raise FileNotFoundError(f"No model weights found in {save_directory}")
 
 
 def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
@@ -182,8 +196,10 @@ class Checkpoint:
         if meta_path.exists():
             return cls.load(save_dir, broadcast=broadcast)
 
-        if weights_path.exists():
+        weights_path = save_path / _WEIGHTS_FILE
-            state_dict = load_state_dict(weights_path, broadcast=broadcast)
+        index_path = save_path / "model.safetensors.index.json"
+        if weights_path.exists() or index_path.exists():
+            state_dict = load_model_weights(save_dir)
             config = {}
             config_path = save_path / _CONFIG_FILE
             if config_path.exists():

astrai/serialization/hf_adapter.py

@@ -0,0 +1,241 @@
+"""HuggingFace checkpoint adaptation for LLaMA-style decoder models.
+
+AstrAI stores weights with its own key names (``layers.<i>.input_norm``,
+``layers.<i>.mlp.gate``), while HuggingFace decoder-only checkpoints use
+``model.layers.<i>.input_layernorm`` / ``model.layers.<i>.mlp.gate_proj``.
+This module translates HF configs and state dicts so external checkpoints
+can be loaded directly.
+
+Supported families (LLaMA layout, dense and MoE):
+- dense FFN: llama, mistral, qwen2, gemma, gemma2, phi3
+- MoE FFN (Mixtral / Qwen2-MoE / DeepSeek-V3 layout): router
+  ``mlp.gate``, routed experts ``mlp.experts.<j>``, shared experts
+  ``mlp.shared_experts.<j>``
+
+Not supported:
+- MLA attention (DeepSeek-V2/V3 ``kv_a_proj_with_mqa``) uses a different
+  KV factorization and cannot be converted numerically.
+- Attention/MLP bias (``attention_bias`` / ``mlp_bias``) — AstrAI
+  projections are bias-free.
+"""
+
+import logging
+import re
+from typing import Any, Dict, Mapping, Union
+
+import torch
+
+from astrai.config.base import BaseConfig
+
+logger = logging.getLogger(__name__)
+
+HF_MODEL_TYPES = frozenset(
+    {
+        "llama",
+        "mistral",
+        "mixtral",
+        "qwen2",
+        "qwen2_moe",
+        "gemma",
+        "gemma2",
+        "phi3",
+    }
+)
+
+_EMBED = re.compile(r"^model\.embed_tokens\.weight$")
+_ATTN = re.compile(r"^model\.layers\.(\d+)\.self_attn\.(q|k|v|o)_proj\.(weight|bias)$")
+_INPUT_NORM = re.compile(r"^model\.layers\.(\d+)\.input_layernorm\.weight$")
+_POST_NORM = re.compile(r"^model\.layers\.(\d+)\.post_attention_layernorm\.weight$")
+_FINAL_NORM = re.compile(r"^model\.norm\.weight$")
+_LM_HEAD = re.compile(r"^lm_head\.weight$")
+_DENSE_MLP = re.compile(
+    r"^model\.layers\.(\d+)\.mlp\.(gate|up|down)_proj\.(weight|bias)$"
+)
+_MOE_ROUTER = re.compile(r"^model\.layers\.(\d+)\.mlp\.gate\.weight$")
+_MOE_EXPERTS = re.compile(
+    r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight|bias)$"
+)
+_MOE_SHARED = re.compile(
+    r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(\d+)\."
+    r"(gate|up|down)_proj\.(weight|bias)$"
+)
+
+_ASTR_PREFIXES = ("embed_tokens.", "layers.", "norm.", "lm_head.")
+
+
+def looks_like_hf_state_dict(state_dict: Mapping[str, Any]) -> bool:
+    """Return True if *state_dict* uses HuggingFace key names."""
+    return any(
+        key.startswith("model.")
+        or "self_attn." in key
+        or "input_layernorm" in key
+        or "mlp.experts." in key
+        for key in state_dict
+    )
+
+
+def adapt_config(raw: Dict[str, Any]) -> Dict[str, Any]:
+    """Translate *raw* for AstrAI if it looks like an HF model config."""
+    if raw.get("model_type") in HF_MODEL_TYPES:
+        return convert_hf_config(raw)
+    return raw
+
+
+def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert an HF LLaMA-style config dict to AstrAI field names."""
+    if raw.get("attention_bias") or raw.get("mlp_bias"):
+        raise NotImplementedError(
+            "attention_bias / mlp_bias checkpoints are not supported; "
+            "AstrAI projections are bias-free"
+        )
+
+    cfg: Dict[str, Any] = {}
+    for key in (
+        "vocab_size",
+        "hidden_size",
+        "num_hidden_layers",
+        "intermediate_size",
+        "rms_norm_eps",
+        "tie_word_embeddings",
+        "max_position_embeddings",
+        "rope_theta",
+        "rope_scaling",
+        "num_attention_heads",
+        "num_key_value_heads",
+        "use_qk_norm",
+        "use_gated_attention",
+        "kv_lora_rank",
+        "qk_nope_head_dim",
+        "qk_rope_head_dim",
+        "moe_intermediate_size",
+        "shared_expert_intermediate_size",
+        "topk_method",
+        "norm_topk_prob",
+        "moe_aux_loss_coef",
+        "neftune_alpha",
+    ):
+        if key in raw:
+            cfg[key] = raw[key]
+
+    if "qk_norm" in raw and "use_qk_norm" not in cfg:
+        cfg["use_qk_norm"] = raw["qk_norm"]
+
+    n_heads = raw.get("num_attention_heads")
+    if cfg.get("num_key_value_heads") is None and n_heads is not None:
+        cfg["num_key_value_heads"] = n_heads
+
+    if raw.get("head_dim") is not None and n_heads and raw.get("hidden_size"):
+        expected = raw["hidden_size"] // n_heads
+        if raw["head_dim"] != expected:
+            raise NotImplementedError(
+                f"HF head_dim={raw['head_dim']} differs from the computed "
+                f"head dim {expected}; AstrAI derives head_dim from "
+                "hidden_size / num_attention_heads"
+            )
+
+    if "kv_lora_rank" in raw:
+        cfg["attn_type"] = "mla"
+
+    n_experts = raw.get("num_local_experts") or raw.get("n_routed_experts")
+    if n_experts:
+        cfg["ffn_type"] = "moe"
+        cfg["n_routed_experts"] = n_experts
+        if "num_experts_per_tok" in raw:
+            cfg["n_activated_experts"] = raw["num_experts_per_tok"]
+        if "n_activated_experts" in raw:
+            cfg["n_activated_experts"] = raw["n_activated_experts"]
+        if "n_shared_experts" in raw:
+            cfg["n_shared_experts"] = raw["n_shared_experts"]
+        else:
+            # Mixtral has no shared experts; AstrAI defaults to one.
+            cfg["n_shared_experts"] = 0
+        if cfg.get("moe_intermediate_size") is None and "intermediate_size" in raw:
+            # MoE configs store the per-expert FFN size in intermediate_size.
+            cfg["moe_intermediate_size"] = raw["intermediate_size"]
+        first_k_dense = raw.get("first_k_dense_replace")
+        if isinstance(first_k_dense, int) and first_k_dense > 0:
+            cfg["mlp_only_layers"] = list(range(first_k_dense))
+            cfg["decoder_sparse_step"] = 1
+
+    cfg["model_type"] = "autoregressive_lm"
+    return cfg
+
+
+def convert_hf_weights(
+    state_dict: Mapping[str, Any],
+    config: BaseConfig,
+) -> Dict[str, torch.Tensor]:
+    """Rename HF state dict keys to AstrAI names.
+
+    Keys that are already AstrAI-style pass through unchanged; unmapped
+    HF keys are dropped with a warning. Use with ``strict=True`` to fail
+    loudly when the checkpoint does not match the config.
+    """
+    if getattr(config, "attn_type", "gqa") == "mla":
+        if any("kv_a_proj_with_mqa" in key for key in state_dict):
+            raise NotImplementedError(
+                "MLA attention (DeepSeek-V2/V3 kv_a_proj_with_mqa) uses a "
+                "different KV factorization and cannot be converted"
+            )
+
+    ffn_type = getattr(config, "ffn_type", "mlp")
+    converted: Dict[str, torch.Tensor] = {}
+    skipped: list[str] = []
+    for key, tensor in state_dict.items():
+        if key.startswith(_ASTR_PREFIXES):
+            converted[key] = tensor
+            continue
+
+        new_key = None
+        if ffn_type == "moe":
+            m = _MOE_ROUTER.match(key)
+            if m:
+                new_key = f"layers.{m.group(1)}.mlp.router.weight"
+            else:
+                m = _MOE_EXPERTS.match(key)
+                if m:
+                    new_key = (
+                        f"layers.{m.group(1)}.mlp.routed_experts.{m.group(2)}."
+                        f"{m.group(3)}.{m.group(4)}"
+                    )
+                else:
+                    m = _MOE_SHARED.match(key)
+                    if m:
+                        new_key = (
+                            f"layers.{m.group(1)}.mlp.shared_experts.{m.group(2)}."
+                            f"{m.group(3)}.{m.group(4)}"
+                        )
+        else:
+            m = _DENSE_MLP.match(key)
+            if m:
+                new_key = f"layers.{m.group(1)}.mlp.{m.group(2)}.{m.group(3)}"
+
+        if new_key is None:
+            m = _ATTN.match(key)
+            if m:
+                new_key = (
+                    f"layers.{m.group(1)}.attention.{m.group(2)}_proj.{m.group(3)}"
+                )
+            elif (m := _INPUT_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.input_norm.weight"
+            elif (m := _POST_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.post_attention_norm.weight"
+            elif (m := _EMBED.match(key)) is not None:
+                new_key = "embed_tokens.weight"
+            elif (m := _FINAL_NORM.match(key)) is not None:
+                new_key = "norm.weight"
+            elif (m := _LM_HEAD.match(key)) is not None:
+                new_key = "lm_head.weight"
+
+        if new_key is None:
+            skipped.append(key)
+        else:
+            converted[new_key] = tensor
+
+    if skipped:
+        logger.warning(
+            "Dropped %d unmapped HuggingFace weight key(s): %s",
+            len(skipped),
+            ", ".join(sorted(skipped)[:10]),
+        )
+    return converted

astrai/trainer/train_context.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader, random_split
 
+from astrai.config.model_config import ConfigFactory
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
 from astrai.inference.scheduler import InferenceScheduler
@@ -15,7 +16,13 @@ from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
-from astrai.serialization import Checkpoint, load_json
+from astrai.serialization import (
+    Checkpoint,
+    adapt_config,
+    convert_hf_weights,
+    load_json,
+    looks_like_hf_state_dict,
+)
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer.metric_util import GradSNRTracker
 from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
@@ -126,9 +133,18 @@ class TrainContextBuilder:
         if self._param_path:
             config_path = Path(self._param_path) / "config.json"
             if config_path.exists():
-                state.model_config = load_json(config_path)
+                state.model_config = adapt_config(load_json(config_path))
             checkpoint = Checkpoint.load_any(self._param_path)
             if checkpoint is not None:
+                if checkpoint.config:
+                    checkpoint.config = adapt_config(checkpoint.config)
+                if checkpoint.state_dict and looks_like_hf_state_dict(
+                    checkpoint.state_dict
+                ):
+                    checkpoint.state_dict = convert_hf_weights(
+                        checkpoint.state_dict,
+                        ConfigFactory.load(checkpoint.config or state.model_config),
+                    )
                 state.state_dict = checkpoint.state_dict
                 state.model_config = checkpoint.config or state.model_config
                 if self._resume:

docs/get-started.md

@@ -58,6 +58,14 @@ The model directory contains:
 - `model.safetensors` — model weights
 - `tokenizer.json` + `tokenizer_config.json` — tokenizer files (including chat template)
 
+External HuggingFace checkpoints of the LLaMA layout (e.g. `meta-llama/...`,
+`mistralai/...`, `Qwen/Qwen2-...`) can be loaded directly: `AutoModel.from_pretrained`
+auto-detects HF `model_type` / key names (`input_layernorm`, `gate_proj`, MoE
+`experts.<j>` ...) and converts config and weights in place. Dense and MoE
+(Mixtral / DeepSeek-V3 layout) FFNs are supported; MLA attention
+(DeepSeek-V2/V3) and biased projections (`attention_bias`) are not. Pass
+`weights_format="astrai"` to skip conversion, or `"hf"` to force it.
+
 ## 3. Run Inference
 
 ### Interactive Chat (Simplest)

scripts/tools/benchmark.py

@@ -13,6 +13,7 @@ from astrai.inference.engine import InferenceEngine
 from astrai.inference.runtime.graph import CudaGraphContext
 from astrai.inference.workspace import InferenceWorkspace
 from astrai.model import AutoModel, AutoRegressiveLM
+from astrai.serialization import adapt_config
 from astrai.tokenize import AutoTokenizer
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -478,8 +479,10 @@ def benchmark_command(
     if ckpt is not None:
         click.echo(f"Loading model from {ckpt} ...")
         config = ConfigFactory.load(
+            adapt_config(
                 json.loads((Path(ckpt) / "config.json").read_text(encoding="utf-8-sig"))
             )
+        )
         model = AutoModel.from_pretrained(ckpt)
     else:
         raw = dict(_DEFAULT_CONFIG)

tests/serialization/test_hf_adapter.py

@@ -0,0 +1,303 @@
+"""Tests for HuggingFace checkpoint/config adaptation."""
+
+import json
+
+import pytest
+import safetensors.torch as st
+import torch
+
+from astrai.config.model_config import ConfigFactory
+from astrai.model import AutoModel, AutoRegressiveLM
+from astrai.serialization import (
+    adapt_config,
+    convert_hf_config,
+    convert_hf_weights,
+    looks_like_hf_state_dict,
+    save_model,
+)
+from tests.helpers import assert_state_dicts_equal, make_tiny_config
+
+LLAMA_RAW = {
+    "architectures": ["LlamaForCausalLM"],
+    "model_type": "llama",
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.44.0",
+    "vocab_size": 1000,
+    "hidden_size": 8,
+    "num_hidden_layers": 2,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "intermediate_size": 16,
+    "max_position_embeddings": 64,
+    "rms_norm_eps": 1e-5,
+    "tie_word_embeddings": False,
+    "rope_theta": 10000.0,
+    "attention_bias": False,
+    "mlp_bias": False,
+    "head_dim": 4,
+}
+
+MOE_RAW = {
+    **LLAMA_RAW,
+    "model_type": "mixtral",
+    "intermediate_size": 16,
+    "num_local_experts": 2,
+    "num_experts_per_tok": 1,
+    "n_shared_experts": 1,
+}
+
+
+def to_hf_keys(state_dict):
+    """Rename AstrAI state dict keys to HuggingFace LLaMA-style names."""
+    out = {}
+    for key, tensor in state_dict.items():
+        if key == "embed_tokens.weight":
+            out["model.embed_tokens.weight"] = tensor
+        elif key == "norm.weight":
+            out["model.norm.weight"] = tensor
+        elif key.startswith("layers."):
+            parts = key.split(".")
+            layer = parts[1]
+            if parts[2] == "attention":
+                out[f"model.layers.{layer}.self_attn.{parts[3]}.{parts[4]}"] = tensor
+            elif parts[2] == "input_norm":
+                out[f"model.layers.{layer}.input_layernorm.weight"] = tensor
+            elif parts[2] == "post_attention_norm":
+                out[f"model.layers.{layer}.post_attention_layernorm.weight"] = tensor
+            elif parts[2] == "mlp":
+                if parts[3] in ("gate", "up", "down"):
+                    out[f"model.layers.{layer}.mlp.{parts[3]}_proj.weight"] = tensor
+                elif parts[3] == "router":
+                    out[f"model.layers.{layer}.mlp.gate.weight"] = tensor
+                elif parts[3] == "routed_experts":
+                    sub, name = parts[4], parts[5]
+                    out[
+                        f"model.layers.{layer}.mlp.experts.{sub}.{name}_proj.weight"
+                    ] = tensor
+                elif parts[3] == "shared_experts":
+                    sub, name = parts[4], parts[5]
+                    out[
+                        f"model.layers.{layer}.mlp.shared_experts.{sub}.{name}_proj.weight"
+                    ] = tensor
+        else:
+            out[key] = tensor
+    return out
+
+
+def test_convert_hf_config_llama():
+    cfg = convert_hf_config(LLAMA_RAW)
+    assert cfg["model_type"] == "autoregressive_lm"
+    assert cfg["hidden_size"] == 8
+    assert cfg["num_key_value_heads"] == 1
+    loaded = ConfigFactory.load(cfg)
+    assert loaded.num_attention_heads == 2
+    assert loaded.ffn_type == "mlp"
+
+
+def test_convert_hf_config_defaults_kv_heads():
+    raw = {k: v for k, v in LLAMA_RAW.items() if k != "num_key_value_heads"}
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.num_key_value_heads == 2
+
+
+def test_convert_hf_config_mixtral_moe():
+    cfg = convert_hf_config(MOE_RAW)
+    assert cfg["ffn_type"] == "moe"
+    assert cfg["n_routed_experts"] == 2
+    assert cfg["n_activated_experts"] == 1
+    assert cfg["n_shared_experts"] == 1
+    assert cfg["moe_intermediate_size"] == 16
+    loaded = ConfigFactory.load(cfg)
+    assert loaded.ffn_type == "moe"
+
+
+def test_convert_hf_config_mixtral_without_shared_experts():
+    raw = {k: v for k, v in MOE_RAW.items() if k != "n_shared_experts"}
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.n_shared_experts == 0
+
+
+def test_convert_hf_config_rejects_bias():
+    with pytest.raises(NotImplementedError):
+        convert_hf_config({**LLAMA_RAW, "attention_bias": True})
+
+
+def test_convert_hf_config_rejects_mismatched_head_dim():
+    with pytest.raises(NotImplementedError):
+        convert_hf_config({**LLAMA_RAW, "head_dim": 8})
+
+
+def test_looks_like_hf_state_dict():
+    assert looks_like_hf_state_dict({"model.layers.0.self_attn.q_proj.weight": 1})
+    assert looks_like_hf_state_dict({"model.embed_tokens.weight": 1})
+    assert not looks_like_hf_state_dict({"layers.0.attention.q_proj.weight": 1})
+
+
+def test_adapt_config_passthrough():
+    raw = dict(LLAMA_RAW, model_type="autoregressive_lm")
+    assert adapt_config(raw) is raw
+
+
+def test_convert_hf_weights_dense_roundtrip():
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg)
+    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
+def test_convert_hf_weights_moe_roundtrip():
+    cfg = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=2,
+        n_shared_experts=1,
+        n_activated_experts=1,
+        moe_intermediate_size=16,
+        shared_expert_intermediate_size=16,
+    )
+    model = AutoRegressiveLM(cfg)
+    hf_raw = convert_hf_config(MOE_RAW)
+    hf_cfg = ConfigFactory.load(hf_raw)
+    converted = convert_hf_weights(to_hf_keys(model.state_dict()), hf_cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
+def test_convert_hf_weights_keeps_astrai_keys():
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg)
+    converted = convert_hf_weights(dict(model.state_dict()), cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
+def test_convert_hf_weights_skips_unmapped_keys():
+    cfg = make_tiny_config()
+    sd = {"model.rotary_emb.inv_freq": torch.zeros(4), "model.embed_tokens.weight": 1}
+    converted = convert_hf_weights(sd, cfg)
+    assert "embed_tokens.weight" in converted
+    assert "model.rotary_emb.inv_freq" not in converted
+
+
+def test_convert_hf_weights_rejects_mla():
+    cfg = make_tiny_config(attn_type="mla", kv_lora_rank=2)
+    sd = {"model.layers.0.self_attn.kv_a_proj_with_mqa.weight": 1}
+    with pytest.raises(NotImplementedError):
+        convert_hf_weights(sd, cfg)
+
+
+def test_from_pretrained_hf_directory(tmp_path):
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg).eval()
+    save_model(
+        config=LLAMA_RAW,
+        state_dict=to_hf_keys(model.state_dict()),
+        save_directory=str(tmp_path),
+    )
+    loaded = AutoModel.from_pretrained(tmp_path).eval()
+    input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
+    with torch.no_grad():
+        torch.testing.assert_close(
+            loaded(input_ids)["logits"], model(input_ids)["logits"]
+        )
+
+
+def test_from_pretrained_astrai_directory(tmp_path):
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg).eval()
+    save_model(
+        config=cfg.to_dict(),
+        state_dict=model.state_dict(),
+        save_directory=str(tmp_path),
+    )
+    loaded = AutoModel.from_pretrained(tmp_path, disable_random_init=False)
+    assert_state_dicts_equal(loaded.state_dict(), model.state_dict())
+
+
+def test_from_pretrained_weights_format_hf_on_astrai_dir(tmp_path):
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg)
+    save_model(
+        config=cfg.to_dict(),
+        state_dict=model.state_dict(),
+        save_directory=str(tmp_path),
+    )
+    loaded = AutoModel.from_pretrained(
+        tmp_path, disable_random_init=False, weights_format="hf"
+    )
+    assert_state_dicts_equal(loaded.state_dict(), model.state_dict())
+
+
+def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path):
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg)
+    save_model(
+        config=LLAMA_RAW,
+        state_dict=to_hf_keys(model.state_dict()),
+        save_directory=str(tmp_path),
+    )
+    with pytest.raises(ValueError):
+        AutoModel.from_pretrained(tmp_path, weights_format="astrai")
+
+
+def test_from_pretrained_invalid_weights_format(tmp_path):
+    cfg = make_tiny_config()
+    save_model(
+        config=cfg.to_dict(),
+        state_dict={},
+        save_directory=str(tmp_path),
+    )
+    with pytest.raises(ValueError):
+        AutoModel.from_pretrained(tmp_path, weights_format="llama")
+
+
+def test_from_pretrained_hf_directory_sharded(tmp_path):
+    cfg = make_tiny_config()
+    model = AutoRegressiveLM(cfg).eval()
+    hf_sd = to_hf_keys(model.state_dict())
+    keys = sorted(hf_sd)
+    split = len(keys) // 2
+    shard_a = {k: hf_sd[k] for k in keys[:split]}
+    shard_b = {k: hf_sd[k] for k in keys[split:]}
+    st.save_file(shard_a, str(tmp_path / "model-00001-of-00002.safetensors"))
+    st.save_file(shard_b, str(tmp_path / "model-00002-of-00002.safetensors"))
+    index = {
+        "metadata": {},
+        "weight_map": {
+            k: (
+                "model-00001-of-00002.safetensors"
+                if k in shard_a
+                else "model-00002-of-00002.safetensors"
+            )
+            for k in keys
+        },
+    }
+    (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+    (tmp_path / "config.json").write_text(json.dumps(LLAMA_RAW))
+
+    loaded = AutoModel.from_pretrained(tmp_path).eval()
+    input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
+    with torch.no_grad():
+        torch.testing.assert_close(
+            loaded(input_ids)["logits"], model(input_ids)["logits"]
+        )
+
+
+def test_from_pretrained_hf_directory_with_moe(tmp_path):
+    cfg = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=2,
+        n_shared_experts=1,
+        n_activated_experts=1,
+        moe_intermediate_size=16,
+        shared_expert_intermediate_size=16,
+    )
+    model = AutoRegressiveLM(cfg).eval()
+    save_model(
+        config=MOE_RAW,
+        state_dict=to_hf_keys(model.state_dict()),
+        save_directory=str(tmp_path),
+    )
+    loaded = AutoModel.from_pretrained(tmp_path).eval()
+    input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
+    with torch.no_grad():
+        torch.testing.assert_close(
+            loaded(input_ids)["logits"], model(input_ids)["logits"]
+        )