Function quantize

lora/convert.py:13–34

(weights, config, args)

def quantize(weights, config, args):
    quantized_config = copy.deepcopy(config)

    # Load the model:
    model = models.Model(models.ModelArgs.from_dict(config))
    model.load_weights(list(weights.items()))

    # Quantize the model:
    nn.quantize(
        model,
        args.q_group_size,
        args.q_bits,
    )

    # Update the config:
    quantized_config["quantization"] = {
        "group_size": args.q_group_size,
        "bits": args.q_bits,
    }
    quantized_weights = dict(tree_flatten(model.parameters()))

    return quantized_weights, quantized_config


if __name__ == "__main__":