ViperEkura
75304d084d
- add skip_no_fp8 (CUDA + fp8_mm kernel + cc 8.9+) to tests/conftest.py - use skip_no_cuda / skip_no_kernel / skip_no_fp8 directly in test modules - drop _GPU alias and tests.extension.conftest re-exports - remove unused imports (Union in hf_adapter, make_grpo_config in data conftest)
272 lines
9.7 KiB
Python
272 lines
9.7 KiB
Python
"""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
|
|
|
|
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)$")
|
|
_Q_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.q_norm\.weight$")
|
|
_K_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.k_norm\.weight$")
|
|
_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_expert(?:s)?\.(\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 _is_dense_mlp_layer(config: BaseConfig, layer_id: int) -> bool:
|
|
"""Return whether a layer uses dense MLP instead of routed experts."""
|
|
if getattr(config, "ffn_type", "mlp") != "moe":
|
|
return True
|
|
mlp_only = getattr(config, "mlp_only_layers", None) or []
|
|
if layer_id in mlp_only:
|
|
return True
|
|
step = getattr(config, "decoder_sparse_step", 1) or 1
|
|
return step > 1 and (layer_id + 1) % step != 0
|
|
|
|
|
|
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",
|
|
"decoder_sparse_step",
|
|
"mlp_only_layers",
|
|
"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"]
|
|
if (
|
|
raw.get("model_type") in ("gemma", "gemma2")
|
|
and "use_qk_norm" not in cfg
|
|
and "qk_norm" not in raw
|
|
):
|
|
# Gemma/Gemma2 always apply RMSNorm to Q and K before attention.
|
|
cfg["use_qk_norm"] = True
|
|
|
|
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)}"
|
|
)
|
|
if new_key is None:
|
|
m = _DENSE_MLP.match(key)
|
|
if m and _is_dense_mlp_layer(config, int(m.group(1))):
|
|
new_key = f"layers.{m.group(1)}.mlp.{m.group(2)}.{m.group(3)}"
|
|
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 := _Q_NORM.match(key)) is not None:
|
|
new_key = f"layers.{m.group(1)}.attention.q_norm.weight"
|
|
elif (m := _K_NORM.match(key)) is not None:
|
|
new_key = f"layers.{m.group(1)}.attention.k_norm.weight"
|
|
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
|