Upload GPTJLoraForCausalLM

config.json

@@ -1,6 +1,7 @@
 {
+  "_name_or_path": "gpt-j-6b-8bit-lora",
   "activation_function": "gelu_new",
-  "add_adapters": true,
+  "add_apapters": true,
   "architectures": [
     "GPTJLoraForCausalLM"
   ],
@@ -10,6 +11,7 @@
     "AutoModelForCausalLM": "gptj.GPTJLoraForCausalLM"
   },
   "bos_token_id": 50256,
+  "eight_bit": true,
   "embd_pdrop": 0.0,
   "eos_token_id": 50256,
   "gradient_checkpointing": false,
@@ -39,7 +41,8 @@
   },
   "tie_word_embeddings": false,
   "tokenizer_class": "GPT2Tokenizer",
-  "transformers_version": "4.20.1",
+  "torch_dtype": "float32",
+  "transformers_version": "4.24.0",
   "use_cache": true,
   "vocab_size": 50400
 }

config.py

@@ -1,10 +1,10 @@
-from transformers import GPTJConfig
-
-
-class GPTJLoraConfig(GPTJConfig):
-    model_type = "gptj-lora"
-
-    def __init__(self, add_adapters=False, **kwargs):
-        self.add_apapters = add_adapters
-        super().__init__(**kwargs)
-        self.model_type = "gptj-lora"
+from transformers import GPTJConfig
+
+
+class GPTJLoraConfig(GPTJConfig):
+    model_type = "gptj-lora"
+
+    def __init__(self, add_adapters=False, **kwargs):
+        self.add_apapters = add_adapters
+        super().__init__(**kwargs)
+        self.model_type = "gptj-lora"

gptj.py

@@ -1,86 +1,86 @@
-import torch
-from torch import nn
-from .lora import FrozenBNBLinear, FrozenBNBEmbedding
-from .config import GPTJLoraConfig
-import transformers
-
-
-def add_adapters(model, adapter_dim=16):
-    assert adapter_dim > 0
-
-    for module in model.modules():
-        if isinstance(module, FrozenBNBLinear):
-            module.adapter = nn.Sequential(
-                nn.Linear(module.in_features, adapter_dim, bias=False),
-                nn.Linear(adapter_dim, module.out_features, bias=False),
-            )
-            nn.init.zeros_(module.adapter[1].weight)
-        elif isinstance(module, FrozenBNBEmbedding):
-            module.adapter = nn.Sequential(
-                nn.Embedding(module.num_embeddings, adapter_dim),
-                nn.Linear(adapter_dim, module.embedding_dim, bias=False),
-            )
-            nn.init.zeros_(module.adapter[1].weight)
-
-
-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):
-                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 GPTJLoraBlock(transformers.models.gptj.modeling_gptj.GPTJBlock):
-    config_class = GPTJLoraConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.config_class = GPTJLoraConfig
-
-        convert_to_int8(self.attn)
-        convert_to_int8(self.mlp)
-
-
-class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel):
-    config_class = GPTJLoraConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.config_class = GPTJLoraConfig
-
-        convert_to_int8(self)
-
-
-class GPTJLoraForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM):
-    config_class = GPTJLoraConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.config_class = GPTJLoraConfig
-
-        convert_to_int8(self)
-        if config.add_apapters:
-            add_adapters(self)
-
-
-transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJLoraBlock  # monkey-patch GPT-J
+import torch
+from torch import nn
+from .lora import FrozenBNBLinear, FrozenBNBEmbedding
+from .config import GPTJLoraConfig
+import transformers
+
+
+def add_adapters(model, adapter_dim=16):
+    assert adapter_dim > 0
+
+    for module in model.modules():
+        if isinstance(module, FrozenBNBLinear):
+            module.adapter = nn.Sequential(
+                nn.Linear(module.in_features, adapter_dim, bias=False),
+                nn.Linear(adapter_dim, module.out_features, bias=False),
+            )
+            nn.init.zeros_(module.adapter[1].weight)
+        elif isinstance(module, FrozenBNBEmbedding):
+            module.adapter = nn.Sequential(
+                nn.Embedding(module.num_embeddings, adapter_dim),
+                nn.Linear(adapter_dim, module.embedding_dim, bias=False),
+            )
+            nn.init.zeros_(module.adapter[1].weight)
+
+
+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):
+                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 GPTJLoraBlock(transformers.models.gptj.modeling_gptj.GPTJBlock):
+    config_class = GPTJLoraConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config_class = GPTJLoraConfig
+
+        convert_to_int8(self.attn)
+        convert_to_int8(self.mlp)
+
+
+class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel):
+    config_class = GPTJLoraConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config_class = GPTJLoraConfig
+
+        convert_to_int8(self)
+
+
+class GPTJLoraForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM):
+    config_class = GPTJLoraConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config_class = GPTJLoraConfig
+
+        convert_to_int8(self)
+        if config.add_apapters:
+            add_adapters(self)
+
+
+transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJLoraBlock  # monkey-patch GPT-J

lora.py

@@ -1,99 +1,99 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torch.cuda.amp import custom_fwd, custom_bwd
-
-from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise
-
-
-def quantize_blockwise_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)
-
-
-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).clone()
-        if self.adapter:
-            output += self.adapter(input)
-        return output
-
-    @classmethod
-    def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear":
-        weights_int8, state = quantize_blockwise_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 quantized 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_blockwise_lowmemory(embedding.weight)
-        return cls(weights_int8, *state)
-
-    def __repr__(self):
-        return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})"
-
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.cuda.amp import custom_fwd, custom_bwd
+
+from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise
+
+
+def quantize_blockwise_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)
+
+
+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).clone()
+        if self.adapter:
+            output += self.adapter(input)
+        return output
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear":
+        weights_int8, state = quantize_blockwise_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 quantized 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_blockwise_lowmemory(embedding.weight)
+        return cls(weights_int8, *state)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})"
+

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:153cb853074d3fb66c18f93b78297f0d88e252eb1f2a2e5779dff97453a63124
+oid sha256:10793d174ead92956a981a490ea62ebd2d2109ed944f8fb2fa2815e987988449
 size 6316410080