File size: 2,792 Bytes
2acf66a | 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 | import os
import argparse
import torch
def convert_pt_to_fp8(input_file, output_file):
try:
# 获取当前运行目录的绝对路径，锁定范围
current_dir = os.getcwd()
in_path = os.path.abspath(os.path.join(current_dir, input_file))
out_path = os.path.abspath(os.path.join(current_dir, output_file))
# 严格限制：只允许处理当前目录下的文件
if os.path.dirname(in_path) != current_dir or os.path.dirname(out_path) != current_dir:
print("❌ 错误：基于安全限制，本脚本只允许处理当前目录下的文件！")
return
if not os.path.exists(in_path):
print(f"❌ 错误：找不到输入文件 '{input_file}'")
return
print(f"📦 正在加载 FP16 模型权重: {input_file} ...")
# 视你的硬件情况，可将 map_location 设为 'cuda' 或 'cpu'
state_dict = torch.load(in_path, map_location='cpu')
print("⚡ 正在将浮点数转换为底层 FP8 (E4M3FN) 格式...")
fp8_state_dict = {}
# 判断载入的是直接的 OrderedDict 还是包裹在字典里的权重
if isinstance(state_dict, dict) and 'model' in state_dict:
target_dict = state_dict['model']
is_wrapped = True
else:
target_dict = state_dict
is_wrapped = False
for k, v in target_dict.items():
if isinstance(v, torch.Tensor):
# 只有浮点数（FP32/FP16）需要转为 FP8，整数（如层索引、形状）原样保留
if v.is_floating_point():
# 转换为 PyTorch 标准的 FP8 权重存储格式
fp8_state_dict[k] = v.to(dtype=torch.float8_e4m3fn)
else:
fp8_state_dict[k] = v
else:
fp8_state_dict[k] = v
# 保持原有的包裹外壳结构
output_data = {'model': fp8_state_dict} if is_wrapped else fp8_state_dict
print(f"💾 正在保存 FP8 模型至: {output_file} ...")
torch.save(output_data, out_path)
print(f"✅ 成功生成 FP8 权重文件！")
except Exception as e:
print(f"❌ 转换失败，错误原因: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="将 FP16/FP32 的 .pt 权重转换为 FP8(E4M3) 格式（仅限当前目录）")
parser.add_argument("-i", "--input", required=True, help="输入的 .pt 权重文件名")
parser.add_argument("-o", "--output", required=True, help="输出的 FP8 .pt 文件名")
args = parser.parse_args()
convert_pt_to_fp8(args.input, args.output) |