| 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} ...")
|
|
|
| state_dict = torch.load(in_path, map_location='cpu')
|
|
|
| print("⚡ 正在将浮点数转换为底层 FP8 (E4M3FN) 格式...")
|
| fp8_state_dict = {}
|
|
|
|
|
| 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):
|
|
|
| if v.is_floating_point():
|
|
|
| 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) |
