diff --git a/examples/onnx_ptq/torch_quant_to_onnx.py b/examples/onnx_ptq/torch_quant_to_onnx.py index 4c33b4c1a53..a89497f6782 100644 --- a/examples/onnx_ptq/torch_quant_to_onnx.py +++ b/examples/onnx_ptq/torch_quant_to_onnx.py @@ -19,6 +19,7 @@ import timm import torch import torch.multiprocessing as mp +import torch.nn.functional as F from datasets import load_dataset from download_example_onnx import export_to_onnx from evaluation import evaluate @@ -26,11 +27,12 @@ import modelopt.torch.quantization as mtq """ -This script is used to quantize a timm model using dynamic quantization like MXFP8 or NVFP4. +This script is used to quantize a timm model using dynamic quantization like MXFP8 or NVFP4, +or using auto quantization for optimal per-layer quantization. The script will: 1. Given the model name, create a timm torch model. -2. Quantize the torch model in MXFP8 or NVFP4 mode. +2. Quantize the torch model in MXFP8, NVFP4, INT4_AWQ, or AUTO mode. 3. Export the quantized torch model to ONNX format. """ @@ -55,8 +57,17 @@ def filter_func(name): return pattern.match(name) is not None -def load_calibration_data(model_name, data_size, batch_size, device): - """Load and prepare calibration data.""" +def load_calibration_data(model_name, data_size, batch_size, device, with_labels=False): + """Load and prepare calibration data. + + Args: + model_name: Name of the timm model + data_size: Number of samples to load + batch_size: Batch size for data loader + device: Device to load data to + with_labels: If True, return dict with 'image' and 'label' keys (for auto_quantize) + If False, return just the images (for standard quantize) + """ dataset = load_dataset("zh-plus/tiny-imagenet") model = timm.create_model(model_name, pretrained=True, num_classes=1000) data_config = timm.data.resolve_model_data_config(model) @@ -65,9 +76,18 @@ def load_calibration_data(model_name, data_size, batch_size, device): images = dataset["train"][:data_size]["image"] calib_tensor = [transforms(img) for img in images] calib_tensor = [t.to(device) for t in calib_tensor] - return torch.utils.data.DataLoader( - calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4 - ) + + if with_labels: + labels = dataset["train"][:data_size]["label"] + labels = torch.tensor(labels, device=device) + calib_dataset = [{"image": img, "label": lbl} for img, lbl in zip(calib_tensor, labels)] + return torch.utils.data.DataLoader( + calib_dataset, batch_size=batch_size, shuffle=True, num_workers=4 + ) + else: + return torch.utils.data.DataLoader( + calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4 + ) def quantize_model(model, config, data_loader=None): @@ -86,16 +106,80 @@ def forward_loop(model): return quantized_model -def get_model_input_shape(model_name, batch_size): +def forward_step(model, batch): + """Forward step function for auto_quantize scoring.""" + return model(batch["image"]) + + +def loss_func(output, batch): + """Loss function for auto_quantize gradient computation.""" + return F.cross_entropy(output, batch["label"]) + + +def auto_quantize_model( + model, + data_loader, + quantization_formats, + effective_bits=4.8, + num_calib_steps=512, + num_score_steps=128, +): + """Auto-quantize the model using optimal per-layer quantization search. + + Args: + model: PyTorch model to quantize + data_loader: DataLoader with image-label dict batches + quantization_formats: List of quantization format config names or dicts + effective_bits: Target effective bits constraint + num_calib_steps: Number of calibration steps + num_score_steps: Number of scoring steps for sensitivity analysis + + Returns: + Tuple of (quantized_model, search_state_dict) + """ + constraints = {"effective_bits": effective_bits} + + # Convert string format names to actual config objects + format_configs = [] + for fmt in quantization_formats: + if isinstance(fmt, str): + format_configs.append(getattr(mtq, fmt)) + else: + format_configs.append(fmt) + + print(f"Starting auto-quantization search with {len(format_configs)} formats...") + print(f"Effective bits constraint: {effective_bits}") + print(f"Calibration steps: {num_calib_steps}, Scoring steps: {num_score_steps}") + + quantized_model, search_state = mtq.auto_quantize( + model, + constraints=constraints, + quantization_formats=format_configs, + data_loader=data_loader, + forward_step=forward_step, + loss_func=loss_func, + num_calib_steps=num_calib_steps, + num_score_steps=num_score_steps, + verbose=True, + ) + + # Disable quantization for specified layers + mtq.disable_quantizer(quantized_model, filter_func) + + return quantized_model, search_state + + +def get_model_input_shape(model): """Get the input shape from timm model configuration.""" - model = timm.create_model(model_name, pretrained=True, num_classes=1000) data_config = timm.data.resolve_model_data_config(model) input_size = data_config["input_size"] - return (batch_size, *tuple(input_size)) # Add batch dimension + return tuple(input_size) def main(): - parser = argparse.ArgumentParser(description="Quantize timm models to MXFP8 or NVFP4") + parser = argparse.ArgumentParser( + description="Quantize timm models to FP8, MXFP8, INT8, NVFP4, INT4_AWQ, or use AUTO quantization" + ) # Model hyperparameters parser.add_argument( @@ -106,14 +190,14 @@ def main(): ) parser.add_argument( "--quantize_mode", - choices=["fp8", "mxfp8", "int8", "nvfp4", "int4_awq"], + choices=["fp8", "mxfp8", "int8", "nvfp4", "int4_awq", "auto"], default="mxfp8", - help="Type of quantization to apply (mxfp8, nvfp4, int4_awq)", + help="Type of quantization to apply. Default is MXFP8.", ) parser.add_argument( "--onnx_save_path", required=True, - help="The path to save the ONNX model.", + help="The save path to save the ONNX model.", type=str, ) parser.add_argument( @@ -140,15 +224,43 @@ def main(): help="Number of samples to use for evaluation. If None, use entire validation set.", ) - args = parser.parse_args() + # Auto quantization specific arguments + parser.add_argument( + "--auto_quantization_formats", + nargs="+", + choices=[ + "NVFP4_AWQ_LITE_CFG", + "FP8_DEFAULT_CFG", + "MXFP8_DEFAULT_CFG", + "INT8_DEFAULT_CFG", + "INT4_AWQ_CFG", + ], + default=["NVFP4_AWQ_LITE_CFG", "FP8_DEFAULT_CFG"], + help="Quantization formats to search from for auto mode (e.g., NVFP4_AWQ_LITE_CFG FP8_DEFAULT_CFG)", + ) + parser.add_argument( + "--effective_bits", + type=float, + default=4.8, + help="Target effective bits for auto quantization constraint. Default is 4.8.", + ) + parser.add_argument( + "--num_score_steps", + type=int, + default=128, + help="Number of scoring steps for auto quantization. Default is 128.", + ) - # Get input shape from model config - input_shape = get_model_input_shape(args.timm_model_name, args.batch_size) + args = parser.parse_args() # Create model and move to appropriate device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = timm.create_model(args.timm_model_name, pretrained=True, num_classes=1000).to(device) + # Get input shape from model config + input_size = get_model_input_shape(model) + input_shape = (args.batch_size, *input_size) + # Evaluate base model if requested if args.evaluate: print("\n=== Evaluating Base Model ===") @@ -159,21 +271,44 @@ def main(): ) print(f"Base Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%") - # Select quantization config - config = QUANT_CONFIG_DICT[args.quantize_mode] - data_loader = ( - None - if args.quantize_mode == "mxfp8" - else load_calibration_data( + # Quantize model based on mode + if args.quantize_mode == "auto": + # Auto quantization requires labels for loss computation + data_loader = load_calibration_data( args.timm_model_name, args.calibration_data_size, - input_shape[0], # batch size + args.batch_size, device, + with_labels=True, ) - ) - # Quantize model - quantized_model = quantize_model(model, config, data_loader) + quantized_model, _ = auto_quantize_model( + model, + data_loader, + args.auto_quantization_formats, + args.effective_bits, + args.calibration_data_size, + args.num_score_steps, + ) + else: + # Standard quantization - only load calibration data if needed + config = QUANT_CONFIG_DICT[args.quantize_mode] + if args.quantize_mode == "mxfp8": + data_loader = None + else: + data_loader = load_calibration_data( + args.timm_model_name, + args.calibration_data_size, + args.batch_size, + device, + with_labels=False, + ) + + quantized_model = quantize_model(model, config, data_loader) + + # Print quantization summary + print("\nQuantization Summary:") + mtq.print_quant_summary(quantized_model) # Evaluate quantized model if requested if args.evaluate: @@ -188,8 +323,10 @@ def main(): ) print(f"Quantized Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%") - if args.quantize_mode in ["fp8", "int8"]: - print(f"Exporting to {args.quantize_mode} ONNX model is not supported yet.") + if args.quantize_mode in ["fp8", "int8", "auto"]: + print( + f"The selected quantization mode {args.quantize_mode} is not supported for ONNX export yet." + ) return # Export to ONNX