deleted leftovers

inference_kernels/router.py

@@ -1,29 +0,0 @@
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from src.inference_kernels.triton_kernel import aqlm_gemm_stupid as triton_gemm
-from src.utils import _dequantize_weight, unpack_int_data
-
-
-def forward_pass_quantized_linear(
-    input: torch.Tensor,
-    codes: torch.IntTensor,
-    codebooks: torch.Tensor,
-    scales: torch.Tensor,
-    bias: Optional[torch.Tensor],
-) -> torch.Tensor:
-    if input.is_cuda:
-        matmul_result = triton_gemm(input, codes, codebooks, scales)
-        if bias is not None:
-            matmul_result += bias
-        return matmul_result
-    else:
-        dequantized_weight = _dequantize_weight(
-            unpack_int_data(codes, codebooks.shape[0].bit_length() - 1),
-            codebooks,
-            scales,
-        )
-        return F.linear(input, dequantized_weight, bias)

inference_kernels/triton_kernel.py

@@ -1,170 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-from torch.autograd import Function
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({"UNUSED": 1}, num_stages=num_stages, num_warps=num_warps)
-        for num_stages in (1, 2, 3, 4, 5)
-        for num_warps in (1, 2, 4, 8)
-    ],
-    key=[
-        "in_features",
-        "out_features",
-        "num_codebooks",
-        "codebook_size",
-        "out_group_size",
-        "in_group_size",
-        "num_input_groups",
-        "num_input_groups_next_power_of_2",
-        "compute_in_fp32",
-    ],
-)
-@triton.jit
-def _aqlm_gemv_simple(
-    input_vec_ptr,
-    output_vec_ptr,
-    codes_i16_ptr,
-    codebooks_ptr,
-    scales_ptr,
-    in_features: tl.constexpr,
-    out_features: tl.constexpr,
-    num_codebooks: tl.constexpr,
-    codebook_size: tl.constexpr,
-    out_group_size: tl.constexpr,
-    in_group_size: tl.constexpr,
-    num_input_groups: tl.constexpr,
-    num_input_groups_next_power_of_2: tl.constexpr,
-    compute_in_fp32: tl.constexpr,
-    UNUSED: tl.constexpr,
-):
-    # variables ending with "_i" mean "for i-th output unit"
-    pid = tl.program_id(axis=0)  # [0, 1, ... {out_features-1}]
-
-    # Stage 1: load input data
-    input_vec = tl.load(
-        input_vec_ptr
-        + tl.arange(0, num_input_groups_next_power_of_2)[:, None, None] * in_group_size
-        + tl.arange(0, in_group_size)[None, None, :],
-        mask=tl.arange(0, num_input_groups_next_power_of_2)[:, None, None] < num_input_groups,
-    )
-    # [in_features//in_group_size, 1, group_size]
-    # Note: we could simply load input_vec then reshape
-    #     input_vec = tl.load(input_vec_ptr + tl.arange(0, in_features))  # [in_features]
-    #     input_vec = tl.view(input_vec, [num_input_groups, 1, in_group_size])
-    #     , but this does not work because tl.view may reorder elements arbitrarily; see its docstring
-
-    # Stage 2: load integer codes for the active row
-    # [in_features // in_group_size, num_codebooks]
-    codes_i_ptrs = (
-        codes_i16_ptr
-        + pid * num_input_groups * num_codebooks
-        + tl.arange(0, num_input_groups_next_power_of_2)[:, None] * num_codebooks
-        + tl.arange(0, num_codebooks)[None, :]
-    )
-    codes_i_mask_1d = tl.arange(0, num_input_groups_next_power_of_2) < num_input_groups
-
-    codes_i = tl.load(codes_i_ptrs, mask=codes_i_mask_1d[:, None])  # [in_features//in_group_size, num_codebooks]
-    if codes_i.dtype == tl.int16:
-        codes_i = codes_i.to(tl.int32)
-        codes_i = (codes_i) + (codes_i < 0) * codebook_size  # aka 2 ** nbits_per_codebook
-        # ^-- (because codes are int16 tensors that contain uint data)
-
-        # The following alternative does not work:
-        #     codes_i = codes_i.to(tl.int32) % codebook_size # aka 2 ** nbits_per_codebook
-    else:
-        codes_i = codes_i.to(tl.int32)
-
-    # shift codes_i so that codebooks after 0th point to correct indices in codebooks_ptr
-    codes_i += tl.arange(0, num_codebooks)[None, :] * codebook_size  # aka 2 ** nbits_per_codebook
-    # ^-- [in_group_size, num_codebooks]
-
-    # Stage 3: convert codes to pointers to every individual (activated) weight in codebooks
-    # [in_features // in_group_size, num_codebooks, out_group_size, in_group_size]
-    out_group_ix = tl.arange(0, out_group_size)[None, None, :, None]
-    in_group_ix = tl.arange(0, in_group_size)[None, None, None, :]
-    weight_i_ptrs = (
-        codebooks_ptr
-        + codes_i[:, :, None, None] * out_group_size * in_group_size
-        + out_group_ix * in_group_size
-        + in_group_ix
-    )
-
-    # Stage 4: reconstruct weights, multiply by inputs and write out
-    weights_i = tl.load(weight_i_ptrs, mask=codes_i_mask_1d[:, None, None, None], other=0)
-    if compute_in_fp32:
-        weights_i = weights_i.to(tl.float32)
-        input_vec = input_vec.to(tl.float32)
-    # ^-- [in_features // in_group_size, num_codebooks, out_group_size, in_group_size]
-    weights_i = tl.sum(weights_i, axis=1)  # sum codebooks as per additive quantization
-    # ^-- [in_features // in_group_size, out_group_size, in_group_size]
-
-    if out_group_size == 1:
-        scale = tl.load(scales_ptr + pid).to(weights_i.dtype)  # scalar
-        output_i = tl.sum(weights_i * input_vec) * scale
-        tl.store(output_vec_ptr + pid, output_i.to(input_vec.dtype))
-    else:
-        output_i = tl.sum(tl.sum(weights_i * input_vec, axis=2), axis=0)  # [out_group_size]
-        output_i *= tl.load(scales_ptr + pid).to(weights_i.dtype)
-        tl.store(output_vec_ptr + pid * out_group_size + tl.arange(0, out_group_size), output_i.to(input_vec.dtype))
-
-
-def next_power_of_2(x):
-    return 1 if x == 0 else 2 ** (x - 1).bit_length()
-
-
-def aqlm_gemv_simple(
-    input_vec: torch.Tensor,
-    codes_i16: torch.ShortTensor,
-    codebooks: torch.Tensor,
-    scales: torch.Tensor,
-    compute_in_fp32: bool = True,
-):
-
-    device, dtype = codebooks.device, codebooks.dtype
-    num_codebooks, codebook_size, out_group_size, in_group_size = codebooks.shape
-    in_features = input_vec.shape[1]
-    out_features = codes_i16.shape[0] * out_group_size
-    num_input_groups = codes_i16.shape[1]
-    assert input_vec.ndim == 2 and input_vec.shape[0] == 1, "do reshape; now!"
-    assert scales.shape == (out_features // out_group_size, 1, 1, 1)
-    assert in_features % in_group_size == 0
-    assert codebooks.shape[1] == 2**16
-
-    output_vec = torch.empty(1, out_features, device=device, dtype=dtype)
-    # 1D launch kernel where each block computes output unit
-    grid = lambda META: (out_features // out_group_size,)
-    _aqlm_gemv_simple[grid](
-        input_vec,
-        output_vec,
-        codes_i16,
-        codebooks,
-        scales,
-        in_features,
-        out_features,
-        num_codebooks,
-        codebook_size,
-        out_group_size,
-        in_group_size,
-        num_input_groups,
-        next_power_of_2(num_input_groups),
-        compute_in_fp32,
-    )
-
-    return output_vec
-
-
-def aqlm_gemm_stupid(
-    input: torch.Tensor,
-    codes_i16: torch.ShortTensor,
-    codebooks: torch.Tensor,
-    scales: torch.Tensor,
-    compute_in_fp32: bool = True,
-):
-    original_shape = input.shape
-    input = input.reshape(-1, original_shape[-1])
-    return torch.cat(
-        [aqlm_gemv_simple(input_vec.unsqueeze(0), codes_i16, codebooks, scales, compute_in_fp32) for input_vec in input]
-    ).reshape(original_shape[:-1] + (-1,))

utils.py

@@ -1,159 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import functools
-import os
-from typing import Callable, Iterator, Optional, Sequence
-
-import torch
-import torch.nn.functional as F
-
-ellipsis = type(...)
-
-
-def get_mean_nbits_by_codebook(codes: torch.IntTensor, huffman_group_size: int = 2):
-
-    """
-    Calculates average code length in codebooks.
-    :param codes: codebook codes
-    :param huffman_group_size: huffman compresssion dimension count
-    """
-    import huffman
-
-    _, codebook_size, num_codebooks = codes.shape
-    flat_codes_by_codebook = codes.permute(2, 0, 1).flatten(1, 2)
-    code_counts = torch.zeros(
-        num_codebooks, codebook_size, device=flat_codes_by_codebook.device, dtype=flat_codes_by_codebook.dtype
-    ).scatter_add(
-        -1, flat_codes_by_codebook, torch.ones_like(flat_codes_by_codebook)
-    )  # shape: [current beam_size, num_codebooks, codebook_size], initial beam_size = 1
-    code_probs = code_counts / code_counts.sum(dim=-1, keepdim=True).float()
-    code_probs = code_probs.cpu().numpy()
-    assert num_codebooks % huffman_group_size == 0
-
-    mean_code_lengths = []
-    for group_index in range(num_codebooks // huffman_group_size):
-        group_code_probs = {(): 1}
-
-        for codebook_index in range(group_index * huffman_group_size, (group_index + 1) * huffman_group_size):
-            new_group_code_probs = {}
-            for group, group_prob in group_code_probs.items():
-                for code, code_prob in tuple(enumerate(code_probs[codebook_index])):
-                    new_group_code_probs[group + (code,)] = group_prob * code_prob
-            group_code_probs = new_group_code_probs
-
-        huffman_codebook_i = huffman.codebook(list(group_code_probs.items()))
-        codebook_mean_code_length_i = sum(
-            len(huffman_codebook_i[code]) * prob for code, prob in group_code_probs.items()
-        )
-        mean_code_lengths.append(codebook_mean_code_length_i)
-    return mean_code_lengths
-
-
-def get_int_dtype(nbits: int) -> torch.dtype:
-    if nbits <= 8:
-        return torch.int8
-    if nbits <= 16:
-        return torch.int16
-    if nbits <= 32:
-        return torch.int32
-    if nbits <= 64:
-        return torch.int64
-    raise ValueError(f"No dtype available for {nbits}-bit codebooks")
-
-
-@torch.inference_mode()
-def pack_int_data(data: torch.IntTensor, nbits: int) -> torch.IntTensor:
-    data[data >= 2 ** (nbits - 1)] -= 2**nbits
-    return data.to(get_int_dtype(nbits))
-
-
-@torch.inference_mode()
-def unpack_int_data(data: torch.IntTensor, nbits: int) -> torch.IntTensor:
-    return data.to(torch.int64) % (2**nbits)
-
-
-@functools.lru_cache()
-def maybe_script(fn: callable) -> callable:
-    """Apply torch.jit.script to function unless one is using TPU. TPU does not support torch.jit.script."""
-    using_tpu = bool(os.environ.get("TPU_NAME"))
-    # this is a reserved variable that must be set to TPU address (e.g. grpc://11.22.33.44:1337) for TPU to function
-    should_script = int(os.environ.get("AQ_USE_JIT", not using_tpu))
-    return torch.jit.script(fn) if should_script else fn
-
-
-@contextlib.contextmanager
-def using_tf32(enabled: bool):
-    was_cudnn = torch.backends.cudnn.allow_tf32
-    was_matmul = torch.backends.cuda.matmul.allow_tf32
-    torch.backends.cudnn.allow_tf32 = enabled
-    torch.backends.cuda.matmul.allow_tf32 = enabled
-    yield
-    torch.backends.cudnn.allow_tf32 = was_cudnn
-    torch.backends.cuda.matmul.allow_tf32 = was_matmul
-
-
-def iterate_minibatches(
-    *tensors: torch.Tensor,
-    batch_size: int,
-    allow_incomplete: bool = True,
-    device: Optional[torch.device] = None,
-    callback: Callable[[Sequence[torch.Tensor]], Sequence[torch.Tensor]] = lambda x: x,
-) -> Iterator[Sequence[torch.Tensor]]:
-    """
-    Samples data points *forever*, in random order, with less overhead than DataLoader;
-    Adapted from https://github.com/stanis-morozov/unq/blob/master/lib/utils.py
-    probably implemented over9000 times in transformers, torch, etc
-    :param tensors: one or more tensors with the same 0-th dimension
-    :param batch_size: sample this many points with each yield
-    :param allow_incomplete: if True and if dataset size is not divisible by batch size, the last batch
-        may have less than :batch_size: samples to cover the entire dataset. If False, the last batch is dropped
-    :param callback: optional function to be called on each batch of tensors before it is yielded to the user
-    :returns: generates a tuple of minibatches from each tensor, same length as input *tensors
-        If a batch contains only one tensor, this function will yield a tensor (and not a tuple/list with one tensor)
-    """
-    num_samples = len(tensors[0])
-    assert all(len(x) == num_samples for x in tensors)
-    indices = torch.randperm(num_samples, device=tensors[0].device)
-    while True:
-        prev_batch = None
-        for batch_start in range(0, len(indices), batch_size):
-            if not allow_incomplete and batch_start + batch_size > len(indices):
-                break
-            batch_ix = indices[batch_start : batch_start + batch_size]
-            batch = callback(tuple(tensor[batch_ix].to(device, non_blocking=True) for tensor in tensors))
-            if prev_batch is not None:
-                yield prev_batch
-            prev_batch = batch if isinstance(batch, (list, tuple)) and len(tensors) > 1 else batch[0]
-            del batch
-        yield prev_batch
-
-
-@maybe_script
-def _dequantize_weight(
-    codes: torch.Tensor, codebooks: torch.Tensor, scales: Optional[torch.Tensor] = None
-) -> torch.Tensor:
-    """
-    Decode float weights from quantization codes. Differentiable.
-    :param codes: tensor of integer quantization codes, shape [*dims, num_out_groups, num_in_groups, num_codebooks]
-    :param codebooks: tensor of vectors for each quantization code, [num_codebooks, codebook_size, out_group_size, in_group_size]
-    :param scales: weight will be multiplied by this factor, must be broadcastble with [*dims, out_groups, num_in_groups, out_group_size, in_group_size]
-    :return: reconstructed weight tensor of shape [*dims, num_in_groups*group_size]
-    """
-    num_out_groups, num_in_groups, num_codebooks = codes.shape[-3:]
-    num_codebooks, codebook_size, out_group_size, in_group_size = codebooks.shape
-    out_features = num_out_groups * out_group_size
-    in_features = num_in_groups * in_group_size
-    codebook_offsets = torch.arange(
-        0, num_codebooks * codebook_size, codebook_size, device=codes.device
-    )  # shape: [num_codebooks]
-    reconstructed_weight_flat = F.embedding_bag(
-        codes.flatten(0, -2) + codebook_offsets, codebooks.flatten(0, 1).flatten(-2, -1), mode="sum"
-    )  # [prod(dims) * num_out_groups * num_in_groups, out_group_size * in_group_size]
-
-    reconstructed_weight_groupwise = reconstructed_weight_flat.view(
-        list(codes.shape[:-3]) + [num_out_groups, num_in_groups, out_group_size, in_group_size]
-    )
-    if scales is not None:
-        reconstructed_weight_groupwise = reconstructed_weight_groupwise.mul(scales)
-    return reconstructed_weight_groupwise.swapaxes(-3, -2).reshape(list(codes.shape[:-3]) + [out_features, in_features])