Merge branch 'main' into egor/margin

Author: Egor-Krivov

dl_bench/llm.py

@@ -16,31 +16,31 @@
 
 
 def get_llm(name, dtype):
-    if name == "gptj":
-        model_name = "EleutherAI/gpt-j-6B"
-
-        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=dtype)
-        tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
-    elif name == "llama2-13b":
-        kwargs = {}
-        if "HF_TOKEN" in os.environ:
-            kwargs["token"] = os.environ.get("HF_TOKEN")
-
-        model_name = "meta-llama/Llama-2-13b-hf"
-        model = LlamaForCausalLM.from_pretrained(
-            model_name, torch_dtype=dtype, **kwargs
-        )
-        tokenizer = LlamaTokenizer.from_pretrained(model_name, **kwargs)
-    else:
+    name2params = {
+        "gptj": ("EleutherAI/gpt-j-6B", AutoModelForCausalLM, AutoTokenizer),
+        "llama2-7b": ("meta-llama/Llama-2-7b-hf", LlamaForCausalLM, LlamaTokenizer),
+        "llama2-13b": ("meta-llama/Llama-2-13b-hf", LlamaForCausalLM, LlamaTokenizer),
+    }
+
+    if name not in name2params:
         raise ValueError("Unsupported model name")
+
+    kwargs = {}
+    if name.startswith("llama2") and "HF_TOKEN" in os.environ:
+        kwargs = {"HF_TOKEN": os.environ.get("HF_TOKEN")}
+
+    model_name, M, T = name2params[name]
+
+    model = M.from_pretrained(model_name, torch_dtype=dtype, **kwargs)
+    tokenizer = T.from_pretrained(model_name)
     return tokenizer, model
 
 
 class LlmBenchmark(Benchmark):
     def __init__(self, params) -> None:
         name = params.get("name", "gptj")
         dtype = params.get("dtype")
-        self.batch_size = params.get("batch_size", 1)
+        self.batch_size = int(params.get("batch_size", 1))
         self.n_iter = params.get("n_iter", 5)
         self.warmup_batches = params.get("warmup", 2)
 
@@ -90,12 +90,13 @@ def inference(self, backend):
             with torch.inference_mode(), cast:
                 tokens, total_time = self.generate(backend)
 
+            print(f"Fw time: {total_time:.1f}")
+
             if i < self.warmup_batches:
                 # We restart timer because that was just a warmup
                 start = get_time()
                 continue
 
-            print(f"Fw time: {total_time:.1f}")
             fw_times.append(total_time)
             n_items += math.prod(tokens.shape)
             outputs.append(tokens)

llm.sh

@@ -9,7 +9,7 @@ if [[ -z "${DL_BENCH_ARGS}" ]]; then
   exit 1
 fi
 
-for NAME in llama2-13b gptj
+for NAME in llama2-7b llama2-13b gptj
 do
   for BS in 1 4
   do