TRTLLM ChatGLM CI

deepjavalibrary · Dec 21, 2023 · 8169eb8 · 8169eb8
1 parent 4d4d8c1
commit 8169eb8
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/.github/workflows/llm_integration.yml b/.github/workflows/llm_integration.yml
@@ -756,6 +756,17 @@ jobs:
           python3 llm/client.py trtllm baichuan2-13b
           rm -rf docker_env
           docker rm -f $(docker ps -aq)
+      - name: chatglm3-6b HF model with tp=4
+        working-directory: tests/integration
+        run: |
+          rm -rf models
+          echo -en "CUDA_VISIBLE_DEVICES=0,1,2,3" > docker_env
+          python3 llm/prepare.py trtllm chatglm3-6b
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models trtllm \
+          serve
+          python3 llm/client.py trtllm chatglm3-6b
+          rm -rf docker_env
+          docker rm -f $(docker ps -aq)
       - name: On fail step
         if: ${{ failure() }}
         working-directory: tests/integration

diff --git a/tests/integration/llm/client.py b/tests/integration/llm/client.py
@@ -467,7 +467,13 @@ def get_model_name():
         "batch_size": [1],
         "seq_length": [64, 128, 256],
         "stream_output": True,
-    }
+    },
+    "chatglm3-6b": {
+        "max_memory_per_gpu": [22.0],
+        "batch_size": [1],
+        "seq_length": [64, 128, 256],
+        "stream_output": True,
+    },
 }
 
 deepspeed_rolling_batch_model_spec = {

diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
@@ -702,6 +702,13 @@
         "option.baichuan_model_version": "v2_13b",
         "option.output_formatter": "jsonlines",
         "option.trust_remote_code": True
+    },
+    "chatglm3-6b": {
+        "option.model_id": "THUDM/chatglm3-6b",
+        "option.tensor_parallel_degree": 4,
+        "option.output_formatter": "jsonlines",
+        "option.trust_remote_code": True,
+        "option.chatglm_model_version": "chatglm3_6b"
     }
 }