mirror of
https://github.com/lucasjinreal/Namo-R1.git
synced 2026-01-13 22:07:17 +08:00
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
import torch
|
||
from transformers import AutoModel, AutoTokenizer
|
||
import argparse
|
||
import os
|
||
|
||
|
||
def convert_and_save_bf16(model_path, output_dir=None):
|
||
|
||
try:
|
||
if output_dir is None:
|
||
output_dir = model_path.strip("/") + "_bf16"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
print(f"⏳ 正在加载原始模型来自: {model_path}")
|
||
|
||
model = AutoModel.from_pretrained(
|
||
model_path,
|
||
torch_dtype=torch.bfloat16, # 初始加载为BF16
|
||
# device_map="auto", # 自动分配设备
|
||
low_cpu_mem_usage=True, # 优化内存使用
|
||
trust_remote_code=True,
|
||
)
|
||
|
||
print("🔧 正在转换模型权重到BF16...")
|
||
model = model.to(torch.bfloat16)
|
||
|
||
print(f"💾 正在保存BF16模型到: {output_dir}")
|
||
model.save_pretrained(
|
||
output_dir,
|
||
safe_serialization=True, # 使用safetensors格式
|
||
max_shard_size="6GB", # 分片大小
|
||
)
|
||
|
||
# 保存tokenizer
|
||
try:
|
||
tokenizer = AutoTokenizer.from_pretrained(
|
||
model_path, trust_remote_code=True
|
||
)
|
||
tokenizer.save_pretrained(output_dir)
|
||
except Exception as e:
|
||
print("passing save tokenzier.")
|
||
|
||
print("✅ 转换完成！保存内容：")
|
||
print(f" - 模型权重: {output_dir}/pytorch_model*.bin")
|
||
print(f" - 配置文件: {output_dir}/config.json")
|
||
print(f" - Tokenizer文件: {output_dir}/tokenizer.*")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 错误发生: {str(e)}")
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="转换HF模型到BF16格式")
|
||
parser.add_argument(
|
||
"model_path",
|
||
type=str,
|
||
help="输入模型路径（本地目录或HF Hub名称）",
|
||
)
|
||
parser.add_argument("--output_dir")
|
||
|
||
args = parser.parse_args()
|
||
|
||
convert_and_save_bf16(args.model_path, args.output_dir)
|