Function quantize

SwissArmyTransformer/sat/quantization/kernels.py:237–285

Replace fp16 linear with quantized linear

(model, weight_bit_width)

def quantize(model, weight_bit_width):
    """Replace fp16 linear with quantized linear"""

    print_rank0(f"> Quantizing model weight to {weight_bit_width} bits")

    quantized_cnt = [0] # [] for tracing in closure
    
    def replace_linear(module):
        for name, sub_module in module.named_children():
            if isinstance(sub_module, ColumnParallelLinear) and not isinstance(sub_module, QuantizedColumnParallelLinear):
                setattr(module, name, QuantizedColumnParallelLinear(
                    weight_bit_width=weight_bit_width,
                    weight=sub_module.weight.to(torch.cuda.current_device()),
                    input_size=sub_module.input_size,
                    output_size=sub_module.output_size,
                    bias=hasattr(sub_module, 'bias') and isinstance(sub_module.bias, Parameter),
                    gather_output=sub_module.gather_output,
                    params_dtype=torch.half,
                    name=name,
                    skip_init=True,
                    device=sub_module.weight.device,
                    bias_val=getattr(sub_module, 'bias', None),
                    )
                )
                quantized_cnt[0] += sub_module.weight.numel()

            elif isinstance(sub_module, RowParallelLinear) and not isinstance(sub_module, QuantizedRowParallelLinear):
                setattr(module, name, QuantizedRowParallelLinear(
                    weight_bit_width=weight_bit_width,
                    weight=sub_module.weight.to(torch.cuda.current_device()),
                    input_size=sub_module.input_size,
                    output_size=sub_module.output_size,
                    bias=hasattr(sub_module, 'bias') and isinstance(sub_module.bias, Parameter),
                    input_is_parallel=sub_module.input_is_parallel,
                    params_dtype=torch.half,
                    name=name,
                    skip_init=True,
                    device=sub_module.weight.device,
                    bias_val=getattr(sub_module, 'bias', None),
                    )
                )
                quantized_cnt[0] += sub_module.weight.numel()
            else:
                replace_linear(sub_module)

    replace_linear(model)
    print_rank0(f"> Quantized {quantized_cnt[0]} parameters in total.")

    return model