File size: 5,745 Bytes
296a506 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/usr/bin/env python3
"""
Frox AI Morph 1.1 — Export to HuggingFace / vLLM-loadable format
Morph is a custom architecture, not a HuggingFace AutoModel. To serve it
with vLLM in production (see the backend architecture doc), it needs:
1. A HuggingFace-style directory: config.json, model.safetensors,
tokenizer files, generation_config.json
2. A registered `AutoModelForCausalLM` mapping OR a vLLM model plugin
that knows how to read Morph's config and weight names
This script produces (1) — the HF-format directory — which is the input
vLLM's `--model` flag expects. It does NOT write the vLLM plugin itself
(that's a `vllm.model_executor.models` registration, versioned against a
specific vLLM release), but it lays out weights in a way that a plugin's
`load_weights()` can consume directly since parameter names match the
saved state dict 1:1.
Usage:
python scripts/export_hf.py --model ./frox-morph-1-1-output/sft_final \
--output ./frox-morph-1-1-hf
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import torch
from model.architecture.morph_model import MorphForCausalLM
from utils.common import print_banner
def export_to_hf_format(model_path: str, output_path: str, save_dtype: str = "float16"):
print(f"📂 Loading Morph 1.1 from {model_path}...")
model = MorphForCausalLM.from_saved(model_path, device="cpu")
cfg = model.config
out = Path(output_path)
out.mkdir(parents=True, exist_ok=True)
# ── config.json (HF-style, with Morph-specific fields preserved) ──
hf_config = {
"architectures": ["MorphForCausalLM"],
"model_type": "morph",
"vocab_size": cfg.total_vocab_size,
"hidden_size": cfg.hidden_size,
"intermediate_size": cfg.intermediate_size,
"num_hidden_layers": cfg.num_hidden_layers,
"num_attention_heads": cfg.num_attention_heads,
"num_key_value_heads": cfg.num_key_value_heads,
"head_dim": cfg.head_dim,
"max_position_embeddings": cfg.max_position_embeddings,
"rope_theta": cfg.rope_theta,
"rope_scaling": {
"type": "yarn",
"factor": cfg.rope_scaling_factor,
"original_max_position_embeddings": 2048,
},
"rms_norm_eps": cfg.rms_norm_eps,
"hidden_act": cfg.hidden_act,
"tie_word_embeddings": cfg.tie_word_embeddings,
"torch_dtype": save_dtype,
"pad_token_id": cfg.pad_token_id,
"bos_token_id": cfg.bos_token_id,
"eos_token_id": cfg.eos_token_id,
# Morph-specific — a vLLM plugin reads these to configure attention
"qk_norm": cfg.qk_norm,
"use_sliding_window": cfg.use_sliding_window,
"sliding_window_size": cfg.sliding_window_size,
"transformers_version": "4.47.0",
"morph_version": "1.1.0",
}
with open(out / "config.json", "w") as f:
json.dump(hf_config, f, indent=2)
print(f" ✓ config.json")
# ── generation_config.json ──
gen_config = {
"bos_token_id": cfg.bos_token_id,
"eos_token_id": cfg.eos_token_id,
"pad_token_id": cfg.pad_token_id,
"do_sample": True,
"temperature": 0.7,
"top_p": 0.9,
"top_k": 50,
"max_new_tokens": 4096,
"repetition_penalty": 1.1,
}
with open(out / "generation_config.json", "w") as f:
json.dump(gen_config, f, indent=2)
print(f" ✓ generation_config.json")
# ── model.safetensors ──
try:
from safetensors.torch import save_file
dtype = getattr(torch, save_dtype)
state_dict = {k: v.to(dtype).contiguous() for k, v in model.state_dict().items()}
# If weights are tied, safetensors refuses to save duplicate storage —
# clone the lm_head weight so both tensors have independent storage.
if cfg.tie_word_embeddings and "lm_head.weight" in state_dict:
state_dict["lm_head.weight"] = state_dict["lm_head.weight"].clone()
save_file(state_dict, str(out / "model.safetensors"), metadata={"format": "pt"})
print(f" ✓ model.safetensors ({sum(v.numel() for v in state_dict.values()):,} params)")
except ImportError:
print(" ⚠ safetensors not installed — saving as model.pt instead")
torch.save(model.state_dict(), out / "model.pt")
# ── tokenizer ──
tokenizer_src = Path(model_path) / "tokenizer"
if tokenizer_src.exists():
import shutil
shutil.copytree(tokenizer_src, out / "tokenizer_files", dirs_exist_ok=True)
for f in (out / "tokenizer_files").iterdir():
f.rename(out / f.name)
(out / "tokenizer_files").rmdir()
print(f" ✓ tokenizer files")
else:
print(f" ⚠ No tokenizer found at {tokenizer_src} — copy it manually before serving")
print(f"\n✅ Exported to {output_path}")
print(f"\nNext steps:")
print(f" 1. Register the Morph architecture with vLLM (model_executor/models/morph.py)")
print(f" 2. Serve: vllm serve {output_path} --dtype {save_dtype} "
f"--max-model-len {cfg.max_position_embeddings}")
def main():
parser = argparse.ArgumentParser(description="Export Morph 1.1 to HF/vLLM format")
parser.add_argument("--model", type=str, required=True)
parser.add_argument("--output", type=str, required=True)
parser.add_argument("--dtype", choices=["float16", "bfloat16", "float32"], default="float16")
args = parser.parse_args()
print_banner()
export_to_hf_format(args.model, args.output, args.dtype)
if __name__ == "__main__":
main()
|