Function fp8_rowwise_quantize

tensorrt_llm/quantization/quantize.py:249–321

(model, quant_config: QuantConfig)

def fp8_rowwise_quantize(model, quant_config: QuantConfig):
    assert quant_config.quant_mode.has_fp8_rowwise()

    quant_cls_map = {
        RmsNorm: Fp8RowwiseRmsNorm,
        LayerNorm: Fp8RowwiseLayerNorm,
        GatedMLP: Fp8RowwiseGatedMLP,
        MLP: Fp8RowwiseMLP,
        Attention: Fp8RowwiseAttention,
    }

    exclude_modules = quant_config.exclude_modules
    if exclude_modules is None:
        exclude_modules = []
    # Always exclude these modules for FP8 rowwise
    exclude_modules = list(
        set(exclude_modules + ['*ln_f', '*ln_embed', '*lm_head']))

    def extract_layer_idx(name):
        ss = name.split('.')
        for s in ss:
            if s.isdigit():
                return int(s)
        return None

    # Meta's LLaMA 3.1 recipe:
    # (1) Skip quantization for the first and last Transformer layers
    # (2) Skip quantization for the Attention layers
    if quant_config.use_meta_recipe:
        exclude_modules.extend(['*input_layernorm', '*attention'])

    for name, layer, parent in model.named_modules_with_parent():
        module_name = name.rsplit('.', 1)[-1]

        if quant_config.use_meta_recipe:
            local_layer_idx = extract_layer_idx(name)
            mapping = model.config.mapping
            layers_range = mapping.pp_layers(model.config.num_hidden_layers)
            if mapping.is_first_pp_rank() and local_layer_idx == 0:
                continue
            if mapping.is_last_pp_rank(
            ) and local_layer_idx == len(layers_range) - 1:
                continue

        quant_cls = None
        for cls in quant_cls_map:
            if isinstance(layer, cls):
                quant_cls = quant_cls_map[cls]
                break
        if quant_cls is None:
            continue

        is_excluded = False
        for exclude_module in exclude_modules:
            if fnmatch.fnmatchcase(name, exclude_module):
                is_excluded = True
                break
        if is_excluded: