Speed up quantization in generate.py (OpenGVLab#35)

Co-authored-by: Carlos Mocholí <[email protected]>

generate.py

@@ -1,4 +1,3 @@
-import os
 import sys
 import time
 from pathlib import Path
@@ -7,8 +6,7 @@
 import lightning as L
 import torch
 
-from lit_llama.model import LLaMA
-from lit_llama.tokenizer import Tokenizer
+from lit_llama import LLaMA, Tokenizer, as_8_bit_quantized
 
 
 @torch.no_grad()
@@ -104,21 +102,13 @@ def main(
 
  fabric = L.Fabric(accelerator=accelerator, devices=1)
 
- if quantize:
- from lit_llama.quantization import quantize
-
- print("Running quantization. This may take a minute ...")
- # TODO: Initializing the model directly on the device does not work with quantization
+ with as_8_bit_quantized(fabric.device, enabled=quantize):
+ print("Loading model ...", file=sys.stderr)
+ t0 = time.time()
  model = LLaMA.from_name(model_size)
- # The output layer can be sensitive to quantization, we keep it in default precision
- model = quantize(model, skip=("lm_head", "output"))
  checkpoint = torch.load(checkpoint_path)
  model.load_state_dict(checkpoint)
- else:
- with fabric.device:
- model = LLaMA.from_name(model_size)
- checkpoint = torch.load(checkpoint_path)
- model.load_state_dict(checkpoint)
+ print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
 
  model.eval()
 
@@ -133,6 +123,7 @@ def main(
 
  L.seed_everything(1234)
  t0 = time.time()
+
  for _ in range(num_samples):
  y = generate(
  model,
@@ -144,8 +135,9 @@ def main(
  )[0] # unpack batch dimension
  print(tokenizer.decode(y))
 
- print(f"Time for inference: {time.time() - t0:.02f} seconds", file=sys.stderr)
- print(f"Memory used (GB): {torch.cuda.max_memory_reserved() / 1e9:.02f}", file=sys.stderr)
+ t = time.time() - t0
+ print(f"\n\nTime for inference: {t:.02f} sec total, {max_new_tokens / t:.02f} tokens/sec", file=sys.stderr)
+ print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr)
 
 
 if __name__ == "__main__":

lit_llama/__init__.py

@@ -1,2 +1,3 @@
 from lit_llama.model import LLaMAConfig, LLaMA, RMSNorm, build_rope_cache, apply_rope
+from lit_llama.quantization import as_8_bit_quantized
 from lit_llama.tokenizer import Tokenizer

lit_llama/model.py

@@ -184,9 +184,6 @@ def __init__(self, config: LLaMAConfig) -> None:
  )
  )
 
- # init all weights
- self.apply(self._init_weights)
-
  def _init_weights(self, module: nn.Module) -> None:
  if isinstance(module, nn.Linear):
  torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))

lit_llama/quantization.py

@@ -1,19 +1,73 @@
 import os
-from typing import Tuple
+from contextlib import contextmanager
+import warnings
 
-import torch.nn as nn
+import torch
 
+# configuration for bitsandbytes before import
 os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+warnings.filterwarnings(
+ "ignore", 
+ message="MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization"
+)
+warnings.filterwarnings(
+ "ignore", 
+ message="The installed version of bitsandbytes was compiled without GPU support. 8-bit optimizers and GPU quantization are unavailable."
+)
 import bitsandbytes as bnb # noqa: E402
 
 
-def quantize(model: nn.Module, threshold: float = 6.0, skip: Tuple[str, ...] = ()) -> nn.Module:
- for name, module in model.named_children():
- if isinstance(module, nn.Linear) and name not in skip:
- model._modules[name] = bnb.nn.Linear8bitLt(
- module.in_features, module.out_features, bias=module.bias, has_fp16_weights=False, threshold=threshold
- )
+class Linear8bitLt(bnb.nn.Linear8bitLt):
+ """Wraps `bnb.nn.Linear8bitLt` and enables instantiation directly on the device and
+ re-quantizaton when loading the state dict.
+ 
+ 
+ This should only be used for inference. For training, use `bnb.nn.Linear8bitLt` directly.
+ """
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs, has_fp16_weights=False, threshold=6.0)
+ # We quantize the initial weight here so we don't end up filling the device
+ # memory with float32 weights which could lead to OOM.
+ self._quantize_weight(self.weight.data)
 
- if module.children():
- quantize(module, threshold=threshold, skip=skip)
- return model
+ def _load_from_state_dict(self, local_state_dict, *args, **kwargs):
+ # There is only one key that ends with `*.weight`, the other one is the bias
+ weight_key = next(name for name in local_state_dict.keys() if name.endswith("weight"))
+
+ # Load the weight from the state dict and re-quantize it
+ weight = local_state_dict.pop(weight_key)
+ self._quantize_weight(weight)
+
+ # If there is a bias, let nn.Module load it
+ if local_state_dict:
+ super()._load_from_state_dict(local_state_dict, *args, **kwargs)
+
+ def _quantize_weight(self, weight: torch.Tensor) -> None:
+ # This code is taken and adapted from `bnb.nn.Int8Params.cuda()`
+ B = weight.contiguous().half().cuda()
+ CB, CBt, SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B)
+ del CBt
+ del SCBt
+ self.weight.data = CB
+ setattr(self.weight, "CB", CB)
+ setattr(self.weight, "SCB", SCB)
+
+
+@contextmanager
+def as_8_bit_quantized(device: torch.device, enabled: bool = True):
+ """A context manager under which you can instantiate the model with 8-bit quantized tensors
+ being created directly on the given device.
+ """
+
+ with torch.device(device):
+ if not enabled:
+ yield
+ return
+
+ if device.type != "cuda":
+ raise ValueError("Quantization is only supported on the GPU.")
+
+ torch_linear_cls = torch.nn.Linear
+ torch.nn.Linear = Linear8bitLt
+ yield
+ torch.nn.Linear = torch_linear_cls

tests/test_model.py

@@ -52,6 +52,7 @@ def test_to_orig_llama(lit_llama, orig_llama) -> None:
  )
 
  llama_model = lit_llama.LLaMA(llama_config)
+ llama_model.apply(llama_model._init_weights)
  orig_llama_model = orig_llama.Transformer(orig_llama_config)
 
  copy_weights(llama_model, orig_llama_model)