quant_utils#

Quantization helpers used only by the TensorRT-LLM checkpoint export path.

Functions

get_scaling_factor_from_weight

Calculate the weight scaling factor for a given group size.

resmooth_and_get_scale

Resmooths weights from a single or multiple ranks and get scaling factors and amax.
get_scaling_factor_from_weight(weight, group_size)#

Calculate the weight scaling factor for a given group size.

Return type:

tensor

resmooth_and_get_scale(merged_weights, pre_quant_scales, ranks, group_size, new_pre_quant_scale=None, quantization=None)#

Resmooths weights from a single or multiple ranks and get scaling factors and amax.

Parameters:

  • merged_weights (Tensor) – Merged weights from ranks.

  • pre_quant_scales (list[Tensor]) – List of pre-quantization scales for each rank.

  • ranks (int) – Number of ranks.

  • group_size (int) – Group size of the quantization block.

  • new_pre_quant_scale (optional) – If not provided, weights will be resmoothed using the average of pre_quant_scales.

  • quantization (str | None)

Returns:

Resmoothed weights. weight_scaling_factors: Resmoothed scaling factors. avg_pre_quant_scale: Calculated average of the quantization scale.

Return type:

weights