Explorar o código

feat: auto-scale TP when recipe runs on larger cluster

When a recipe's tensor_parallel_size == required_nodes (the common
pattern), auto-scale TP to match the actual cluster size.  Same approach
as pipeline_parallel_size (auto-scales when 0).

- Added tp_size parameter to build_docker_command
- In start_cluster: if TP == required_nodes, scale TP to nnodes
- This way deepseek-v4-flash (TP=2, req=2) auto-scales to TP=4
  on a 4-node cluster
Alexander Korolev hai 4 meses
pai
achega
04413a8991
Modificáronse 2 ficheiros con 35 adicións e 3 borrados
  1. 9 2
      src/docker/mod.rs
  2. 26 1
      src/docker/tests.rs

+ 9 - 2
src/docker/mod.rs

@@ -251,6 +251,11 @@ async fn start_cluster(
     } else {
         recipe.pipeline_parallel_size.unwrap()
     };
+    let tp_size = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
+        nnodes
+    } else {
+        recipe.tensor_parallel_size
+    };
     let mut handles = Vec::new();
     for node in &nodes {
         let node = (*node).clone();
@@ -280,6 +285,7 @@ async fn start_cluster(
                 &node,
                 is_head,
                 pp_size,
+                tp_size,
                 nnodes,
                 &cd_s,
                 keep,
@@ -321,6 +327,7 @@ fn build_docker_command(
     node: &Node,
     is_head: bool,
     pp_size: u32,
+    tp_size: u32,
     nnodes: u32,
     cache_dir: &str,
     keep: bool,
@@ -416,7 +423,7 @@ fn build_docker_command(
     let mut flags: Vec<String> = match recipe.distributed_backend.as_str() {
         "ray" | "mp" if is_sglang => vec![
             "--tp".into(),
-            recipe.tensor_parallel_size.to_string(),
+            tp_size.to_string(),
             "--nnodes".into(),
             nnodes.to_string(),
             "--node-rank".into(),
@@ -427,7 +434,7 @@ fn build_docker_command(
         "ray" | "mp" => {
             let mut fp = vec![
                 "--tensor-parallel-size".into(),
-                recipe.tensor_parallel_size.to_string(),
+                tp_size.to_string(),
                 "--pipeline-parallel-size".into(),
                 pp_size.to_string(),
                 "--nnodes".into(),

+ 26 - 1
src/docker/tests.rs

@@ -57,6 +57,11 @@ fn trec() -> Recipe {
 }
 
 fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
+    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
+        2
+    } else {
+        recipe.tensor_parallel_size
+    };
     build_docker_command(
         &tcls(),
         &tnccl(),
@@ -65,6 +70,7 @@ fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
         node,
         head,
         2,
+        tp,
         2,
         "/home/user/.cache/huggingface",
         false,
@@ -74,6 +80,11 @@ fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
 }
 
 fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
+    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
+        2
+    } else {
+        recipe.tensor_parallel_size
+    };
     build_docker_command(
         &tcls(),
         &tnccl(),
@@ -82,6 +93,7 @@ fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
         node,
         head,
         2,
+        tp,
         2,
         "/home/user/.cache/huggingface",
         true,
@@ -237,6 +249,11 @@ fn sglang_recipe() -> Recipe {
 }
 
 fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
+    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
+        4
+    } else {
+        recipe.tensor_parallel_size
+    };
     build_docker_command(
         &tcls(),
         &tnccl(),
@@ -245,6 +262,7 @@ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
         node,
         head,
         4,
+        tp,
         4,
         "/home/user/.cache/huggingface",
         false,
@@ -372,14 +390,21 @@ fn recipe_env_overrides_cluster_defaults() {
 fn custom_init_port() {
     // Non-default init_port should propagate to vLLM master-port and
     // SGLang dist-init-addr.
+    let r = trec();
+    let tp = if r.required_nodes == Some(r.tensor_parallel_size) {
+        2
+    } else {
+        r.tensor_parallel_size
+    };
     let cmd = build_docker_command(
         &tcls(),
         &tnccl(),
-        &trec(),
+        &r,
         "r",
         &tcls().nodes[0],
         true,
         2,
+        tp,
         2,
         "/home/user/.cache/huggingface",
         false,