[ci][test] add correctness test for cpu offloading (#6549)

2024-07-18 16:41:06 -07:00
parent 2d4733ba2d
commit f53b8f0d05
4 changed files with 105 additions and 85 deletions
--- a/tests/basic_correctness/test_cpu_offload.py
+++ b/tests/basic_correctness/test_cpu_offload.py
@@ -0,0 +1,8 @@
+from ..utils import compare_two_settings
+
+
+def test_cpu_offload():
+    compare_two_settings("meta-llama/Llama-2-7b-hf", [],
+                         ["--cpu-offload-gb", "4"])
+    compare_two_settings("nm-testing/llama7b-one-shot-2_4-w4a16-marlin24-t",
+                         [], ["--cpu-offload-gb", "1"])