def quantize_model_(
    model,
    size_tracker,
    layers_to_quantize,
    block_sizes_config,
    n_centroids_config,
    step=0,
    n_iter=15,
    eps=1e-6,
    max_tentatives=100,
    remove_weights=False,
    verbose=True,
    state_dict=None,
):
    """
    Quantize a model in-place by stages. All the targeted
    layers are replaced by their quantized counterpart,
    and the model is ready for the finetuning of the
    centroids in a standard training loop (no modifications
    required). Note that we do not quantize biases.

    Args:
        - model: a nn.Module
        - size_tracker: useful for tracking quatization statistics
        - layers_to_quantize: a list containing regexps for
          filtering the layers to quantize at each stage according
          to their name (as in model.named_parameters())
        - block_sizes_config: dict like
          {
              'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
              'Linear': ('in_features', {'*': 8})
          }
          For instance, all conv2d layers with kernel size 3x3 have
          a block size of 9 and all Linear layers are quantized with
          a block size of 8, irrespective of their size.
        - n_centroids_config: dict like
          {
              'Conv2d': ('kernel_size', {'*': 256}),
              'Linear': ('in_features', {'*': 256})
          }
          For instance, all conv2d layers are quantized with 256 centroids
        - step: the layers to quantize inplace corresponding
          to layers_to_quantize[step]
    """

    quantized_layers = get_layers(
        model, layers_to_quantize[step], remove_weights=remove_weights
    )

    for layer in quantized_layers:

        # book-keeping
        is_master_process = (not dist.is_initialized()) or (
            dist.is_initialized() and dist.get_rank() == 0
        )
        verbose = verbose and is_master_process

        # get block size and centroids
        module = attrgetter(layer)(model)
        block_size = get_param(module, layer, block_sizes_config)
        n_centroids = get_param(module, layer, n_centroids_config)
        if verbose:
            logging.info(
                f"Quantizing layer {layer} with block size {block_size} and {n_centroids} centroids"
            )

        # quantize layer
        weight = module.weight.data.clone()
        is_bias = "bias" in [x[0] for x in module.named_parameters()]
        bias = module.bias.data.clone() if is_bias else None
        quantizer = PQ(
            weight,
            block_size,
            n_centroids=n_centroids,
            n_iter=n_iter,
            eps=eps,
            max_tentatives=max_tentatives,
            verbose=verbose,
        )

        # quantization performed on all GPUs with same seed
        quantizer.encode()
        centroids = quantizer.centroids.contiguous()
        assignments = quantizer.assignments.contiguous()

        # If n_iter = 0 and state_dict is provided, then
        # we initialize random assignments and centroids to
        # random values of the appropriate dimensions
        # because the quantized model parameters will
        # overwritten by the state_dict later on.
        if n_iter == 0 and state_dict:
            # Initialize random centroids of the correct size
            centroids = torch.rand(centroids.size())
            centroids.cuda()
            # Get counts and assignment keys from layer in loaded checkpoint.
            counts_key = layer + "." + "counts"
            assignment_key = layer + "." + "assignments"
            # Get number of different bins to include.
            counts = list(state_dict[counts_key].shape)[0]
            print(layer)
            print(state_dict[counts_key])
            print(counts)
            # Initialize random assignments of the correct size
            # with an appropriate number of bins.
            num_assignments = list(state_dict[assignment_key].shape)[0]
            num_extra = num_assignments - counts
            print(num_assignments)
            print(num_extra)
            assignments_bins = torch.arange(counts)
            assignments_rand = torch.randint(0, counts - 1, (num_extra,))
            assignments = torch.cat((assignments_bins, assignments_rand), 0)
            # assignments = assignments.type(torch.IntTensor)
            assignments.cuda()
            print("assignments")
            print(assignments)

        # broadcast results to make sure weights are up-to-date
        if dist.is_initialized():
            dist.broadcast(centroids, 0)
            dist.broadcast(assignments, 0)

        # instantiate the quantized counterpart
        if isinstance(module, nn.Linear):
            out_features, in_features = map(
                lambda k: module.__dict__[k], ["out_features", "in_features"]
            )
            quantized_module = PQLinear(
                centroids, assignments, bias, in_features, out_features
            )
        elif isinstance(module, nn.Embedding):
            num_embeddings, embedding_dim = map(
                lambda k: module.__dict__[k], ["num_embeddings", "embedding_dim"]
            )
            quantized_module = PQEmbedding(
                centroids, assignments, num_embeddings, embedding_dim
            )
        elif isinstance(module, nn.Conv2d):
            out_channels, in_channels, kernel_size = map(
                lambda k: module.__dict__[k],
                ["out_channels", "in_channels", "kernel_size"],
            )
            stride, padding, dilation, groups, padding_mode = map(
                lambda k: module.__dict__[k],
                ["stride", "padding", "dilation", "groups", "padding_mode"],
            )

            quantized_module = PQConv2d(
                centroids,
                assignments,
                bias,
                in_channels,
                out_channels,
                kernel_size,
                stride=stride,
                padding=padding,
                dilation=dilation,
                groups=groups,
                padding_mode=padding_mode,
            )
        else:
            raise ValueError(f"Module {module} not yet supported for quantization")

        # replace layer by its quantized counterpart
        attrsetter(layer)(model, quantized_module)

        # update statistics
        size_tracker.update(weight, block_size, n_centroids)

    # return name of quantized layers
    return quantized_layers

    Quantize a model in-place by stages. All the targeted
    layers are replaced by their quantized counterpart,
    and the model is ready for the finetuning of the
    centroids in a standard training loop (no modifications
    required). Note that we do not quantize biases.

    Args:
        - model: a nn.Module
        - size_tracker: useful for tracking quatization statistics
        - layers_to_quantize: a list containing regexps for
          filtering the layers to quantize at each stage according
          to their name (as in model.named_parameters())
        - block_sizes_config: dict like
          {
              'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
              'Linear': ('in_features', {'*': 8})
          }
          For instance, all conv2d layers with kernel size 3x3 have
          a block size of 9 and all Linear layers are quantized with
          a block size of 8, irrespective of their size.
        - n_centroids_config: dict like
          {
              'Conv2d': ('kernel_size', {'*': 256}),
              'Linear': ('in_features', {'*': 256})
          }
          For instance, all conv2d layers are quantized with 256 centroids
        - step: the layers to quantize inplace corresponding
          to layers_to_quantize[step]
    
understand_icl/fairseq/fairseq/modules/quantization/pq/utils.py