Update llama/model.py

llama/model.py

@@ -1,20 +1,18 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # This software may be used and distributed according to the terms of the GNU General Public License version 3.
 
-from typing import Optional, Tuple
+from contextvars import ContextVar
+
+from typing import Optional, Tuple, Type
 from dataclasses import dataclass
 import math
 
 import torch
 from torch import nn
 import torch.nn.functional as F
+import bitsandbytes as bnb
 
-import fairscale.nn.model_parallel.initialize as fs_init
-from fairscale.nn.model_parallel.layers import (
-    ParallelEmbedding,
-    RowParallelLinear,
-    ColumnParallelLinear,
-)
+import tqdm
 
 
 @dataclass
@@ -73,40 +71,57 @@ def apply_rotary_emb(
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
 
+class UninitializedLinear(nn.Linear):
+    def reset_parameters(self) -> None:
+        pass
+
+
+class InferenceQuantizedLinear(bnb.nn.Linear8bitLt):
+    def __init__(self, *args, **kwargs):
+        super().__init__(has_fp16_weights=False, *args, **kwargs)
+
+    def reset_parameters(self) -> None:
+        pass
+
+
+default_quantize: ContextVar[bool] = ContextVar("default_quantize", default=False)
+
+
+def get_linear_class() -> Type[nn.Linear]:
+    if default_quantize.get():
+        return InferenceQuantizedLinear
+    return UninitializedLinear
+
+
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
 
-        self.n_local_heads = args.n_heads // fs_init.get_model_parallel_world_size()
+        self.n_local_heads = (
+            args.n_heads // 1
+        )  # fs_init.get_model_parallel_world_size()
         self.head_dim = args.dim // args.n_heads
 
-        self.wq = ColumnParallelLinear(
+        Linear = get_linear_class()
+        self.wq = Linear(
             args.dim,
             args.n_heads * self.head_dim,
             bias=False,
-            gather_output=False,
-            init_method=lambda x: x,
         )
-        self.wk = ColumnParallelLinear(
+        self.wk = Linear(
             args.dim,
             args.n_heads * self.head_dim,
             bias=False,
-            gather_output=False,
-            init_method=lambda x: x,
         )
-        self.wv = ColumnParallelLinear(
+        self.wv = Linear(
             args.dim,
             args.n_heads * self.head_dim,
             bias=False,
-            gather_output=False,
-            init_method=lambda x: x,
         )
-        self.wo = RowParallelLinear(
-            args.n_heads * self.head_dim,
+        self.wo = Linear(
             args.dim,
+            args.n_heads * self.head_dim,
             bias=False,
-            input_is_parallel=True,
-            init_method=lambda x: x,
         )
 
         self.cache_k = torch.zeros(
@@ -116,7 +131,13 @@ class Attention(nn.Module):
             (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
         ).cuda()
 
-    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs_cis: torch.Tensor,
+        mask: Optional[torch.Tensor],
+    ):
         bsz, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
@@ -143,9 +164,7 @@ class Attention(nn.Module):
             scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
         scores = F.softmax(scores.float(), dim=-1).type_as(xq)
         output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
-        output = output.transpose(
-            1, 2
-        ).contiguous().view(bsz, seqlen, -1)
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
 
         return self.wo(output)
 
@@ -161,14 +180,17 @@ class FeedForward(nn.Module):
         hidden_dim = int(2 * hidden_dim / 3)
         hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
 
-        self.w1 = ColumnParallelLinear(
-            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
-        )
-        self.w2 = RowParallelLinear(
-            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
+        Linear = get_linear_class()
+        self.w1 = Linear(dim, hidden_dim, bias=False)
+        self.w2 = Linear(
+            hidden_dim,
+            dim,
+            bias=False,
         )
-        self.w3 = ColumnParallelLinear(
-            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
+        self.w3 = Linear(
+            dim,
+            hidden_dim,
+            bias=False,
         )
 
     def forward(self, x):
@@ -189,12 +211,36 @@ class TransformerBlock(nn.Module):
         self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
         self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
 
-    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
-        h = x + self.attention.forward(self.attention_norm(x), start_pos, freqs_cis, mask)
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs_cis: torch.Tensor,
+        mask: Optional[torch.Tensor],
+    ):
+        h = x + self.attention.forward(
+            self.attention_norm(x), start_pos, freqs_cis, mask
+        )
         out = h + self.feed_forward.forward(self.ffn_norm(h))
         return out
 
 
+def convert_linear_to_bnb(float_linear):
+    new_layer = InferenceQuantizedLinear(
+        float_linear.in_features,
+        float_linear.out_features,
+        bias=float_linear.bias is not None,
+    )
+    new_layer._parameters["weight"] = bnb.nn.Int8Params(
+        float_linear.weight.data.cpu(),
+        requires_grad=False,
+        has_fp16_weights=False,
+    )
+    if float_linear.bias is not None:
+        new_layer._parameters["bias"] = float_linear.bias
+    return new_layer
+
+
 class Transformer(nn.Module):
     def __init__(self, params: ModelArgs):
         super().__init__()
@@ -202,18 +248,16 @@ class Transformer(nn.Module):
         self.vocab_size = params.vocab_size
         self.n_layers = params.n_layers
 
-        self.tok_embeddings = ParallelEmbedding(
-            params.vocab_size, params.dim, init_method=lambda x: x
-        )
+        self.tok_embeddings = torch.nn.Embedding(params.vocab_size, params.dim)
 
         self.layers = torch.nn.ModuleList()
         for layer_id in range(params.n_layers):
             self.layers.append(TransformerBlock(layer_id, params))
 
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
-        self.output = ColumnParallelLinear(
-            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
-        )
+
+        Linear = get_linear_class()
+        self.output = Linear(params.dim, params.vocab_size, bias=False)
 
         self.freqs_cis = precompute_freqs_cis(
             self.params.dim // self.params.n_heads, self.params.max_seq_len * 2
@@ -228,11 +272,39 @@ class Transformer(nn.Module):
 
         mask = None
         if seqlen > 1:
-            mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=tokens.device)
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=tokens.device
+            )
             mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
 
         for layer in self.layers:
             h = layer(h, start_pos, freqs_cis, mask)
         h = self.norm(h)
         output = self.output(h[:, -1, :])  # only compute last logits
-        return output.float()
+        return output.float()
+
+    def quantize(self):
+        # https://github.com/pytorch/vision/issues/2391#issuecomment-653900218
+        def get_layer(model, name):
+            layer = model
+            for attr in name.split("."):
+                layer = getattr(layer, attr)
+            return layer
+
+        def set_layer(model, name, layer):
+            try:
+                attrs, name = name.rsplit(".", 1)
+                model = get_layer(model, attrs)
+            except ValueError:
+                pass
+            setattr(model, name, layer)
+
+        linear_layers = {
+            k: v for k, v in self.named_modules() if isinstance(v, nn.Linear)
+        }
+
+        print("Quantizing", len(linear_layers), "layers")
+        for name, layer in tqdm.tqdm(linear_layers.items()):
+            new_layer = convert_linear_to_bnb(layer)
+            set_layer(self, name, new_layer)
+        self.cuda()