refactor(distributed): 重写分布式推理

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 26, 2024 · 3e69ec6 · 3e69ec6
1 parent 36a72df
commit 3e69ec6
Show file tree

Hide file tree

Showing 7 changed files with 443 additions and 301 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -18,6 +18,7 @@ resolver = "2"
 [workspace.dependencies]
 half = "2.4"
 rayon = "1.9"
+itertools = "0.12"
 serde = "1.0"
 serde_json = "1.0"
 log = "0.4"

diff --git a/nvidia/distributed/Cargo.toml b/nvidia/distributed/Cargo.toml
@@ -7,11 +7,13 @@ authors = ["YdrMaster <[email protected]>"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+causal-lm = { path = "../../causal-lm" }
 transformer = { path = "../../transformer" }
 common-nv = { path = "../common" }
 nccl.workspace = true
 log.workspace = true
 half.workspace = true
+itertools.workspace = true
 
 [dev-dependencies]
 simple_logger = "4.3"