| import os |
| import json |
| from argparse import ArgumentParser |
| from glob import glob |
| from tqdm import tqdm |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
| from huggingface_hub import snapshot_download |
|
|
| from kernel import weight_quant |
|
|
| def main(bf16_path, int8_path, model_name="deepseek-ai/DeepSeek-R1"): |
| torch.set_default_dtype(torch.bfloat16) |
| os.makedirs(int8_path, exist_ok=True) |
| model_index_file = os.path.join(int8_path, "model.safetensors.index.json") |
| |
| if not os.path.exists(model_index_file): |
| snapshot_download( |
| repo_id=model_name, |
| allow_patterns=["model.safetensors.index.json"], |
| local_dir=int8_path, |
| local_dir_use_symlinks=False |
| ) |
| print(f"model index file downloaded to {model_index_file}") |
|
|
| with open(model_index_file, "r") as f: |
| model_index = json.load(f) |
| weight_map = model_index["weight_map"] |
| scale_count = len([key for key in weight_map.keys() if key.endswith("_scale_inv")]) |
| |
| safetensor_files = list(glob(os.path.join(bf16_path, "*.safetensors"))) |
| safetensor_files.sort() |
| quant_count = 0 |
| for safetensor_file in tqdm(safetensor_files): |
| file_name = os.path.basename(safetensor_file) |
| state_dict = load_file(safetensor_file, device="cuda") |
| new_state_dict = {} |
| for weight_name, weight in state_dict.items(): |
| scale_inv_name = f"{weight_name}_scale_inv" |
| if scale_inv_name in weight_map: |
| assert weight.element_size() == 2 |
| quant_count += 1 |
| int8_weight, scale_inv = weight_quant(weight) |
| new_state_dict[weight_name] = int8_weight |
| new_state_dict[scale_inv_name] = scale_inv |
| else: |
| new_state_dict[weight_name] = weight |
| new_safetensor_file = os.path.join(int8_path, file_name) |
| save_file(new_state_dict, new_safetensor_file) |
| assert quant_count == scale_count |
| print(f"{quant_count} weights are quantized.") |
| |
|
|
| if __name__ == "__main__": |
| parser = ArgumentParser() |
| parser.add_argument("--input-bf16-hf-path", type=str, required=True) |
| parser.add_argument("--output-int8-hf-path", type=str, required=True) |
| parser.add_argument("--model-name", type=str, default="deepseek-ai/DeepSeek-R1") |
| args = parser.parse_args() |
| main(args.input_bf16_hf_path, args.output_int8_hf_path, args.model_name) |
| print("done") |
| |
