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
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"]
|
||||
)
|
||||
Reference in New Issue
Block a user