初始化项目,由ModelHub XC社区提供模型
Model: divakar-yadav/transformer-1b-chat Source: Original Platform
This commit is contained in:
167
training_code/export_to_hf.py
Normal file
167
training_code/export_to_hf.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""
|
||||
Export the trained model to HuggingFace-compatible format.
|
||||
|
||||
Creates:
|
||||
- model.safetensors (weights)
|
||||
- config.json (architecture config)
|
||||
- generation_config.json
|
||||
- tokenizer.json, tokenizer_config.json, special_tokens_map.json
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import torch
|
||||
from collections import OrderedDict
|
||||
from safetensors.torch import save_file
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from model.config import ModelConfig
|
||||
from model.transformer import Transformer
|
||||
from model.data import get_tokenizer
|
||||
|
||||
CHECKPOINT = "/jfs/deepak-kumar/checkpoints_dpo/dpo_final.pt"
|
||||
OUTPUT_DIR = "/home/jovyan/training/hf_model"
|
||||
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
print("=" * 60)
|
||||
print(" EXPORTING MODEL TO HUGGING FACE FORMAT")
|
||||
print("=" * 60)
|
||||
|
||||
# --- 1. Load model ---
|
||||
print("\n[1/4] Loading checkpoint...")
|
||||
tokenizer = get_tokenizer()
|
||||
special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
|
||||
vocab = tokenizer.get_vocab()
|
||||
new_tokens = [t for t in special_tokens if t not in vocab]
|
||||
if new_tokens:
|
||||
tokenizer.add_tokens(new_tokens, special_tokens=True)
|
||||
|
||||
model_config = ModelConfig()
|
||||
model_config.vocab_size = len(tokenizer)
|
||||
|
||||
model = Transformer(model_config)
|
||||
ckpt = torch.load(CHECKPOINT, map_location="cpu", weights_only=False)
|
||||
model.load_state_dict(ckpt["model"])
|
||||
step = ckpt.get("step", 0)
|
||||
del ckpt
|
||||
print(f" Loaded DPO model (step {step}, vocab {model_config.vocab_size})")
|
||||
|
||||
# --- 2. Convert state dict keys to HF-style naming ---
|
||||
print("\n[2/4] Converting weights to safetensors...")
|
||||
|
||||
state_dict = model.state_dict()
|
||||
hf_state = OrderedDict()
|
||||
|
||||
KEY_MAP = {
|
||||
"tok_embeddings.weight": "model.embed_tokens.weight",
|
||||
"norm.weight": "model.norm.weight",
|
||||
"output.weight": "lm_head.weight",
|
||||
}
|
||||
|
||||
for key, tensor in state_dict.items():
|
||||
if key in KEY_MAP:
|
||||
hf_state[KEY_MAP[key]] = tensor
|
||||
continue
|
||||
|
||||
if key.startswith("layers."):
|
||||
parts = key.split(".")
|
||||
layer_idx = parts[1]
|
||||
rest = ".".join(parts[2:])
|
||||
|
||||
layer_map = {
|
||||
"attention_norm.weight": f"model.layers.{layer_idx}.input_layernorm.weight",
|
||||
"ffn_norm.weight": f"model.layers.{layer_idx}.post_attention_layernorm.weight",
|
||||
"attention.wq.weight": f"model.layers.{layer_idx}.self_attn.q_proj.weight",
|
||||
"attention.wk.weight": f"model.layers.{layer_idx}.self_attn.k_proj.weight",
|
||||
"attention.wv.weight": f"model.layers.{layer_idx}.self_attn.v_proj.weight",
|
||||
"attention.wo.weight": f"model.layers.{layer_idx}.self_attn.o_proj.weight",
|
||||
"ffn.w_gate.weight": f"model.layers.{layer_idx}.mlp.gate_proj.weight",
|
||||
"ffn.w_up.weight": f"model.layers.{layer_idx}.mlp.up_proj.weight",
|
||||
"ffn.w_down.weight": f"model.layers.{layer_idx}.mlp.down_proj.weight",
|
||||
}
|
||||
|
||||
if rest in layer_map:
|
||||
hf_state[layer_map[rest]] = tensor
|
||||
else:
|
||||
print(f" WARNING: unmapped key {key}")
|
||||
hf_state[key] = tensor
|
||||
elif key == "freqs_cis":
|
||||
continue
|
||||
else:
|
||||
print(f" WARNING: unmapped key {key}")
|
||||
hf_state[key] = tensor
|
||||
|
||||
# Convert all to bfloat16 for storage
|
||||
for k in hf_state:
|
||||
if hf_state[k].dtype == torch.float32:
|
||||
hf_state[k] = hf_state[k].to(torch.bfloat16)
|
||||
|
||||
safetensors_path = os.path.join(OUTPUT_DIR, "model.safetensors")
|
||||
save_file(hf_state, safetensors_path)
|
||||
size_gb = os.path.getsize(safetensors_path) / 1e9
|
||||
print(f" Saved {len(hf_state)} tensors -> {safetensors_path} ({size_gb:.2f} GB)")
|
||||
|
||||
# --- 3. Write config files ---
|
||||
print("\n[3/4] Writing config files...")
|
||||
|
||||
config_json = {
|
||||
"architectures": ["LlamaForCausalLM"],
|
||||
"model_type": "llama",
|
||||
"vocab_size": model_config.vocab_size,
|
||||
"hidden_size": model_config.hidden_dim,
|
||||
"intermediate_size": model_config.intermediate_dim,
|
||||
"num_hidden_layers": model_config.num_layers,
|
||||
"num_attention_heads": model_config.num_attention_heads,
|
||||
"num_key_value_heads": model_config.num_kv_heads,
|
||||
"max_position_embeddings": model_config.max_seq_len,
|
||||
"rope_theta": model_config.rope_theta,
|
||||
"rms_norm_eps": model_config.rms_norm_eps,
|
||||
"hidden_act": "silu",
|
||||
"initializer_range": 0.02,
|
||||
"tie_word_embeddings": False,
|
||||
"torch_dtype": "bfloat16",
|
||||
"transformers_version": "4.40.0",
|
||||
"use_cache": True,
|
||||
"bos_token_id": tokenizer.bos_token_id,
|
||||
"eos_token_id": tokenizer.eos_token_id,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
|
||||
with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
|
||||
json.dump(config_json, f, indent=2)
|
||||
print(" config.json")
|
||||
|
||||
gen_config = {
|
||||
"bos_token_id": tokenizer.bos_token_id,
|
||||
"eos_token_id": tokenizer.eos_token_id,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
"do_sample": True,
|
||||
"temperature": 0.7,
|
||||
"top_k": 50,
|
||||
"top_p": 0.9,
|
||||
"repetition_penalty": 1.15,
|
||||
"max_new_tokens": 512,
|
||||
"transformers_version": "4.40.0",
|
||||
}
|
||||
|
||||
with open(os.path.join(OUTPUT_DIR, "generation_config.json"), "w") as f:
|
||||
json.dump(gen_config, f, indent=2)
|
||||
print(" generation_config.json")
|
||||
|
||||
# --- 4. Export tokenizer ---
|
||||
print("\n[4/4] Exporting tokenizer...")
|
||||
tokenizer.save_pretrained(OUTPUT_DIR)
|
||||
print(" Tokenizer files saved")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(" EXPORT COMPLETE -> " + OUTPUT_DIR)
|
||||
print("=" * 60)
|
||||
print("\nFiles:")
|
||||
for f in sorted(os.listdir(OUTPUT_DIR)):
|
||||
size = os.path.getsize(os.path.join(OUTPUT_DIR, f))
|
||||
if size > 1e6:
|
||||
print(f" {f:40s} {size/1e6:.1f} MB")
|
||||
else:
|
||||
print(f" {f:40s} {size/1e3:.1f} KB")
|
||||
Reference in New Issue
Block a user