[CI] Add APC tests

HabanaAI · Feb 25, 2025 · 730a6a2 · 730a6a2
1 parent 7a16eb9
commit 730a6a2
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/.jenkins/lm-eval-harness/run-tests.sh b/.jenkins/lm-eval-harness/run-tests.sh
@@ -9,11 +9,14 @@ usage() {
     echo
     echo "  -c    - path to the test data config (e.g. configs/small-models.txt)"
     echo "  -t    - tensor parallel size"
+    echo "  -a    - enable automatic prefix caching"
+
     echo
 }
 
 SUCCESS=0
-
+APC_ENABLED="false"
+TP_SIZE=1
 while getopts "c:t:" OPT; do
   case ${OPT} in
     c ) 
@@ -22,6 +25,9 @@ while getopts "c:t:" OPT; do
     t )
         TP_SIZE="$OPTARG"
         ;;
+    a )
+        APC_ENABLED="true"
+        ;;
     \? )
         usage
         exit 1
@@ -40,6 +46,7 @@ do
 
     export LM_EVAL_TEST_DATA_FILE=$PWD/configs/${MODEL_CONFIG}
     export LM_EVAL_TP_SIZE=$TP_SIZE
+    export LM_EVAL_APC_ENABLED=$APC_ENABLED
     export PT_HPU_ENABLE_LAZY_COLLECTIVES=true
     export VLLM_SKIP_WARMUP=true
     export TQDM_BAR_FORMAT="{desc}: {percentage:3.0f}% {bar:10} | {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" 

diff --git a/.jenkins/lm-eval-harness/test_lm_eval_correctness.py b/.jenkins/lm-eval-harness/test_lm_eval_correctness.py
@@ -5,6 +5,7 @@
 
 * export LM_EVAL_TEST_DATA_FILE=configs/Meta-Llama-3-70B-Instruct.yaml
 * export LM_EVAL_TP_SIZE=4 
+* export LM_EVAL_APC_ENABLED=true 
 * pytest -s test_lm_eval_correctness.py
 """
 import atexit
@@ -29,6 +30,7 @@
                                     "false") in ['1', 'true']
 
 TP_SIZE = os.environ.get("LM_EVAL_TP_SIZE", 1)
+APC_ENABLED = os.environ.get("LM_EVAL_APC_ENABLED", 'false')
 
 
 def setup_fp8():
@@ -47,6 +49,7 @@ def launch_lm_eval(eval_config):
     max_num_seqs = eval_config.get('max_num_seqs', 128)
     model_args = f"pretrained={eval_config['model_name']}," \
                  f"tensor_parallel_size={TP_SIZE}," \
+                 f"enable_prefix_caching={APC_ENABLED}," \
                  f"add_bos_token=true," \
                  f"dtype={dtype}," \
                  f"max_model_len=4096," \

diff --git a/.jenkins/test_config.yaml b/.jenkins/test_config.yaml
@@ -14,6 +14,14 @@ stages:
       - name: gsm8k_small_g2_tp2
         flavor: g2.s
         command: cd .jenkins/lm-eval-harness && bash run-tests.sh -c configs/models-small.txt -t 2
+  - name: test_gsm8k_small_models_apc
+    steps:
+      - name: gsm8k_small_g3_tp1
+        flavor: g3
+        command: cd .jenkins/lm-eval-harness && bash run-tests.sh -c configs/models-small.txt -t 1 -a
+      - name: gsm8k_small_g2_tp1
+        flavor: g2
+        command: cd .jenkins/lm-eval-harness && bash run-tests.sh -c configs/models-small.txt -t 1 -a
   - name: test_gsm8k_large_models
     steps:
       - name: gsm8k_large_g3_tp2