build(llama): 添加 LLaMA 寒武纪实现

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Jun 20, 2024 · aef4f67 · aef4f67
1 parent c6b782e
commit aef4f67
Show file tree

Hide file tree

Showing 18 changed files with 62 additions and 16 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,9 +14,10 @@ members = [
     "devices/nvidia",
 
     "models/llama/common",
-    "models/llama/cpu",
-    "models/llama/nvidia",
-    "models/llama/nvidia-distributed",
+    "models/llama/common-cpu",
+    "models/llama/nvidia-gpu",
+    "models/llama/nvidia-gpu-distributed",
+    "models/llama/cambricon-mlu",
     "models/mixtral/common",
     "models/mixtral/cpu",
 ]
@@ -38,3 +39,6 @@ operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "cac4307"
 cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a4b3cf5" }
 nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a4b3cf5" }
 search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a4b3cf5" }
+
+cndrv = { git = "https://github.com/YdrMaster/cndrv", rev = "d94e170" }
+search-neuware-tools = { git = "https://github.com/YdrMaster/cndrv", rev = "d94e170" }
diff --git a/models/llama/cambricon-mlu/Cargo.toml b/models/llama/cambricon-mlu/Cargo.toml
@@ -0,0 +1,14 @@
+[package]
+name = "llama-cn"
+version = "0.0.0"
+edition = "2021"
+authors = ["YdrMaster <[email protected]>"]
+
+[dependencies]
+common = { path = "../../../common" }
+causal-lm = { path = "../../../causal-lm" }
+llama = { path = "../common" }
+
+[build-dependencies]
+build-script-cfg.workspace = true
+search-neuware-tools.workspace = true
diff --git a/models/llama/cambricon-mlu/build.rs b/models/llama/cambricon-mlu/build.rs
@@ -0,0 +1,9 @@
+fn main() {
+    use build_script_cfg::Cfg;
+    use search_neuware_tools::find_neuware_home;
+
+    let neuware = Cfg::new("detected_neuware");
+    if find_neuware_home().is_some() {
+        neuware.define();
+    }
+}
diff --git a/models/llama/cambricon-mlu/src/lib.rs b/models/llama/cambricon-mlu/src/lib.rs
@@ -0,0 +1 @@
+#![cfg(detected_neuware)]
diff --git a/models/llama/cpu/Cargo.toml → models/llama/common-cpu/Cargo.toml b/models/llama/cpu/Cargo.toml → models/llama/common-cpu/Cargo.toml
diff --git a/models/llama/cpu/src/lib.rs → models/llama/common-cpu/src/lib.rs b/models/llama/cpu/src/lib.rs → models/llama/common-cpu/src/lib.rs
diff --git a/models/llama/nvidia-distributed/Cargo.toml → ...s/llama/nvidia-gpu-distributed/Cargo.toml b/models/llama/nvidia-distributed/Cargo.toml → ...s/llama/nvidia-gpu-distributed/Cargo.toml
diff --git a/models/llama/nvidia-distributed/build.rs → models/llama/nvidia-gpu-distributed/build.rs b/models/llama/nvidia-distributed/build.rs → models/llama/nvidia-gpu-distributed/build.rs
diff --git a/...lama/nvidia-distributed/src/distribute.rs → .../nvidia-gpu-distributed/src/distribute.rs b/...lama/nvidia-distributed/src/distribute.rs → .../nvidia-gpu-distributed/src/distribute.rs
diff --git a/models/llama/nvidia-distributed/src/lib.rs → ...s/llama/nvidia-gpu-distributed/src/lib.rs b/models/llama/nvidia-distributed/src/lib.rs → ...s/llama/nvidia-gpu-distributed/src/lib.rs
diff --git a/...lama/nvidia-distributed/src/parameters.rs → .../nvidia-gpu-distributed/src/parameters.rs b/...lama/nvidia-distributed/src/parameters.rs → .../nvidia-gpu-distributed/src/parameters.rs
diff --git a/models/llama/nvidia/Cargo.toml → models/llama/nvidia-gpu/Cargo.toml b/models/llama/nvidia/Cargo.toml → models/llama/nvidia-gpu/Cargo.toml
diff --git a/models/llama/nvidia/build.rs → models/llama/nvidia-gpu/build.rs b/models/llama/nvidia/build.rs → models/llama/nvidia-gpu/build.rs
diff --git a/models/llama/nvidia/src/lib.rs → models/llama/nvidia-gpu/src/lib.rs b/models/llama/nvidia/src/lib.rs → models/llama/nvidia-gpu/src/lib.rs
diff --git a/models/llama/nvidia/src/resource.rs → models/llama/nvidia-gpu/src/resource.rs b/models/llama/nvidia/src/resource.rs → models/llama/nvidia-gpu/src/resource.rs
diff --git a/service/Cargo.toml b/service/Cargo.toml
@@ -16,4 +16,4 @@ tokio.workspace = true
 
 [dev-dependencies]
 colored = "2.1"
-llama-cpu = { path = "../models/llama/cpu" }
+llama-cpu = { path = "../models/llama/common-cpu" }
diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml
@@ -15,9 +15,10 @@ web-api = { path = "../web-api" }
 
 # models
 llama = { path = "../models/llama/common" }
-llama-cpu = { path = "../models/llama/cpu" }
-llama-nv = { path = "../models/llama/nvidia", optional = true }
-llama-nv-distributed = { path = "../models/llama/nvidia-distributed", optional = true }
+llama-cpu = { path = "../models/llama/common-cpu" }
+llama-nv = { path = "../models/llama/nvidia-gpu", optional = true }
+llama-nv-distributed = { path = "../models/llama/nvidia-gpu-distributed", optional = true }
+llama-cn = { path = "../models/llama/cambricon-mlu", optional = true }
 mixtral = { path = "../models/mixtral/common" }
 mixtral-cpu = { path = "../models/mixtral/cpu" }