From 492e44c0f6893e1b03425e7424ad5a91193d2273 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 7 May 2026 02:04:54 +0000 Subject: [PATCH] Fix dataloader API: max_sample_length not seq_len, proper create_forward_loop --- quantize_modelopt.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/quantize_modelopt.py b/quantize_modelopt.py index 8da3a96..55b3271 100644 --- a/quantize_modelopt.py +++ b/quantize_modelopt.py @@ -108,22 +108,23 @@ def main(): # Build calibration dataloader print("Building calibration dataset...") - calib_dataloader = create_forward_loop( - model, - dataloader=get_dataloader( - tokenizer=tokenizer, - calib_size=args.calib_size, - batch_size=args.batch_size, - calib_seq=args.calib_seq, - ), + calib_dataloader = get_dataloader( + tokenizer=tokenizer, + calib_size=args.calib_size, + batch_size=args.batch_size, + calib_seq=args.calib_seq, ) + # Create forward loop from dataloader (modelopt helper) + from modelopt.torch.utils.dataset_utils import create_forward_loop + forward_loop = create_forward_loop(calib_dataloader) + # Quantize quant_cfg = QUANT_CONFIGS[args.qformat] print(f"Running PTQ with {args.qformat}...") t0 = time.time() - model = mtq.quantize(model, quant_cfg, calib_dataloader) + model = mtq.quantize(model, quant_cfg, forward_loop) elapsed = time.time() - t0 print(f"Quantization complete in {elapsed/60:.1f} min") @@ -149,7 +150,7 @@ def get_dataloader(tokenizer, calib_size, batch_size, calib_seq): tokenizer=tokenizer, num_samples=calib_size[0], batch_size=batch_size, - seq_len=calib_seq, + max_sample_length=calib_seq, )