# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ python tools/convert_hf_to_fp8.py [-h] [--model-dir MODEL_DIR] [--save-dir SAVE_DIR] [--strategy {block,channel,tensor}] [--block-size [BLOCK_SIZE ...]] [--max-workers MAX_WORKERS] options: -h, --help show this help message and exit --model-dir MODEL_DIR Path to the directory of the HF safetensors model. --save-dir SAVE_DIR Path to the directory to save the converted model. --strategy {block,channel,tensor} --block-size [BLOCK_SIZE ...] eg. --block-size 32 32 --max-workers MAX_WORKERS Number of worker threads for parallel processing """ import argparse import gc import json import os import shutil import threading from concurrent.futures import ThreadPoolExecutor import safetensors import safetensors.torch import torch import torch.nn.functional as F from tqdm import tqdm FP8_INFO = torch.finfo(torch.float8_e4m3fn) FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min def ceildiv(a, b): return -(-a // b) def block_fp8(weight, block_size): # per block quant block_n, block_k = block_size[0], block_size[1] shape_0, shape_1 = weight.shape n_tiles = ceildiv(shape_0, block_n) k_tiles = ceildiv(shape_1, block_k) q_weight = F.pad( weight, (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0), mode="constant", value=0.0, ) qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k) block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0] block_max = torch.max(block_max, dim=3, keepdim=True)[0] scale = block_max.to(torch.float32) / FP8_MAX qweight = ( (qweight / scale) .clamp(min=FP8_MIN, max=FP8_MAX) .reshape((n_tiles * block_n, k_tiles * block_k)) .to(torch.float8_e4m3fn) ) qweight = qweight[:shape_0, :shape_1].clone().detach() scale = scale.squeeze() return qweight, scale def channel_fp8(weight): channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0] scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) qweight = qweight.to(torch.float8_e4m3fn) return qweight, scale def tensor_fp8(weight): scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) qweight = qweight.to(torch.float8_e4m3fn) scale = scale.view(1) return qweight, scale def quant_fp8(weight, strategy, block_size=None): if strategy == "tensor": return tensor_fp8(weight) elif strategy == "channel": return channel_fp8(weight) else: return block_fp8(weight, block_size) class ConversionResult: def __init__(self): self.lock = threading.Lock() self.weight_map = {} self.param_count = 0 self.modules_to_not_convert = [] def add_result(self, filename, q_weights, module_names): with self.lock: for k, v in q_weights.items(): self.weight_map[k] = filename self.param_count += len(v) self.modules_to_not_convert.extend(module_names) def process_file(input_path, output_path, filename, strategy, block_size, result_collector): if not filename.endswith(".safetensors"): return print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}") weights = {} q_weights = {} with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f: for k in f.keys(): weights[k] = f.get_tensor(k) modules_to_not_convert = [] for key in weights.keys(): if ( "weight" in key and "layernorm" not in key and "embed" not in key and "router" not in key and "mlp.gate." not in key and "norm" not in key and "lm_head" not in key and "eh_proj" not in key ): qw, s = quant_fp8(weights[key], strategy, block_size) q_weights[key] = qw if block_size: scale_name = key.replace(".weight", ".weight_scale_inv") else: scale_name = key.replace(".weight", ".weight_scale") q_weights[scale_name] = s else: modules_to_not_convert.append(key.replace(".weight", "")) q_weights[key] = weights[key] safetensors.torch.save_file(q_weights, os.path.join(output_path, filename), metadata={"format": "pt"}) result_collector.add_result(filename, q_weights, modules_to_not_convert) def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4): input_path = os.path.abspath(input_path) os.makedirs(output_path, exist_ok=True) for filename in os.listdir(input_path): if not filename.endswith(".safetensors") and not os.path.isdir(os.path.join(input_path, filename)): shutil.copyfile(os.path.join(input_path, filename), os.path.join(output_path, filename)) safetensors_files = [f for f in os.listdir(input_path) if f.endswith(".safetensors")] result_collector = ConversionResult() with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for filename in safetensors_files: future = executor.submit( process_file, input_path, output_path, filename, strategy, block_size, result_collector ) futures.append(future) for future in tqdm(futures, desc="Processing files"): future.result() if strategy == "block" or strategy == "tensor": quantization_config = { "activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", } if block_size: quantization_config["weight_block_size"] = block_size if len(result_collector.modules_to_not_convert) > 0: quantization_config["modules_to_not_convert"] = list(set(result_collector.modules_to_not_convert)) else: quant_group = { "group_0": { "input_activations": { "actorder": None, "block_structure": None, "dynamic": True, "group_size": None, "num_bits": 8, "observer": None, "observer_kwargs": {}, "strategy": "token", "symmetric": True, "type": "float", }, "output_activations": None, "targets": ["Linear"], "weights": { "actorder": None, "block_structure": None, "dynamic": False, "group_size": None, "num_bits": 8, "observer": "minmax", "observer_kwargs": {}, "strategy": strategy, "symmetric": True, "type": "float", }, }, } quantization_config = { "config_groups": quant_group, "format": "float-quantized", "ignore": list(set(result_collector.modules_to_not_convert)), "quant_method": "compressed-tensors", "quantization_status": "compressed", } config_path = os.path.join(input_path, "config.json") if os.path.exists(config_path): cfg = json.load(open(config_path)) cfg["quantization_config"] = quantization_config json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2) index_dict = {"weight_map": result_collector.weight_map, "metadata": {"total_size": result_collector.param_count}} json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2) gc.collect() torch.cuda.empty_cache() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model-dir", type=str, help="Path to the directory of the HF safetensors model.") parser.add_argument("--save-dir", type=str, help="Path to the directory to save the converted model.") parser.add_argument("--strategy", type=str, default="block", choices=["block", "channel", "tensor"]) parser.add_argument("--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32") parser.add_argument("--max-workers", type=int, default=1, help="Number of worker threads for parallel processing") args = parser.parse_args() if not os.path.exists(args.save_dir): print(f"Creating directory {args.save_dir}") os.makedirs(args.save_dir) elif not os.path.isdir(args.save_dir): raise ValueError("The save_dir should be a directory.") convert_fp8(args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers)