Add lit-llama benchmarks (logits, autoregressive generation, lora fine tuning) (#1730)

Summary:
Signed-off-by: Edward Z. Yang <ezyang@meta.com>

Pull Request resolved: #1730

Reviewed By: xuzhao9

Differential Revision: D47295535

Pulled By: ezyang

fbshipit-source-id: 9e5998569b6dc58c6c918c4caaf1e6c896600430

Author: ezyang

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "submodules/FAMBench"]
 	path = submodules/FAMBench
 	url = https://github.com/facebookresearch/FAMBench.git
+[submodule "submodules/lit-llama"]
+	path = submodules/lit-llama
+	url = https://github.com/Lightning-AI/lit-llama.git

submodules/lit-llama

@@ -0,0 +1,52 @@
+from ...util.model import BenchmarkModel
+from torchbenchmark.tasks import NLP
+import torch
+import os
+from torchbenchmark import add_path, REPO_PATH
+import sys
+import lightning as L
+
+LIT_LLAMA_PATH = os.path.join(REPO_PATH, "submodules", "lit-llama")
+
+with add_path(LIT_LLAMA_PATH):
+    from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+    from lit_llama import LLaMA, Tokenizer
+
+class Model(BenchmarkModel):
+    task = NLP.LANGUAGE_MODELING
+    DEFAULT_EVAL_BSIZE = 1
+
+    def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):
+        super().__init__(test=test, device=device, jit=jit, batch_size=batch_size, extra_args=extra_args)
+
+        checkpoint_path = os.path.join(LIT_LLAMA_PATH, "checkpoints/lit-llama/7B/lit-llama.pth")
+        if not os.path.exists(checkpoint_path):
+            raise NotImplementedError("checkpoint doesn't exist")
+        with lazy_load(checkpoint_path) as checkpoint:
+            name = llama_model_lookup(checkpoint)
+
+            with EmptyInitOnDevice(device=device):
+                model = LLaMA.from_name(name)
+            model.load_state_dict(checkpoint)
+
+        self.model = model
+        self.seq_len = 32
+        self.max_seq_len = 64
+        self.example_inputs = (
+            torch.ones([self.batch_size, self.seq_len], dtype=torch.int32, device=self.device),
+            self.max_seq_len,
+            torch.arange(self.seq_len, dtype=torch.int64, device=self.device)  # positions
+        )
+
+
+    def get_module(self):
+        return self.model, self.example_inputs
+
+    def train(self):
+        return NotImplementedError("you will OOM trying to train directly")
+
+    def eval(self):
+        self.model.eval()
+        with torch.no_grad():
+            logits = self.model(*self.example_inputs)
+        return (logits,)

torchbenchmark/models/lit_llama/__init__.py

@@ -0,0 +1,4 @@
+from torchbenchmark.util.framework.lit_llama import install_lit_llama
+
+if __name__ == '__main__':
+    install_lit_llama()

torchbenchmark/models/lit_llama/install.py

@@ -0,0 +1,11 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 32
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+train_benchmark: false
+train_deterministic: false
+not_implemented:
+  - test: eval
+  - test: example

torchbenchmark/models/lit_llama/metadata.yaml

@@ -0,0 +1,41 @@
+from .. import lit_llama as lit_llama
+from ..lit_llama import LIT_LLAMA_PATH
+import importlib.util
+import os.path
+import torch.nn as nn
+import sys
+from lit_llama import Tokenizer
+
+def import_from_file_path(module_name, file_path):
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    sys.modules[module_name] = module
+    return module
+
+lit_llama_generate = import_from_file_path("lit_llama_generate", os.path.join(LIT_LLAMA_PATH, 'generate.py'))
+
+class GenerationWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, idx, max_new_tokens):
+        return lit_llama_generate.generate(self.model, idx, max_new_tokens)
+
+class Model(lit_llama.Model):
+    def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):
+        super().__init__(test=test, device=device, jit=jit, batch_size=batch_size, extra_args=extra_args)
+        self.model = GenerationWrapper(self.model)
+        tokenizer = Tokenizer(os.path.join(LIT_LLAMA_PATH, "checkpoints/lit-llama/tokenizer.model"))
+        # max_new_tokens matches lit-llama/generate.py
+        self.example_inputs = (tokenizer.encode("The meaning of life is", bos=True, eos=False, device=device), 50)
+
+    def train(self):
+        return NotImplementedError("cannot train on autoregressive generation")
+
+    def eval(self):
+        self.model.eval()
+        with torch.no_grad():
+            y = self.model(*self.example_inputs)
+        return (y,)

torchbenchmark/models/lit_llama_generate/__init__.py

@@ -0,0 +1,4 @@
+from torchbenchmark.util.framework.lit_llama import install_lit_llama
+
+if __name__ == '__main__':
+    install_lit_llama()

torchbenchmark/models/lit_llama_generate/install.py

@@ -0,0 +1,11 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 32
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+train_benchmark: false
+train_deterministic: false
+not_implemented:
+  - test: eval
+  - test: example

torchbenchmark/models/lit_llama_generate/metadata.yaml

@@ -0,0 +1,66 @@
+from ...util.model import BenchmarkModel
+from torchbenchmark.tasks import NLP
+import torch
+from ..lit_llama import LIT_LLAMA_PATH
+import importlib.util
+import os.path
+import torch.nn as nn
+import sys
+from lit_llama.lora import mark_only_lora_as_trainable, lora, lora_state_dict
+from torchbenchmark import REPO_PATH
+
+LIT_LLAMA_PATH = os.path.join(REPO_PATH, "submodules", "lit-llama")
+
+sys.path.insert(0, LIT_LLAMA_PATH)
+
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+from lit_llama import LLaMA, Tokenizer
+
+class Model(BenchmarkModel):
+    task = NLP.LANGUAGE_MODELING
+    DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 4  # micro_batch_size in lora.py
+
+    def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):
+        super().__init__(test=test, device=device, jit=jit, batch_size=batch_size, extra_args=extra_args)
+
+        # From finetune/lora.py hyperparameters
+        lora_r = 8
+        lora_alpha = 16
+        lora_dropout = 0.05
+
+        checkpoint_path = os.path.join(LIT_LLAMA_PATH, "checkpoints/lit-llama/7B/lit-llama.pth")
+        if not os.path.exists(checkpoint_path):
+            raise NotImplementedError("checkpoint doesn't exist")
+        with lazy_load(checkpoint_path) as checkpoint, lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
+            name = llama_model_lookup(checkpoint)
+
+            with EmptyInitOnDevice(device=device):
+                model = LLaMA.from_name(name)
+            # LoRA weights won't be in base checkpoint
+            model.load_state_dict(checkpoint, strict=False)
+
+        mark_only_lora_as_trainable(model)
+
+        self.model = model
+        self.seq_len = 32
+        self.max_seq_len = 64
+        self.example_inputs = (
+            torch.ones([self.batch_size, self.seq_len], dtype=torch.int32, device=self.device),
+            self.max_seq_len,
+        )
+
+
+    def get_module(self):
+        return self.model, self.example_inputs
+
+    def train(self):
+        logits = self.model(*self.example_inputs)
+        logits.sum().backward()
+        # meh this sucks
+
+    def eval(self):
+        self.model.eval()
+        with torch.no_grad():
+            logits = self.model(*self.example_inputs)
+        return (logits,)

torchbenchmark/models/lit_llama_lora/__init__.py

@@ -0,0 +1,4 @@
+from torchbenchmark.util.framework.lit_llama import install_lit_llama
+
+if __name__ == '__main__':
+    install_lit_llama()