Update gen.py

gen.py

@@ -1,15 +1,13 @@
 from typing import Tuple
 
 import os
+import torch
 import time
 import json
 from pathlib import Path
 
-import torch
-from fairscale.nn.model_parallel.initialize import initialize_model_parallel
-from llama.generation import LLaMA
-from llama.model import ModelArgs, Transformer
-from llama.tokenizer import Tokenizer
+os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+from llama import ModelArgs, Transformer, Tokenizer, LLaMA, default_quantize
 
 from google.cloud import storage
 
@@ -18,18 +16,6 @@ bucket_name = os.environ.get("GCS_BUCKET")
 llama_weight_path = "weights/llama"
 tokenizer_weight_path = "weights/tokenizer"
 
-def setup_model_parallel() -> Tuple[int, int]:
-    local_rank = int(os.environ.get("LOCAL_RANK", -1))
-    world_size = int(os.environ.get("WORLD_SIZE", -1))
-
-    torch.distributed.init_process_group("nccl")
-    initialize_model_parallel(world_size)
-    torch.cuda.set_device(local_rank)
-
-    # seed must be the same in all processes
-    torch.manual_seed(1)
-    return local_rank, world_size
-
 def download_pretrained_models(
     ckpt_path: str,
     tokenizer_path: str
@@ -52,33 +38,92 @@ def download_pretrained_models(
 
 def get_pretrained_models(
     ckpt_path: str, 
-    tokenizer_path: str, 
-    local_rank: int, 
-    world_size: int) -> LLaMA:
+    tokenizer_path: str) -> LLaMA:
 
     download_pretrained_models(ckpt_path, tokenizer_path)
+    # max_seq_len: int = 512, max_batch_size: int = 32
+    generator = load(
+        ckpt_dir=llama_weight_path, 
+        tokenizer_path=tokenizer_weight_path,
+        max_seq_len=512,
+        max_batch_size=1
+    )
 
+    return generator
+
+def load(
+    ckpt_dir: str,
+    tokenizer_path: str,
+    max_seq_len: int,
+    max_batch_size: int,
+) -> LLaMA:
     start_time = time.time()
-    checkpoints = sorted(Path(llama_weight_path).glob("*.pth"))
+    checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
+    print(checkpoints)
 
-    llama_ckpt_path = checkpoints[local_rank]
-    print("Loading")
-    checkpoint = torch.load(llama_ckpt_path, map_location="cpu")
-    with open(Path(llama_weight_path) / "params.json", "r") as f:
+    with open(Path(ckpt_dir) / "params.json", "r") as f:
         params = json.loads(f.read())
 
-    model_args: ModelArgs = ModelArgs(max_seq_len=512, max_batch_size=1, **params)
-    tokenizer = Tokenizer(model_path=f"{tokenizer_weight_path}/tokenizer.model")
+    model_args: ModelArgs = ModelArgs(
+        max_seq_len=max_seq_len, max_batch_size=max_batch_size, **params
+    )
+    tokenizer = Tokenizer(model_path=f"{tokenizer_path}/tokenizer.model")
     model_args.vocab_size = tokenizer.n_words
-    torch.set_default_tensor_type(torch.cuda.HalfTensor)
-    model = Transformer(model_args).cuda().half()
+
+    torch.set_default_tensor_type(torch.HalfTensor)
+    print("Allocating transformer on host")
+    ctx_tok = default_quantize.set(True)
+    model = Transformer(model_args)
+    default_quantize.set(ctx_tok)
+    key_to_dim = {
+        "w1": 0,
+        "w2": -1,
+        "w3": 0,
+        "wo": -1,
+        "wq": 0,
+        "wk": 0,
+        "wv": 0,
+        "output": 0,
+        "tok_embeddings": -1,
+        "ffn_norm": None,
+        "attention_norm": None,
+        "norm": None,
+        "rope": None,
+    }
+
+    # ?
     torch.set_default_tensor_type(torch.FloatTensor)
-    model.load_state_dict(checkpoint, strict=False)
+
+    # load the state dict incrementally, to avoid memory problems
+    for i, ckpt in enumerate(checkpoints):
+        print(f"Loading checkpoint {i}")
+        checkpoint = torch.load(ckpt, map_location="cpu")
+        for parameter_name, parameter in model.named_parameters():
+            short_name = parameter_name.split(".")[-2]
+            if key_to_dim[short_name] is None and i == 0:
+                parameter.data = checkpoint[parameter_name]
+            elif key_to_dim[short_name] == 0:
+                size = checkpoint[parameter_name].size(0)
+                parameter.data[size * i : size * (i + 1), :] = checkpoint[
+                    parameter_name
+                ]
+            elif key_to_dim[short_name] == -1:
+                size = checkpoint[parameter_name].size(-1)
+                parameter.data[:, size * i : size * (i + 1)] = checkpoint[
+                    parameter_name
+                ]
+            del checkpoint[parameter_name]
+        del checkpoint
+
+    model.cuda()
 
     generator = LLaMA(model, tokenizer)
-    print(f"Loaded in {time.time() - start_time:.2f} seconds")
+    print(
+        f"Loaded in {time.time() - start_time:.2f} seconds with {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GiB"
+    )
     return generator
 
+
 def get_output(
     generator: LLaMA,
     prompt: str,
@@ -94,4 +139,4 @@ def get_output(
         top_p=top_p
     )
 
-    return results
+    return results