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:
2026-08-20 11:34:59 +08:00
parent 84753d3e08
commit 7d27f3e078
8 changed files with 645 additions and 8 deletions
+4 -1
View File
@@ -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: