fix: add support for passing calib sequence length, and num samples + fixing use of custom calibration dataset for smoothquant in llama

examples/llama/convert_checkpoint.py

@@ -91,6 +91,18 @@ def parse_arguments():
         help=
         "The huggingface dataset name or the local directory of the dataset for calibration."
     )
+    parser.add_argument(
+        "--calib_size",
+        type=int,
+        default=512,
+        help="Number of datapoints to use for calibration. Default is 512. Set to -1 to use the whole dataset.",
+    )
+    parser.add_argument(
+        "--calib_max_seq_length",
+        type=int,
+        default=512,
+        help="Max Sequence length to use for calibration. Default is 512.",
+    )
     parser.add_argument(
         "--smoothquant",
         "-sq",
@@ -402,6 +414,8 @@ def convert_and_save_hf(args):
             quant_config=quant_config,
             device='cpu' if args.load_model_on_cpu else 'cuda',
             calib_dataset=args.calib_dataset,
+            calib_batches=args.calib_size,
+            calib_max_seq_length=args.calib_max_seq_length,
             **override_fields)
     else:
         # When not loading by shard, preload one complete model and then slice per rank weights from this

tensorrt_llm/models/convert_utils.py

@@ -254,8 +254,8 @@ def has_safetensors(model_dir: str):
 
 def load_calib_dataset(dataset_name_or_dir: str,
                        config_name: Optional[str] = None,
-                       split: Optional[str] = None,
-                       key: Optional[str] = None,
+                       split: Optional[str] = "train", # default split value in hf datasets object
+                       key: Optional[str] = "text", # default key value in hf datasets object
                        trust_remote_code=True,
                        **kwargs):
     if config_name is None:

tensorrt_llm/models/llama/convert.py

@@ -1084,7 +1084,10 @@ def quantize(hf_model_dir: str,
              output_dir: str,
              config: LLaMAConfig,
              device: str = 'cuda',
-             calib_dataset: str = 'cnn_dailymail'):
+             calib_dataset: str = 'cnn_dailymail',
+             calib_batches: int = 512,
+             calib_max_seq_length: int = 512,
+             ):
     '''
         Quantize the save the model as TRT-LLM checkpoint to output_dir
     '''
@@ -1118,7 +1121,10 @@ def quantize(hf_model_dir: str,
 
     dataset = load_calib_dataset(calib_dataset)
 
-    act_range = capture_activation_range(hf_model, tokenizer, dataset)
+    if calib_batches == -1: # use the whole dataset if calib_batches is -1
+        calib_batches = len(dataset)
+
+    act_range = capture_activation_range(hf_model, tokenizer, dataset, num_samples=calib_batches, seq_len=calib_max_seq_length)
     qkv_para, smoother = {}, {}
     if use_smooth_quant:
         smooth_llama_model(hf_model, act_range, quant_config.smoothquant_val,

tensorrt_llm/models/llama/model.py

@@ -446,7 +446,9 @@ def quantize(
                              output_dir,
                              config=config,
                              device=device,
-                             calib_dataset=calib_dataset)
+                             calib_dataset=calib_dataset,
+                             calib_batches=calib_batches,
+                             calib_max_seq_length=calib_max_seq_length,)
         else:
             raise ValueError(
                 f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead."