Generate llama instead of downloading it (#250)

* Generate llama instead of downloading it

* Generate safe checkpoints for llama3 70B

* Fix llm requirements

* rename huggingface token to MILABENCH_* to automatically forward the env var to a remote in such cases

.github/workflows/tests_unit.yml

@@ -74,7 +74,7 @@ jobs:
 
       - name: tests
         env:
-          HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN}}
+          MILABENCH_HF_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN}}
         run: |
           source $(poetry env info -p)/bin/activate
           coverage run --source=milabench -m pytest --ignore=tests/integration tests/ 

benchmarks/llm/benchfile.py

@@ -1,3 +1,5 @@
+import tempfile
+from milabench.fs import XPath
 from milabench.pack import Package
 
 
@@ -38,6 +40,28 @@ class Llm(Package):
     prepare_script = "prepare.py"
 
     async def install(self):
+        llama3_dir = XPath(__file__).resolve().parent
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_dir = XPath(tmp_dir)
+            tmp_dir.clone_subtree(
+                "https://github.com/meta-llama/llama3.git",
+                "11817d47e1ba7a4959b025eb1ca308572e0e3963",
+            )
+            tmp_dir.merge_into(
+                llama3_dir,
+                manifest="\n".join(
+                    [
+                        "/llama/",
+                        "/requirements.txt",
+                    ]
+                )
+            )
+        # Fix conflict with tiktoken. As we only need llama/model.py, we don't
+        # need to care about a compatible tiktoken for the llama3 module
+        requirements = (llama3_dir / "requirements.txt").read_text().splitlines()
+        requirements = [l for l in requirements if not l.startswith("tiktoken==")]
+        (llama3_dir / "requirements.txt").write_text("\n".join(requirements))
+
         await super().install()  # super() call installs the requirements
 
     def build_run_plan(self):

benchmarks/llm/configs/llama3_70B_full.yaml

@@ -31,40 +31,73 @@ shuffle: True
 model:
   _component_: torchtune.models.llama3_1.llama3_1_70b
 
+safetensors: true
 checkpointer:
   _component_: torchtune.utils.FullModelHFCheckpointer
   checkpoint_dir: /tmp/Meta-Llama-3.1-70B-Instruct/
   checkpoint_files: [
-    model-00001-of-00030.safetensors,
-    model-00002-of-00030.safetensors,
-    model-00003-of-00030.safetensors,
-    model-00004-of-00030.safetensors,
-    model-00005-of-00030.safetensors,
-    model-00006-of-00030.safetensors,
-    model-00007-of-00030.safetensors,
-    model-00008-of-00030.safetensors,
-    model-00009-of-00030.safetensors,
-    model-00010-of-00030.safetensors,
-    model-00011-of-00030.safetensors,
-    model-00012-of-00030.safetensors,
-    model-00013-of-00030.safetensors,
-    model-00014-of-00030.safetensors,
-    model-00015-of-00030.safetensors,
-    model-00016-of-00030.safetensors,
-    model-00017-of-00030.safetensors,
-    model-00018-of-00030.safetensors,
-    model-00019-of-00030.safetensors,
-    model-00020-of-00030.safetensors,
-    model-00021-of-00030.safetensors,
-    model-00022-of-00030.safetensors,
-    model-00023-of-00030.safetensors,
-    model-00024-of-00030.safetensors,
-    model-00025-of-00030.safetensors,
-    model-00026-of-00030.safetensors,
-    model-00027-of-00030.safetensors,
-    model-00028-of-00030.safetensors,
-    model-00029-of-00030.safetensors,
-    model-00030-of-00030.safetensors,
+    model-00001-of-00062.safetensors,
+    model-00002-of-00062.safetensors,
+    model-00003-of-00062.safetensors,
+    model-00004-of-00062.safetensors,
+    model-00005-of-00062.safetensors,
+    model-00006-of-00062.safetensors,
+    model-00007-of-00062.safetensors,
+    model-00008-of-00062.safetensors,
+    model-00009-of-00062.safetensors,
+    model-00010-of-00062.safetensors,
+    model-00011-of-00062.safetensors,
+    model-00012-of-00062.safetensors,
+    model-00013-of-00062.safetensors,
+    model-00014-of-00062.safetensors,
+    model-00015-of-00062.safetensors,
+    model-00016-of-00062.safetensors,
+    model-00017-of-00062.safetensors,
+    model-00018-of-00062.safetensors,
+    model-00019-of-00062.safetensors,
+    model-00020-of-00062.safetensors,
+    model-00021-of-00062.safetensors,
+    model-00022-of-00062.safetensors,
+    model-00023-of-00062.safetensors,
+    model-00024-of-00062.safetensors,
+    model-00025-of-00062.safetensors,
+    model-00026-of-00062.safetensors,
+    model-00027-of-00062.safetensors,
+    model-00028-of-00062.safetensors,
+    model-00029-of-00062.safetensors,
+    model-00030-of-00062.safetensors,
+    model-00031-of-00062.safetensors,
+    model-00032-of-00062.safetensors,
+    model-00033-of-00062.safetensors,
+    model-00034-of-00062.safetensors,
+    model-00035-of-00062.safetensors,
+    model-00036-of-00062.safetensors,
+    model-00037-of-00062.safetensors,
+    model-00038-of-00062.safetensors,
+    model-00039-of-00062.safetensors,
+    model-00040-of-00062.safetensors,
+    model-00041-of-00062.safetensors,
+    model-00042-of-00062.safetensors,
+    model-00043-of-00062.safetensors,
+    model-00044-of-00062.safetensors,
+    model-00045-of-00062.safetensors,
+    model-00046-of-00062.safetensors,
+    model-00047-of-00062.safetensors,
+    model-00048-of-00062.safetensors,
+    model-00049-of-00062.safetensors,
+    model-00050-of-00062.safetensors,
+    model-00051-of-00062.safetensors,
+    model-00052-of-00062.safetensors,
+    model-00053-of-00062.safetensors,
+    model-00054-of-00062.safetensors,
+    model-00055-of-00062.safetensors,
+    model-00056-of-00062.safetensors,
+    model-00057-of-00062.safetensors,
+    model-00058-of-00062.safetensors,
+    model-00059-of-00062.safetensors,
+    model-00060-of-00062.safetensors,
+    model-00061-of-00062.safetensors,
+    model-00062-of-00062.safetensors,
   ]
   recipe_checkpoint: null
   output_dir: /tmp/Meta-Llama-3.1-70B-Instruct/

benchmarks/llm/configs/llama3_70B_lora.yaml

@@ -21,42 +21,73 @@ tokenizer:
   _component_: torchtune.models.llama3.llama3_tokenizer
   path: /tmp/Meta-Llama-3.1-70B-Instruct/original/tokenizer.model
 
-
 safetensors: true
 checkpointer:
   _component_: torchtune.utils.FullModelHFCheckpointer
   checkpoint_dir: /tmp/Meta-Llama-3.1-70B-Instruct/
   checkpoint_files: [
-    model-00001-of-00030.safetensors,
-    model-00002-of-00030.safetensors,
-    model-00003-of-00030.safetensors,
-    model-00004-of-00030.safetensors,
-    model-00005-of-00030.safetensors,
-    model-00006-of-00030.safetensors,
-    model-00007-of-00030.safetensors,
-    model-00008-of-00030.safetensors,
-    model-00009-of-00030.safetensors,
-    model-00010-of-00030.safetensors,
-    model-00011-of-00030.safetensors,
-    model-00012-of-00030.safetensors,
-    model-00013-of-00030.safetensors,
-    model-00014-of-00030.safetensors,
-    model-00015-of-00030.safetensors,
-    model-00016-of-00030.safetensors,
-    model-00017-of-00030.safetensors,
-    model-00018-of-00030.safetensors,
-    model-00019-of-00030.safetensors,
-    model-00020-of-00030.safetensors,
-    model-00021-of-00030.safetensors,
-    model-00022-of-00030.safetensors,
-    model-00023-of-00030.safetensors,
-    model-00024-of-00030.safetensors,
-    model-00025-of-00030.safetensors,
-    model-00026-of-00030.safetensors,
-    model-00027-of-00030.safetensors,
-    model-00028-of-00030.safetensors,
-    model-00029-of-00030.safetensors,
-    model-00030-of-00030.safetensors,
+    model-00001-of-00062.safetensors,
+    model-00002-of-00062.safetensors,
+    model-00003-of-00062.safetensors,
+    model-00004-of-00062.safetensors,
+    model-00005-of-00062.safetensors,
+    model-00006-of-00062.safetensors,
+    model-00007-of-00062.safetensors,
+    model-00008-of-00062.safetensors,
+    model-00009-of-00062.safetensors,
+    model-00010-of-00062.safetensors,
+    model-00011-of-00062.safetensors,
+    model-00012-of-00062.safetensors,
+    model-00013-of-00062.safetensors,
+    model-00014-of-00062.safetensors,
+    model-00015-of-00062.safetensors,
+    model-00016-of-00062.safetensors,
+    model-00017-of-00062.safetensors,
+    model-00018-of-00062.safetensors,
+    model-00019-of-00062.safetensors,
+    model-00020-of-00062.safetensors,
+    model-00021-of-00062.safetensors,
+    model-00022-of-00062.safetensors,
+    model-00023-of-00062.safetensors,
+    model-00024-of-00062.safetensors,
+    model-00025-of-00062.safetensors,
+    model-00026-of-00062.safetensors,
+    model-00027-of-00062.safetensors,
+    model-00028-of-00062.safetensors,
+    model-00029-of-00062.safetensors,
+    model-00030-of-00062.safetensors,
+    model-00031-of-00062.safetensors,
+    model-00032-of-00062.safetensors,
+    model-00033-of-00062.safetensors,
+    model-00034-of-00062.safetensors,
+    model-00035-of-00062.safetensors,
+    model-00036-of-00062.safetensors,
+    model-00037-of-00062.safetensors,
+    model-00038-of-00062.safetensors,
+    model-00039-of-00062.safetensors,
+    model-00040-of-00062.safetensors,
+    model-00041-of-00062.safetensors,
+    model-00042-of-00062.safetensors,
+    model-00043-of-00062.safetensors,
+    model-00044-of-00062.safetensors,
+    model-00045-of-00062.safetensors,
+    model-00046-of-00062.safetensors,
+    model-00047-of-00062.safetensors,
+    model-00048-of-00062.safetensors,
+    model-00049-of-00062.safetensors,
+    model-00050-of-00062.safetensors,
+    model-00051-of-00062.safetensors,
+    model-00052-of-00062.safetensors,
+    model-00053-of-00062.safetensors,
+    model-00054-of-00062.safetensors,
+    model-00055-of-00062.safetensors,
+    model-00056-of-00062.safetensors,
+    model-00057-of-00062.safetensors,
+    model-00058-of-00062.safetensors,
+    model-00059-of-00062.safetensors,
+    model-00060-of-00062.safetensors,
+    model-00061-of-00062.safetensors,
+    model-00062-of-00062.safetensors,
   ]
   recipe_checkpoint: null
   output_dir: /tmp/Meta-Llama-3.1-70B-Instruct/

benchmarks/llm/llama/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed in accordance with the terms of the Llama 3 Community License Agreement.
+
+from .generation import Llama
+from .model import ModelArgs, Transformer
+from .tokenizer import Dialog, Tokenizer