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

[Bhuvanesh09]

Issue in the past code:

• Unable to use SmoothQuant with custom calibration dataset in the current code of convert_checkpoint.py. This is because:
  • load_calib_dataset function in tensorrt_llm/models/convert_utils.py has the default value of split and key set to be None. In the past, this function would work only for cnn_dailymail, lambada. But with better default values of "train" split and "text" key, this function would be able to load any dataset configured correctly.
• The number of samples for calibration in smoothquant was hard coded to 512 samples and max sequence length of 512 which lead to errors and incorrect calibration with custom calibration datasets.
  • This PR fixes the above with small changes in the files :
    • tensorrt_llm/models/llama/convert.py.
    • examples/llama/convert_checkpoint.py.
    • tensorrt_llm/models/llama/model.py.

Results:

With the above changes, to the release v0.12.0. I was able to check that the calibration works well with much better quality of quantized model in comparison to default calibration dataset. For SmoothQuant especially, when per_token is set to true, it is important that the calibration sequence length matches the distribution of the samples which are used during deployment/production.

[Bhuvanesh09]

@Barry-Delaney : Kindly take a look. Thanks in advance!

[Bhuvanesh09]

Any help needed here @Barry-Delaney ?

[Barry-Delaney]

Thanks for the PR @Bhuvanesh09!
Overall, the change looks good to me. Left some minor comments.
We will integrate the calibration process into examples/quantization/quantize.py in the future, and the args for calib_size and the usage of customized dataset is already implemented there.
For now, we will merge it first, and thanks for your contribution!

[Barry-Delaney]

The change looks good to me. We'll merge your changes into internal code base.
Thanks for the contribution!