From 49cce0c4373db3f5bc0b7748269eb6ac03576e18 Mon Sep 17 00:00:00 2001
From: chenxl <chenxl6436@outlook.com>
Date: Fri, 30 Aug 2024 03:24:26 +0000
Subject: [PATCH] [fix] bugs about Qwen57B, install requirement, Dockerfile

---
 Dockerfile                                                 | 3 ++-
 .../optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml   | 7 -------
 pyproject.toml                                             | 1 +
 3 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 707c1a8..d56a607 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -12,6 +12,7 @@ EOF
 
 FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel as compile_server
 WORKDIR /workspace
+ENV CUDA_HOME /usr/local/cuda
 COPY --from=web_compile /home/ktransformers /workspace/ktransformers
 RUN <<EOF
 apt update -y &&  apt install -y  --no-install-recommends \
@@ -27,7 +28,7 @@ git submodule init &&
 git submodule update &&
 pip install ninja pyproject numpy cpufeature &&
 pip install flash-attn &&
-CPU_INSTRUCT=NATIVE  KTRANSFORMERS_FORCE_BUILD=TRUE TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9" pip install . --no-build-isolation --verbose &&
+CPU_INSTRUCT=NATIVE  KTRANSFORMERS_FORCE_BUILD=TRUE TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9;9.0+PTX" pip install . --no-build-isolation --verbose &&
 pip cache purge
 EOF
 
diff --git a/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml b/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml
index 989e4b8..0cc2edf 100644
--- a/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml
+++ b/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml
@@ -1,10 +1,3 @@
-- match:
-    name: "^model\\.layers\\..*\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
 - match:
     class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding
   replace:
diff --git a/pyproject.toml b/pyproject.toml
index adeb8a9..8070241 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,7 @@ dependencies = [
   "torch >= 2.3.0",
   "transformers == 4.43.2",
   "fastapi >= 0.111.0",
+  "uvicorn >= 0.30.1",
   "langchain >= 0.2.0",
   "blessed >= 1.20.0",
   "accelerate >= 0.31.0",