add gguf path option and gguf tests (pytorch#168)

mikekgfb · malfet · commit 592a307ff8cf · 2024-07-16T23:03:11.000-07:00
* add gguf path option and gguf tests

* tokenizer path for GGUF

* tab-&gt;spc

* GGUF_PATH definition

* pip install gguf
diff --git a/.github/workflows/compile-gguf.yml b/.github/workflows/compile-gguf.yml
@@ -0,0 +1,80 @@
+name: Compile main using GGUF
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  run-tinystories:
+    strategy:
+      matrix:
+        runner: [ubuntu-latest, macos-14]
+    runs-on: ${{matrix.runner}}
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v2
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.11
+      - name: Print machine info
+        run: |
+          uname -a
+          if [ $(uname -s) == Darwin ]; then
+            sysctl machdep.cpu.brand_string
+            sysctl machdep.cpu.core_count
+          fi
+      - name: Install requirements
+        run: |
+          pip install gguf
+          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+          pip install -r requirements.txt
+      - name: Download GGUF
+        run: |
+          mkdir gguf_files
+          export GGUF_PATH=gguf_files/llama-2-7b.Q4_0.gguf
+          export TOKENIZER_PATH=gguf_files/tokenizer.model
+          wget -O ${GGUF_PATH} "https://huggingface.co/TheBloke/Llama-2-7B-GGUF/resolve/main/llama-2-7b.Q4_0.gguf?download=true"
+          wget -O ${TOKENIZER_PATH} https://github.com/karpathy/llama2.c/raw/master/tokenizer.model
+      - name: Run inference
+        run: |          
+          export GGUF_PATH=gguf_files/llama-2-7b.Q4_0.gguf
+          export TOKENIZER_PATH=gguf_files/tokenizer.model
+          export MODEL_NAME=llama-2-7b.Q4_0.gguf
+          export MODEL_DIR=/tmp
+          python generate.py --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_eager
+          cat ./output_eager
+          python generate.py --compile --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_compiled
+          cat ./output_compiled
+          python export.py --gguf-path ${GGUF_PATH} --output-dso-path ${MODEL_DIR}/${MODEL_NAME}.so
+          python generate.py --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so  > ./output_aoti
+          cat ./output_aoti
+
+          echo "******************************************"
+          echo "******* Emb: channel-wise quantized ******"
+          echo "******************************************"
+          python generate.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 0}}' --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_eager
+          cat ./output_eager
+          python generate.py --compile --quant '{"embedding" : {"bitwidth": 8, "group_size": 0}}' --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_compiled
+          cat ./output_compiled
+          python export.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 0}}' --gguf-path ${GGUF_PATH} --output-dso-path ${MODEL_DIR}/${MODEL_NAME}.so
+          python generate.py --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so  > ./output_aoti
+          cat ./output_aoti
+
+          echo "******************************************"
+          echo "******** Emb: group-wise quantized *******"
+          echo "******************************************"
+          python generate.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_eager
+          cat ./output_eager
+          python generate.py --compile --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 > ./output_compiled
+          cat ./output_compiled
+          python export.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --gguf-path ${GGUF_PATH} --output-dso-path ${MODEL_DIR}/${MODEL_NAME}.so
+          python generate.py --gguf-path ${GGUF_PATH} --tokenizer-path ${TOKENIZER_PATH} --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so  > ./output_aoti
+          cat ./output_aoti
+
+          echo "tests complete"
+          echo "******************************************"
+
diff --git a/cli.py b/cli.py
@@ -152,6 +152,12 @@ def cli_args():
         default=None,
         help="Parameter file path.",
     )
+    parser.add_argument(
+        "--gguf-path",
+        type=Path,
+        default=None,
+        help="GGUF file path.",
+    )
     parser.add_argument(
         "--tokenizer-path",
         type=Path,
diff --git a/eval.py b/eval.py
@@ -215,8 +215,8 @@ def eval_main(args) -> None:
     checkpoint_dir = args.checkpoint_dir
     params_path = args.params_path
     params_table = args.params_table
+    gguf_path = args.gguf_path
     tokenizer_path = args.tokenizer_path
-    params_path = args.params_path
     dso_path = args.dso_path
     pte_path = args.pte_path
     quantize = args.quantize
@@ -241,6 +241,7 @@ def eval_main(args) -> None:
         checkpoint_dir,
         params_path,
         params_table,
+        gguf_path,
         dso_path,
         pte_path,
         quantize,
diff --git a/export.py b/export.py
@@ -79,6 +79,7 @@ def main(args):
         args.checkpoint_dir,
         args.params_path,
         args.params_table,
+        args.gguf_path,
         device=device,
         precision=precision,
         use_tp=False
diff --git a/generate.py b/generate.py
@@ -279,16 +279,19 @@ def _load_model(
         checkpoint_dir,
         params_path,
         params_table,
+        gguf_path,
         device,
         precision,
-        use_tp=False
+        use_tp # =False
 ):
     use_cuda = "cuda" in device
     with torch.device("meta"):
         if params_path:
             model = Transformer.from_params(params_path)
         elif params_table:
-            model = Transformer.from_table(params_path)            
+            model = Transformer.from_table(params_path)
+        elif gguf_path:
+            model = Transformer.from_gguf(gguf_path)            
         else:
             model = Transformer.from_name(checkpoint_path.parent.name)
 
@@ -343,22 +346,29 @@ def _load_inference_model(
         checkpoint_dir,
         params_path,
         params_table,
+        gguf_path,
         dso_path,
         pte_path,
         quantize,
         device,
         precision,
-        use_tp=False
+        use_tp # =False
 ):
     assert (
         (checkpoint_path and checkpoint_path.is_file()) or
+        (checkpoint_dir and checkpoint_path.is_dir()) or
+        (gguf_path and gguf_path.is_file()) or
         (dso_path and Path(dso_path).is_file()) or
         (pte_path and Path(pte_path).is_file())
-    ), "need to specified a valid checkpoint path, DSO path, or PTE path"
+    ), "need to specified a valid checkpoint path, checkpoint dir, gguf path, DSO path, or PTE path"
     assert not (dso_path and pte_path), "specify either DSO path or PTE path, but not both"
 
     if (checkpoint_path and (dso_path or pte_path)):
         print("Warning: checkpoint path ignored because an exported DSO or PTE path specified")
+    if (checkpoint_dir and (dso_path or pte_path)):
+        print("Warning: checkpoint dir ignored because an exported DSO or PTE path specified")
+    if (gguf_path and (dso_path or pte_path)):
+        print("Warning: GGUF path ignored because an exported DSO or PTE path specified")
 
     print("Loading model ...")
     t0 = time.time()    
@@ -367,6 +377,7 @@ def _load_inference_model(
         checkpoint_dir,
         params_path,
         params_table,
+        gguf_path,
         device,
         precision,
         use_tp
@@ -423,6 +434,7 @@ def _main(
     checkpoint_dir: Optional[Path] = None,
     params_path: Optional[Path] = None,
     params_table: Optional[str] = None,
+    gguf_path: Optional[Path] = None,
     tokenizer_path: Optional[Path] = None,
     compile: bool = True,
     compile_prefill: bool = False,
@@ -463,6 +475,7 @@ def _main(
         checkpoint_dir,
         params_path,
         params_table,
+        gguf_path,
         dso_path,
         pte_path,
         quantize,
@@ -476,8 +489,10 @@ def _main(
     if is_speculative:
         draft_model = _load_model(
             draft_checkpoint_path,
-            None,
-            None,
+            None, # checkpoint_dir
+            None, # params_path,
+            None, # params_table
+            None, # gguf_path
             device,
             precision,
             use_tp
@@ -619,6 +634,7 @@ def main(args):
         args.checkpoint_dir,
         args.params_path,
         args.params_table,
+        args.gguf_path,
         args.tokenizer_path,
         args.compile,
         args.compile_prefill,
diff --git a/model.py b/model.py
@@ -234,6 +234,12 @@ def from_table(cls, name: str):
     def from_params(cls, params_path: str):
         return cls(ModelArgs.from_params(params_path))
 
+    @classmethod
+    def from_gguf(cls, gguf_path: str):
+        from gguf_loader import load_llama_from_gguf_file
+        model = load_llama_from_gguf_file(gguf_path)
+        return model
+    
 
 class TransformerBlock(nn.Module):
     def __init__(self, config: ModelArgs) -> None:
