| import argparse |
| import copy |
| import json |
| import shutil |
| import time |
| from pathlib import Path |
|
|
| import modelopt.torch.quantization as mtq |
| import torch |
| from modelopt.torch.export import export_hf_checkpoint |
| from safetensors import safe_open |
| from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--source', required=True) |
| parser.add_argument('--output', required=True) |
| args = parser.parse_args() |
| source, output = Path(args.source), Path(args.output) |
| output.mkdir(parents=True, exist_ok=False) |
| torch.set_num_threads(4) |
| torch.manual_seed(20261020) |
| config = AutoConfig.from_pretrained(source, trust_remote_code=True, local_files_only=True) |
| config.backend = 'torch' |
| config.prefill_chunk_size = 512 |
| model = AutoModelForCausalLM.from_pretrained( |
| source, config=config, trust_remote_code=True, local_files_only=True, |
| dtype=torch.bfloat16, device_map={'': 0}, attn_implementation='sdpa', |
| ).eval().requires_grad_(False) |
| tokenizer = AutoTokenizer.from_pretrained(source, trust_remote_code=True, local_files_only=True) |
| tokenizer.chat_template = (Path(__file__).parent/'chat_template.jinja').read_text() |
| model.generation_config.eos_token_id = [248046, 248044] |
| model.generation_config.pad_token_id = tokenizer.pad_token_id |
| modules = {} |
| recipe = copy.deepcopy(mtq.W4A16_NVFP4_CFG) |
| for name, module in model.named_modules(): |
| if not isinstance(module, torch.nn.Linear): |
| continue |
| if name != 'lm_head' and min(module.weight.shape) >= 128 and '.admission_' not in name: |
| assert module.in_features % 16 == 0 and module.out_features % 16 == 0, name |
| modules[name] = list(module.weight.shape) |
| else: |
| recipe['quant_cfg'].append({'quantizer_name': name+'.*', 'enable': False}) |
| assert modules |
| probe = tokenizer('The sum of 17 and 23 is', return_tensors='pt').to('cuda') |
| started = time.perf_counter() |
| print(json.dumps({'stage': 'quantizing', 'modules': len(modules), 'format': 'W4A16_NVFP4'}), flush=True) |
| with torch.inference_mode(): |
| mtq.quantize(model, recipe, forward_loop=lambda current: current(**probe, use_cache=False, logits_to_keep=1)) |
| quantized_probe = model(**probe, use_cache=False, logits_to_keep=1).logits.float().cpu() |
| assert torch.isfinite(quantized_probe).all() |
| torch.save(quantized_probe, output.parent/'quantizer_probe.pt') |
| print(json.dumps({'stage': 'exporting'}), flush=True) |
| export_hf_checkpoint(model, export_dir=output, save_modelopt_state=False, max_shard_size='4GB') |
| tokenizer.save_pretrained(output) |
| for name in ('configuration_edgerock.py', 'modeling_edgerock.py', 'LICENSE', 'LICENSE-CODE', 'NOTICE'): |
| shutil.copy2(source/name, output/name) |
| for name in ('load_nvfp4.py', 'quantize.py'): |
| shutil.copy2(Path(__file__).parent/name, output/name) |
| shutil.copy2(Path(__file__).parent/'chat_template.jinja', output/'chat_template.jinja') |
| hf_config = json.loads((output/'hf_quant_config.json').read_text()) |
| assert hf_config['quantization']['quant_algo'] == 'W4A16_NVFP4', hf_config |
| manifest = { |
| 'format': 'NVFP4', 'recipe': 'W4A16_NVFP4', 'weight_encoding': 'E2M1, two values per byte, low nibble first', |
| 'block_size': 16, 'block_scale_dtype': 'float8_e4m3fn', 'global_scale_dtype': 'float32', |
| 'activation_dtype': 'bfloat16', 'scale_layout': 'row-major', 'modules': modules, |
| 'embedding_and_lm_head': 'unquantized, tied', |
| 'preserved': ['admission networks', 'small recurrent gates', 'normalization', 'history capacities and update rules'], |
| } |
| (output/'nvfp4_config.json').write_text(json.dumps(manifest, indent=2)+'\n') |
| packed = {} |
| for shard in sorted(output.glob('model*.safetensors')): |
| with safe_open(shard, framework='pt', device='cpu') as handle: |
| for name, shape in modules.items(): |
| if name+'.weight' not in handle.keys(): |
| continue |
| weight = handle.get_tensor(name+'.weight') |
| scale = handle.get_tensor(name+'.weight_scale') |
| global_scale = handle.get_tensor(name+'.weight_scale_2') |
| assert weight.dtype == torch.uint8 and list(weight.shape) == [shape[0], shape[1]//2], name |
| assert scale.dtype == torch.float8_e4m3fn and scale.numel() == shape[0]*shape[1]//16, name |
| assert global_scale.dtype == torch.float32 and global_scale.numel() == 1, name |
| assert torch.isfinite(scale.float()).all() and global_scale.item() > 0, name |
| packed[name] = {'shape': shape, 'weight_bytes': weight.numel(), 'scale_bytes': scale.numel()+4} |
| assert set(packed) == set(modules) |
| receipt = { |
| 'format': 'W4A16_NVFP4', 'quantized_modules': len(packed), 'modules': packed, |
| 'conversion_seconds': time.perf_counter()-started, |
| 'weight_bytes': sum(path.stat().st_size for path in output.glob('model*.safetensors')), |
| 'gpu': torch.cuda.get_device_name(0), 'modelopt_version': '0.47.0', |
| } |
| (output/'conversion.json').write_text(json.dumps(receipt, indent=2)+'\n') |
| print(json.dumps({'stage': 'conversion_complete', 'modules': len(packed), 'weight_bytes': receipt['weight_bytes']}), flush=True) |
|
|
