fix oom (#1335)

* decrease batch to 2

* half data size

* get batch size based on vram

* fix bench not found

* fix 'int' object is not callable

* [CI] share GPU default false

* batch 1

* keep cache

Author: CSY-ModelCloud

.github/workflows/unit_tests.yml

@@ -46,7 +46,7 @@ on:
         description: 'one test, one gpu. for collecting statistics'
         type: boolean
         required: false
-        default: false
+        default: true
       server:
         description: 'Choose server (zen4 or xeon5)'
         required: true
@@ -661,7 +661,7 @@ jobs:
       - name: Clean cache
         if: always()
         run: |
-          rm ~/.cache/evalplus/*pkl || true
+          # rm ~/.cache/evalplus/*pkl || true
           pip cache purge && uv cache clean && rm -rf ./* ./.*
 
   show-statistics:

tests/models/model_test.py

@@ -19,6 +19,10 @@
 import sys
 from typing import Dict, List
 
+from device_smi import Device
+
+from gptqmodel.models._const import CUDA_0
+
 if sys.platform == "darwin":
     os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
@@ -131,9 +135,12 @@ def load_tokenizer(self, model_id_or_path, trust_remote_code=False):
         return tokenizer
 
     @classmethod
-    def load_dataset(self, tokenizer, rows: int = DATASET_SIZE):
+    def load_dataset(self, tokenizer=None, rows: int = DATASET_SIZE):
         traindata = load_dataset("json", data_files="/monster/data/model/dataset/c4-train.00000-of-01024.json.gz", split="train")
 
+        if not tokenizer:
+            return traindata.select(range(rows))
+
         datas = []
         for index, sample in enumerate(traindata):
             tokenized = tokenizer(sample['text'])
@@ -369,3 +376,6 @@ def clear_directory(self, directory_path):
                 os.unlink(item_path)
             elif os.path.isdir(item_path):
                 shutil.rmtree(item_path)
+
+    def get_batch_size(self):
+        return 32 if Device(CUDA_0).memory_total / 1024 / 1024 / 1024 > 24 else 2

tests/test_eval.py

@@ -25,11 +25,12 @@
 
 from gptqmodel import GPTQModel  # noqa: E402
 from gptqmodel.utils.eval import EVAL  # noqa: E402
+from models.model_test import ModelTest  # noqa: E402
 from lm_eval.tasks import TaskManager  # noqa: E402
 from parameterized import parameterized  # noqa: E402
 
 
-class TestEval(unittest.TestCase):
+class TestEval(ModelTest):
     @classmethod
     def setUpClass(self):
         self.MODEL_ID = "/monster/data/model/Llama-3.2-1B-Instruct-gptqmodel-4bit-vortex-v1"
@@ -54,7 +55,7 @@ def test_eval_gptqmodel(self, framework: Union[Type[EVAL.LM_EVAL],Type[EVAL.EVAL
             results = GPTQModel.eval(model_or_id_or_path=self.MODEL_ID,
                                      framework=framework,
                                      tasks=[task],
-                                     batch_size=8,
+                                     batch_size=1,
                                      output_path=output_path,
                                      llm_backend=llm_backend,
                                      model_args=model_args,

tests/test_post_quant_eora.py

@@ -76,7 +76,7 @@ def test_post_quant_eora(self):
         desc_act = True
         rank = 256
         batch_size = 1
-        calibration_dataset_rows = 1024
+        calibration_dataset_rows = 512
         calibration_dataset_concat_size = 0  # disable
         auto_gc = False
         adapter_file_name = "eora.safetensors"
@@ -93,11 +93,7 @@ def test_post_quant_eora(self):
             "adapter_file_name": adapter_file_name,
         }
 
-        calibration_dataset = load_dataset(
-            "allenai/c4",
-            data_files="en/c4-train.00001-of-01024.json.gz",
-            split="train"
-        ).select(range(calibration_dataset_rows))["text"]
+        calibration_dataset = self.load_dataset(rows=calibration_dataset_rows)["text"]
 
         with tempfile.TemporaryDirectory() as tmpdir:
             eora = Lora(

tests/test_quant_and_eora.py

@@ -32,30 +32,6 @@
 from tabulate import tabulate  # noqa: E402
 
 
-def bench(path: str, backend: BACKEND, adapter: Optional[Lora]):
-    # test post-quant inference
-    model = GPTQModel.load(
-        model_id_or_path=path,
-        backend=backend,
-        adapter=adapter,
-    )
-
-    tokens = model.generate("Capital of France is")[0]
-    result = model.tokenizer.decode(tokens)
-    print(f"BACKEND: {backend}, Result: {result}")
-    assert "paris" in result.lower(), f"`paris` not found in `{result}`"
-
-    bench_result = GPTQModel.eval(
-        model_or_id_or_path=model,
-        framework=EVAL.LM_EVAL,
-        tasks=[EVAL.LM_EVAL.ARC_CHALLENGE, EVAL.LM_EVAL.MMLU],
-        batch_size=32,
-    )
-
-    del model
-    torch_empty_cache()
-
-    return bench_result
 
 class Test(ModelTest):
     NATIVE_MODEL_ID = "/monster/data/model/Qwen2.5-0.5B-Instruct/"
@@ -140,8 +116,8 @@ def test_quant_and_eora(self):
 
             # BACKEND.EXLLAMA_V2, BACKEND.EXLLAMA_V1, BACKEND.TRITON, BACKEND.CUDA,
             for backend in [ BACKEND.MARLIN ]: # BACKEND.IPEX, BACKEND.BITBLAS, BACKEND.EXLLAMA_V2V BACKEND.MARLIN
-                base_bench = bench(path=tmpdir, backend=backend, adapter=None) # inference using qweights only
-                eora_bench = bench(path=tmpdir, backend=backend, adapter=eora) # inference using eora (lora)
+                base_bench = self.bench(path=tmpdir, backend=backend, adapter=None) # inference using qweights only
+                eora_bench = self.bench(path=tmpdir, backend=backend, adapter=eora) # inference using eora (lora)
 
                 print('--------GPTQModel + EoRA Config ---------')
 
@@ -158,3 +134,28 @@ def test_quant_and_eora(self):
                 print(make_table(eora_bench))
                 if "groups" in eora_bench:
                     print(make_table(eora_bench, "groups"))
+
+    def bench(self, path: str, backend: BACKEND, adapter: Optional[Lora]):
+        # test post-quant inference
+        model = GPTQModel.load(
+            model_id_or_path=path,
+            backend=backend,
+            adapter=adapter,
+        )
+
+        tokens = model.generate("Capital of France is")[0]
+        result = model.tokenizer.decode(tokens)
+        print(f"BACKEND: {backend}, Result: {result}")
+        assert "paris" in result.lower(), f"`paris` not found in `{result}`"
+
+        bench_result = GPTQModel.eval(
+            model_or_id_or_path=model,
+            framework=EVAL.LM_EVAL,
+            tasks=[EVAL.LM_EVAL.ARC_CHALLENGE, EVAL.LM_EVAL.MMLU],
+            batch_size=self.get_batch_size(),
+        )
+
+        del model
+        torch_empty_cache()
+
+        return bench_result

tests/test_quant_formats.py

@@ -78,7 +78,7 @@ def test_quantize(self, method: QUANT_METHOD, backend: BACKEND, sym: bool, forma
             self.pretrained_model_id,
             quantize_config=quantize_config,
         )
-        model.quantize(self.calibration_dataset, batch_size=32)
+        model.quantize(self.calibration_dataset, batch_size=self.get_batch_size())
 
         with tempfile.TemporaryDirectory() as tmpdirname:
             model.save(tmpdirname)

tests/test_quant_formats_auto_round.py

@@ -76,7 +76,7 @@ def test_quantize(self, method: QUANT_METHOD, backend: BACKEND, sym: bool, forma
             self.pretrained_model_id,
             quantize_config=quantize_config,
         )
-        model.quantize(self.calibration_dataset, batch_size=32)
+        model.quantize(self.calibration_dataset, batch_size=self.get_batch_size())
 
         with tempfile.TemporaryDirectory() as tmpdirname:
             model.save(tmpdirname)