58 lines
2.6 KiB
Python
58 lines
2.6 KiB
Python
# Copyright (c) Kyutai, all rights reserved.
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""Quantization based on bitsandbytes, supporting only 8 bits for now.
|
|
We are taking from freedom from the intended use of bnb:
|
|
"""
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class QLinear(nn.Module):
|
|
def __init__(self, linear: nn.Linear):
|
|
super().__init__()
|
|
from bitsandbytes import functional as bnbF # type: ignore
|
|
weight = linear.weight
|
|
assert weight.data.dtype.is_floating_point
|
|
assert linear.bias is None
|
|
CB, SCB, _ = bnbF.int8_vectorwise_quant(weight.data.to(torch.float16)) # type: ignore
|
|
self.weight = nn.Parameter(CB, requires_grad=False)
|
|
self.weight_scb = nn.Parameter(SCB, requires_grad=False)
|
|
|
|
def forward(self, x):
|
|
import bitsandbytes as bnb # type: ignore
|
|
state = bnb.MatmulLtState()
|
|
state.CB = self.weight # type: ignore
|
|
assert isinstance(state.CB, torch.Tensor)
|
|
state.SCB = self.weight_scb # type: ignore
|
|
assert isinstance(state.SCB, torch.Tensor)
|
|
if state.SCB.dtype != torch.float:
|
|
raise RuntimeError(
|
|
"Expected `weight_scb` to have type float, but got bfloat16. "
|
|
"When using quantized models, care should be taken not to change the dtype of "
|
|
"the model once initialized.")
|
|
assert state.SCB.dtype == torch.float, state.SCB.dtype
|
|
state.has_fp16_weights = False
|
|
y = bnb.matmul(x.half(), state.CB, state=state)
|
|
assert isinstance(y, torch.Tensor)
|
|
return y
|
|
|
|
|
|
def replace_linear_with_qlinear(module):
|
|
"""Recursively replace all Linear layers with QLinear layers."""
|
|
for name, child in module.named_children():
|
|
if isinstance(child, nn.Linear):
|
|
setattr(module, name, QLinear(child))
|
|
elif isinstance(child, QLinear):
|
|
# Slight issue with the way we implement things: the scale param
|
|
# might get casted with the rest of the model to bfloat16, altough
|
|
# we most likely want to keep it as float. For the LM model we might call this function twice,
|
|
# first layer by layer to avoid to big of a memory usage, and second, at the end
|
|
# of the LM init, after all other modules are initialized and properly dtyped.
|
|
# In any case that should happen before loading the state dict to avoid a loss of precision.
|
|
child.float()
|
|
else:
|
|
replace_linear_with_qlinear(child)
|