()
| 18 | |
| 19 | |
| 20 | def main(): |
| 21 | # 获取脚本所在目录 |
| 22 | script_dir = Path(__file__).parent |
| 23 | project_root = script_dir.parent.parent |
| 24 | |
| 25 | parser = argparse.ArgumentParser(description="Quantize audio adapter model to FP8") |
| 26 | parser.add_argument( |
| 27 | "--model_path", |
| 28 | type=str, |
| 29 | default=str(project_root / "models" / "SekoTalk-Distill" / "audio_adapter_model.safetensors"), |
| 30 | help="Path to input model file", |
| 31 | ) |
| 32 | parser.add_argument( |
| 33 | "--output_path", |
| 34 | type=str, |
| 35 | default=str(project_root / "models" / "SekoTalk-Distill-fp8" / "audio_adapter_model_fp8.safetensors"), |
| 36 | help="Path to output quantized model file", |
| 37 | ) |
| 38 | args = parser.parse_args() |
| 39 | |
| 40 | model_path = Path(args.model_path) |
| 41 | output_path = Path(args.output_path) |
| 42 | |
| 43 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 44 | |
| 45 | state_dict = {} |
| 46 | with safetensors.safe_open(model_path, framework="pt", device="cpu") as f: |
| 47 | for key in f.keys(): |
| 48 | state_dict[key] = f.get_tensor(key) |
| 49 | |
| 50 | new_state_dict = {} |
| 51 | |
| 52 | for key in state_dict.keys(): |
| 53 | if key.startswith("ca") and ".to" in key and "weight" in key: |
| 54 | print(f"Converting {key} to FP8, dtype: {state_dict[key].dtype}") |
| 55 | |
| 56 | ## fp8 |
| 57 | weight = state_dict[key].to(torch.float32).cuda() |
| 58 | w_quantizer = FloatQuantizer("e4m3", True, "per_channel") |
| 59 | weight, weight_scale, _ = w_quantizer.real_quant_tensor(weight) |
| 60 | weight = weight.to(torch.float8_e4m3fn) |
| 61 | weight_scale = weight_scale.to(torch.float32) |
| 62 | |
| 63 | ## QuantWeightMxFP4, QuantWeightMxFP6, QuantWeightMxFP8 for mxfp4,mxfp6,mxfp8 |
| 64 | # weight = state_dict[key].to(torch.bfloat16).cuda() |
| 65 | # quantizer = QuantWeightMxFP4(weight) |
| 66 | # weight, weight_scale, _ = quantizer.weight_quant_func(weight) |
| 67 | |
| 68 | new_state_dict[key] = weight.cpu() |
| 69 | new_state_dict[key + "_scale"] = weight_scale.cpu() |
| 70 | else: |
| 71 | # 不匹配的权重转换为BF16 |
| 72 | print(f"Converting {key} to BF16, dtype: {state_dict[key].dtype}") |
| 73 | new_state_dict[key] = state_dict[key].to(torch.bfloat16) |
| 74 | |
| 75 | save_file(new_state_dict, str(output_path)) |
| 76 | print(f"Quantized model saved to: {output_path}") |
| 77 |
no test coverage detected