Decode float weights from quantization codes. Differentiable. :param codes: tensor of integer quantization codes, shape [*dims, num_out_group

Update outs_tensor with new activations and optionally compute sample-wise mse loss with previous activations :param layer: transformer layer

Parallel version of update_outs_and_compute_losses; works on lists of input/output tensors

Compute the activation MSE error between transformer layers :param

Compute MSE in parallel over multiple GPUs, each GPU processes a portion of samples

Calculate the effective number of bits per original matrix parameters

optimized batch k-means for 1d datapoint using sort :param groupwise_data: stuff to be compressed, shape: [num_groups, group_size] :param

Get quantization codebooks or reconstruct them from second level quantization (see codebook_values_nbits)

Get a non view to codes, regardless of how codes are stored

mocks model launch to collect inputs to the first model layer :returns: a list of torch tensors with activations for each device in args.devi

Loads and prepares data for a Transformers model. Args: name (str): The name of the dataset to load. This can be one of 'wiki

find transformer block classes that should be wrapped with inner FullyShardedDataParallel (auto_wrap_policy)

Return True iff an integer dtype is signed

Iterate over (name, QuantizedWeight) pairs for all quantized weights trained by this optimizer and rank

A hacky way to substitute an already registered parameter with a non-parameter tensor. Breaks future use.

:param sharded_tensors_by_name: a dictionary from string to flat (1d) tensors available on the current shard :note: the keys should be the sa

Compute the activation MSE error = ||X @ quantized_weight - X @ reference_weight||^2 Use the square-of-difference formula to avoid ma

Given a matrix delta [out_features, in_features], compute a tensor [num_output_groups, num_input_groups] that contains the squared su

update _flat_codes [num_groups, num_codebooks] to approximate _flat_reference [num_groups, group_size]

Update quantized_weight codes in-place via beam search

Compute token-wise KL divergence loss without materializing all logits/logprobs simultaneously :param student_hidden_states: input hidden sta

Get per-channel or per-group quantization scales or reconstruct those scales based on scales_nbits

Samples data points *forever*, in random order, with less overhead than DataLoader; Adapted from https://github.com/stanis-morozov/unq/blob/m

Hacky way to allow compatibility between old *pickled* layers and new transformers

Save model, optimizer state dict and training metadata to be loaded via load_training_state

Update codes[selection] to new_codes, regardless of their dtype and whether they are wrapped as storage

Wrap a model *ForCausalLM components: transformer and lm_head are wrapped as FSDP instances

collect full (possibly optimizer-updated) dequantized weights

collect full parameter gradients from fsdp-sharded parameters, return dict[name -> grad]

Select top-:beam_size: and reorder beam accordingly, return new beam :param beam_codes: a tensor with best weight codes, shape: [beam_size, n

Compute MSE or sum-of-squared-error losses for all possible ways to replace quantization codes for one input group and one codebook. Works i

:param reference: [num_groups, group_size] :param codes: [num_groups, num_codebooks] :param codebooks: [num_codebooks, codebook_size, gro

Compute per-channel squared error between X @ weight_or_weights and X @ reference_weight :param XTX: pairwise products of input features matm

Compute MSE in parallel over output channels

Decode float weights from quantization codes. Differentiable. :param codes: tensor of integer quantization codes, shape [*dims, num_out_group

For each parameter name, return a tuple of sizes (numbers of elements) this parameter across all FSDP ranks

:param reference: [num_groups, group_size] :param codebook: [codebook_size, group_size] :param chunk_size_values: how many values can be

Get initial clusters by iteratively choosing a vector that is the farthest from already selected clusters

Prepare auxiliary data structures for quickly copying parameters to replicas for data-parallel training.

Update discrete state representations to approximate straight through buffers

Ensure that every optimized parameter receives gradient

Choose which version of parameter to optimize: the parameter itself or a straight-through buffer

Assign dequantized weight buffers to latest quantized weights after codebook/scale/code updates

Accumulate a minibatch of layer inputs and update the X.T @ X (aka half hessian)

Update own codes in-place via beam search so as to minimize squared errors. Return the updated codes. :param reference_weight: origin

Create a version of the model where all QuanizedWeight and derivative layers are de-quantized and cast to dtype. :param model: model to be de

Create optimizer for training only continuous parameters of a quantized model

Create optimizer for PV-Tuning using a de-quantized student model and a dictionary of quantized weights

Perplexity evaluation as per https://github.com/IST-DASLab/gptq (standard among quantization research)

Find nearest clusters for each batch of data and return their indices

Fine-tune a module with pre-quantized linear layers so as to minimize MSE between layer-wise inps/outs :param layer: a trainable module wher

:param data: [nsamples, dim] :param k: number of centroids :param max_iter: run at most this many iterations :param gpu: if True, run

:param data: [nsamples, dim] :param k: number of centroids :param max_iter: run at most this many iterations :param check_every: chec