From 7d27f3e0786228440ddfc349f27bc8b6b7600430 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 20 Aug 2026 11:34:48 +0800 Subject: [PATCH] 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 | 42 +++- astrai/serialization/__init__.py | 12 + astrai/serialization/checkpoint.py | 22 +- astrai/serialization/hf_adapter.py | 241 ++++++++++++++++++++ astrai/trainer/train_context.py | 20 +- docs/get-started.md | 8 + scripts/tools/benchmark.py | 5 +- tests/serialization/test_hf_adapter.py | 303 +++++++++++++++++++++++++ 8 files changed, 645 insertions(+), 8 deletions(-) create mode 100644 astrai/serialization/hf_adapter.py create mode 100644 tests/serialization/test_hf_adapter.py diff --git a/astrai/model/automodel.py b/astrai/model/automodel.py index 99b1d7d..0a00996 100644 --- a/astrai/model/automodel.py +++ b/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 diff --git a/astrai/serialization/__init__.py b/astrai/serialization/__init__.py index 501f886..0ac5e0d 100644 --- a/astrai/serialization/__init__.py +++ b/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", diff --git a/astrai/serialization/checkpoint.py b/astrai/serialization/checkpoint.py index 2bef9ae..d72c699 100644 --- a/astrai/serialization/checkpoint.py +++ b/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(): - state_dict = load_state_dict(weights_path, broadcast=broadcast) + weights_path = save_path / _WEIGHTS_FILE + 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(): diff --git a/astrai/serialization/hf_adapter.py b/astrai/serialization/hf_adapter.py new file mode 100644 index 0000000..591ed76 --- /dev/null +++ b/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..input_norm``, +``layers..mlp.gate``), while HuggingFace decoder-only checkpoints use +``model.layers..input_layernorm`` / ``model.layers..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.``, shared experts + ``mlp.shared_experts.`` + +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 diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index c29f1e6..e41a695 100644 --- a/astrai/trainer/train_context.py +++ b/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: diff --git a/docs/get-started.md b/docs/get-started.md index c09cd55..8e23196 100644 --- a/docs/get-started.md +++ b/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.` ...) 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) diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index 45a4ad6..c4d4031 100644 --- a/scripts/tools/benchmark.py +++ b/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,7 +479,9 @@ def benchmark_command( if ckpt is not None: click.echo(f"Loading model from {ckpt} ...") config = ConfigFactory.load( - json.loads((Path(ckpt) / "config.json").read_text(encoding="utf-8-sig")) + adapt_config( + json.loads((Path(ckpt) / "config.json").read_text(encoding="utf-8-sig")) + ) ) model = AutoModel.from_pretrained(ckpt) else: diff --git a/tests/serialization/test_hf_adapter.py b/tests/serialization/test_hf_adapter.py new file mode 100644 index 0000000..f52ca1b --- /dev/null +++ b/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"] + )