Function blockwise_quantize

convert.py:262–275

(weight: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype)

  return out

def blockwise_quantize(weight: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
  assert weight.ndim == block_size.numel() and weight.ndim == 2
  scale_shape = torch.Size((torch.tensor(list(weight.shape)) / block_size).ceil().long())
  scale = torch.empty(scale_shape, dtype=torch.float32)
  out = torch.empty_like(weight, dtype=dtype)
  for i in range(scale.shape[0]):
    for j in range(scale.shape[1]):
      block_size_i = block_size[0]
      block_size_j = block_size[1]
      w_block = weight[i*block_size_i:(i+1)*block_size_i, j*block_size_j:(j+1)*block_size_j]
      qw_block, scale_block = per_tensor_quantize(w_block, dtype)
      out[i*block_size_i:(i+1)*block_size_i, j*block_size_j:(j+1)*block_size_j] = qw_block
      scale[i, j] = scale_block
  return out, scale

def per_expert_blockwise_quantize(expert_weights: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
  assert expert_weights.ndim == 3