--- inference: false license: bigscience-bloom-rail-1.0 language: - ak - ar - as - bm - bn - ca - en - es - eu - fon - fr - gu - hi - id - ig - ki - kn - lg - ln - ml - mr - ne - nso - ny - or - pa - pt - rn - rw - sn - st - sw - ta - te - tn - ts - tum - tw - ur - vi - wo - xh - yo - zh - zu pipeline_tag: text-generation --- ### Quantized bigscience/bloom 1B3 with 8-bit weights Heavily inspired by [Hivemind's GPT-J-6B with 8-bit weights](https://huggingface.co/hivemind/gpt-j-6B-8bit), this is a version of [bigscience/bloom](https://huggingface.co/bigscience/bloom-1b3) a ~1 billion parameters language model that you run and fine-tune with less memory. Here, we also apply [LoRA (Low Rank Adaptation)](https://arxiv.org/abs/2106.09685) to reduce model size. ### How to fine-tune TBA ### How to use This model can be used by adapting Bloom original implementation. This is an adaptation from [Hivemind's GPT-J 8-bit](https://nbviewer.org/urls/huggingface.co/hivemind/gpt-j-6B-8bit/raw/main/convert-gpt-j.ipynb): ```python import transformers import torch import torch.nn as nn import torch.nn.functional as F from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise from typing import Tuple from torch.cuda.amp import custom_fwd, custom_bwd class FrozenBNBLinear(nn.Module): def __init__(self, weight, absmax, code, bias=None): assert isinstance(bias, nn.Parameter) or bias is None super().__init__() self.out_features, self.in_features = weight.shape self.register_buffer("weight", weight.requires_grad_(False)) self.register_buffer("absmax", absmax.requires_grad_(False)) self.register_buffer("code", code.requires_grad_(False)) self.adapter = None self.bias = bias def forward(self, input): output = DequantizeAndLinear.apply(input, self.weight, self.absmax, self.code, self.bias) if self.adapter: output += self.adapter(input) return output @classmethod def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear": weights_int8, state = quantize_blockise_lowmemory(linear.weight) return cls(weights_int8, *state, linear.bias) def __repr__(self): return f"{self.__class__.__name__}({self.in_features}, {self.out_features})" class DequantizeAndLinear(torch.autograd.Function): @staticmethod @custom_fwd def forward(ctx, input: torch.Tensor, weights_quantized: torch.ByteTensor, absmax: torch.FloatTensor, code: torch.FloatTensor, bias: torch.FloatTensor): weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) ctx.save_for_backward(input, weights_quantized, absmax, code) ctx._has_bias = bias is not None return F.linear(input, weights_deq, bias) @staticmethod @custom_bwd def backward(ctx, grad_output: torch.Tensor): assert not ctx.needs_input_grad[1] and not ctx.needs_input_grad[2] and not ctx.needs_input_grad[3] input, weights_quantized, absmax, code = ctx.saved_tensors # grad_output: [*batch, out_features] weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) grad_input = grad_output @ weights_deq grad_bias = grad_output.flatten(0, -2).sum(dim=0) if ctx._has_bias else None return grad_input, None, None, None, grad_bias class FrozenBNBEmbedding(nn.Module): def __init__(self, weight, absmax, code): super().__init__() self.num_embeddings, self.embedding_dim = weight.shape self.register_buffer("weight", weight.requires_grad_(False)) self.register_buffer("absmax", absmax.requires_grad_(False)) self.register_buffer("code", code.requires_grad_(False)) self.adapter = None def forward(self, input, **kwargs): with torch.no_grad(): # note: both quantuized weights and input indices are *not* differentiable weight_deq = dequantize_blockwise(self.weight, absmax=self.absmax, code=self.code) output = F.embedding(input, weight_deq, **kwargs) if self.adapter: output += self.adapter(input) return output @classmethod def from_embedding(cls, embedding: nn.Embedding) -> "FrozenBNBEmbedding": weights_int8, state = quantize_blockise_lowmemory(embedding.weight) return cls(weights_int8, *state) def __repr__(self): return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})" def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2 ** 20): assert chunk_size % 4096 == 0 code = None chunks = [] absmaxes = [] flat_tensor = matrix.view(-1) for i in range((matrix.numel() - 1) // chunk_size + 1): input_chunk = flat_tensor[i * chunk_size: (i + 1) * chunk_size].clone() quantized_chunk, (absmax_chunk, code) = quantize_blockwise(input_chunk, code=code) chunks.append(quantized_chunk) absmaxes.append(absmax_chunk) matrix_i8 = torch.cat(chunks).reshape_as(matrix) absmax = torch.cat(absmaxes) return matrix_i8, (absmax, code) def convert_to_int8(model): """Convert linear and embedding modules to 8-bit with optional adapters""" for module in list(model.modules()): for name, child in module.named_children(): if isinstance(child, nn.Linear): print(name, child) setattr( module, name, FrozenBNBLinear( weight=torch.zeros(child.out_features, child.in_features, dtype=torch.uint8), absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), code=torch.zeros(256), bias=child.bias, ), ) elif isinstance(child, nn.Embedding): setattr( module, name, FrozenBNBEmbedding( weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8), absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), code=torch.zeros(256), ) ) class BloomBlock(transformers.models.bloom.modeling_bloom.BloomBlock): def __init__(self, config, layer_number=None): super().__init__(config, layer_number) convert_to_int8(self.self_attention) convert_to_int8(self.mlp) class BloomModel(transformers.models.bloom.modeling_bloom.BloomModel): def __init__(self, config): super().__init__(config) convert_to_int8(self) class BloomForCausalLM(transformers.models.bloom.modeling_bloom.BloomForCausalLM): def __init__(self, config): super().__init__(config) convert_to_int8(self) transformers.models.bloom.modeling_bloom.BloomBlock = BloomBlock model_name = 'mrm8488/bloom-1b3-8bit' model = BloomForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True) tokenizer = BloomTokenizerFast.from_pretrained(model_name) prompt = tokenizer("Given a table named salaries and columns id, created_at, salary, age. Creates a SQL to answer What is the average salary for 22 years old:", return_tensors='pt') out = model.generate(**prompt, min_length=10, do_sample=True) tokenizer.decode(out[0]) ```