Răsfoiți Sursa

feat: extend NcclConfig with per-cluster NCCL overrides

- Add gloo_socket_ifname, ucx_net_devices, nccl_cross_nic, nccl_proto,
  nccl_p2p_level fields to NcclConfig for per-cluster NCCL tuning
- Use gloo_socket_ifname for GLOO/TP/MN/OMPI interfaces (falls back to
  socket_ifname when unset)
- Support ucx_net_devices override (e.g. "rocep1s0f0:1,roceP2p1s0f0:1")
- Support optional NCCL_CROSS_NIC, NCCL_PROTO, NCCL_P2P_LEVEL env vars
Alexander Korolev 3 luni în urmă
părinte
comite
ff23a80a81
4 a modificat fișierele cu 181 adăugiri și 51 ștergeri
  1. 123 0
      recipes/glm52-int4-int8/recipe.toml
  2. 27 27
      src/config.rs
  3. 26 24
      src/docker/mod.rs
  4. 5 0
      src/docker/tests.rs

+ 123 - 0
recipes/glm52-int4-int8/recipe.toml

@@ -0,0 +1,123 @@
+# GLM-5.2 QuantTrio Int4-Int8Mix — 4× DGX Spark / GB10 (sm_121), TP=4.
+# MTP k=4, cudagraph FULL, prefix caching, fp8_ds_mla KV, 350k ctx.
+#
+# Model: QuantTrio/GLM-5.2-Int4-Int8Mix (mixed Int4/Int8, 10% expert-pruned)
+#   → https://huggingface.co/QuantTrio/GLM-5.2-Int4-Int8Mix
+# MTP draft: CosmicRaisins/GLM-5.2-MTP-INT4-aligned
+#   → https://huggingface.co/CosmicRaisins/GLM-5.2-MTP-INT4-aligned
+#
+# BUILD THE IMAGE FIRST:
+#   Reference: https://github.com/CosmicRaisins/glm-5.2-gb10
+#   The custom image vllm-node-tf5-glm52-b12x:probe is built from eugr/spark-vllm-docker
+#   fork at pinned ref ab666069935c1f23e8ef56038b4659ac9e8f19f8, then applies mods
+#   glm52-sm12x-sparse and glm52-b12x-sparse on top. See the repo's bootstrap.sh
+#   for the full build-and-copy workflow.
+#
+#   1) git clone https://github.com/eugr/spark-vllm-docker
+#   2) cd spark-vllm-docker && ./build-and-copy.sh \
+#        --vllm-ref ab666069935c1f23e8ef56038b4659ac9e8f19f8 \
+#        -t vllm-node-tf5-glm52-b12x:probe --tf5
+#   3) Bake mods glm52-sm12x-sparse + glm52-b12x-sparse into the image
+#      (see Dockerfile.glm52-consolidated in the repo)
+#
+# PREREQUISITES:
+#   - NCCL 2.30.4 staged on every node at /root/.cache/huggingface/hub/nccl-2.30.4/libnccl.so.2
+#     (the image LD_PRELOADs it). Get it with:
+#       pip install --no-deps -d /tmp/ncclwheel nvidia-nccl-cu13==2.30.4 && mkdir -p ... && unzip ...
+#   - Weights cached on every node:
+#       huggingface-cli download QuantTrio/GLM-5.2-Int4-Int8Mix
+#       huggingface-cli download CosmicRaisins/GLM-5.2-MTP-INT4-aligned
+#
+# TUNING NOTES (from community):
+#   - gpu-memory-utilization 0.92 is the sweet spot for 350k ctx on 4× GB10 with MTP
+#   - MTP k=4 is more reliable on startup than k=5 on GB10
+#   - cudagraph FULL requires b12x installed; without it, use PIECEWISE
+
+engine = "vllm"
+name = "GLM-5.2 Int4-Int8Mix (vLLM)"
+docker_image = "vllm-node-tf5-glm52-b12x:probe"
+model = "QuantTrio/GLM-5.2-Int4-Int8Mix"
+port = 8210
+required_nodes = 4
+tensor_parallel_size = 4
+pipeline_parallel_size = 1
+distributed_backend = "mp"
+shm_size = "64gb"
+
+entrypoint = ["vllm", "serve"]
+
+args = [
+    # Reasoning & tools
+    "--reasoning-parser", "glm45",
+    "--tool-call-parser", "glm47",
+    "--enable-auto-tool-choice",
+    # KV cache & memory
+    "--kv-cache-dtype", "fp8_ds_mla",
+    "--max-model-len", "262144",
+    "--max-num-seqs", "1",
+    "--max-num-batched-tokens", "4096",
+    "--gpu-memory-utilization", "0.92",
+    # Compilation (b12x required for FULL cudagraph)
+    "--compilation-config", """
+'{"cudagraph_mode":"FULL"}'""",
+    # MTP speculative decoding
+    "--speculative-config", """
+'{"model":"/root/.cache/huggingface/hub/glm52-mtp-int4-aligned","method":"mtp","num_speculative_tokens":4,"attention_backend":"FLASHMLA_SPARSE"}'""",
+    # Other
+    "--trust-remote-code",
+    "--enable-prefix-caching",
+    "--enable-flashinfer-autotune",
+    "--served-model-name", "glm52",
+]
+
+[env]
+# Timeout safety net for cold boot (cudagraph + JIT)
+VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS = "1800"
+# NCCL shm_broadcast fix
+LD_PRELOAD = "/root/.cache/huggingface/hub/nccl-2.30.4/libnccl.so.2"
+
+# Offline mode — weights must be cached on each node
+HF_HOME = "/root/.cache/huggingface"
+TRITON_CACHE_DIR = "/root/.cache/huggingface/.tritoncache"
+HF_HUB_OFFLINE = "1"
+TRANSFORMERS_OFFLINE = "1"
+
+# GLM-5.2 sparse-MLA suite (these guard the Triton fallback paths on sm_121)
+VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"
+VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"
+GLM52_BIND_HOST_TRITON = "1"
+GLM52_MQA_LOGITS_TRITON = "1"
+GLM52_PAGED_MQA_TRITON = "1"
+GLM52_PAGED_MQA_TOPK_CHUNK_SIZE = "8192"
+GLM52_B12X_MLA = "1"
+TORCH_CUDA_ARCH_LIST = "12.1a"
+
+# NCCL / RoCE (launcher auto-sets most; recipe env wins last)
+NCCL_NET = "IB"
+NCCL_IB_DISABLE = "0"
+NCCL_IB_GID_INDEX = "3"
+NCCL_CROSS_NIC = "1"
+NCCL_CUMEM_ENABLE = "0"
+NCCL_IGNORE_CPU_AFFINITY = "1"
+NCCL_DEBUG = "WARN"
+
+# MoE & compute
+OMP_NUM_THREADS = "8"
+
+[mods]
+commands = [
+    # Copy Triton sparse-MLA kernel files into the vLLM tree inside the container.
+    # These are mounted via the mods system at /mods/glm52-sm12x-sparse/ and copied
+    # into the vLLM Python package paths before the engine starts.
+    "cp /mods/glm52-sm12x-sparse/sparse_mla_kernels.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
+    "cp /mods/glm52-sm12x-sparse/sparse_mla_env.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
+    "cp /mods/glm52-sm12x-sparse/sm12x_sparse_mla_attn.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
+    "cp /mods/glm52-sm12x-sparse/patch_flashmla_ops.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
+    "cp /mods/glm52-sm12x-sparse/flashmla_sparse.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
+    "cp /mods/glm52-sm12x-sparse/sm12x_deep_gemm_fallbacks.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
+    "cp /mods/glm52-sm12x-sparse/sm12x_mqa.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
+    "cp /mods/glm52-sm12x-sparse/b12x_sparse_helpers.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
+    # upstream vLLM #46862: fused indexer Q rope+fp8-quant
+    "cp /mods/glm52-sm12x-sparse/sparse_attn_indexer.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/",
+    "cp /mods/glm52-sm12x-sparse/deepseek_v2.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/",
+]

+ 27 - 27
src/config.rs

@@ -3,8 +3,7 @@ pub use crate::pre_serve::{
 };
 use anyhow::{Context, Result};
 use serde::Deserialize;
-use std::collections::HashMap;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use tracing::info;
 #[derive(Clone, Deserialize)]
 pub struct Node {
@@ -23,17 +22,26 @@ pub struct NcclConfig {
     pub socket_ifname: String,
     pub ib_gid_index: u32,
     pub ib_disable: u32,
+    #[serde(default)]
+    pub gloo_socket_ifname: Option<String>,
+    #[serde(default)]
+    pub ucx_net_devices: Option<String>,
+    #[serde(default)]
+    pub nccl_cross_nic: Option<u32>,
+    #[serde(default)]
+    pub nccl_proto: Option<String>,
+    #[serde(default)]
+    pub nccl_p2p_level: Option<String>,
 }
 #[derive(Clone, Deserialize)]
 pub struct Cluster {
     pub name: String,
     pub master_ib_ip: String,
     pub nodes: Vec<Node>,
-    /// Extra Home Assistant switches to toggle alongside the nodes during power on/off.
+    /// Extra HA switches.
     #[serde(default)]
     pub extra_entities: Vec<String>,
 }
-/// Top-level config (clusters + per-cluster pre-serve commands).
 #[derive(Clone, Deserialize)]
 pub struct Config {
     pub ssh_user: String,
@@ -43,23 +51,21 @@ pub struct Config {
     pub nccl: HashMap<String, NcclConfig>,
     #[serde(default)]
     pub hf_cache_dir: Option<String>,
-    /// TCP port for cluster init (vLLM/SGLang)
+    /// Init port
     #[serde(default = "default_init_port")]
     pub init_port: u16,
-    /// Per-cluster pre-serve commands, command aliases, HA config
+    /// Pre-serve, aliases, HA
     #[serde(default)]
-    pub before_serve: HashMap<String, BeforeServeTable>, // <-- move here
+    pub before_serve: HashMap<String, BeforeServeTable>,
     #[serde(default)]
     pub exec_aliases: HashMap<String, String>,
     #[serde(default)]
     pub home_assistant: Option<HomeAssistantConfig>,
-    /// Global mods directory — script files merged with recipe mods/ into
-    /// the container's /mods/.  Recipe-level [mods] commands explicitly
-    /// reference files from either source via /mods/<file>.
+    /// Global mods dir.
     #[serde(default)]
     pub global_mods_dir: Option<std::path::PathBuf>,
 }
-/// Home Assistant config for `power on`/`power off`.
+/// Home Assistant config.
 #[derive(Clone, Deserialize)]
 pub struct HomeAssistantConfig {
     pub url: String,
@@ -116,18 +122,17 @@ pub struct Recipe {
     #[serde(default)]
     pub entrypoint: Vec<String>,
     pub distributed_backend: Option<DistributedBackend>,
-    /// Shared memory size for the container (e.g. "64gb"). Omitted when empty.
+    /// Shm size
     #[serde(default = "default_shm_size")]
     pub shm_size: String,
-    /// Recipe-level pre-serve commands (appended after cluster-level).
+    /// Recipe pre-serve.
     #[serde(default)]
     pub before_serve: Vec<PreServeCommand>,
-    /// Mods: commands that run *inside* the container via entrypoint wrapping
-    /// before the engine process starts. Requires `entrypoint` to be set.
+    /// Mods before engine.
     #[serde(default)]
     pub mods: ModsTable,
 }
-/// Resolves cluster by exact name, then unique prefix. Errors if none or ambiguous.
+/// Resolve cluster.
 fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
     if let Some(c) = clusters.iter().find(|c| c.name == name) {
         return Ok(c);
@@ -160,8 +165,7 @@ fn default_shm_size() -> String {
 fn default_tensor_parallel_size() -> u32 {
     1
 }
-/// Selects a cluster + NCCL config by name or auto-picks the smallest
-/// cluster with NCCL config and enough free nodes.
+/// Select cluster.
 pub fn select_cluster<'c>(
     config: &'c Config,
     cluster_name: &str,
@@ -202,11 +206,9 @@ pub fn select_cluster<'c>(
     }
     anyhow::bail!(fluent_i18n::t!("cluster-auto-fail"));
 }
-/// Looks up a cluster by name (or unique prefix) without requiring NCCL config.
 pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
     Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
 }
-/// Resolves the cluster for a recipe, applying recipe-level cluster pinning.
 pub fn resolve_cluster<'c>(
     config: &'c Config,
     recipe: &Recipe,
@@ -222,8 +224,6 @@ pub fn resolve_cluster<'c>(
     };
     select_cluster(config, &choice, recipe.required_nodes, occupied)
 }
-/// Resolve the path to a recipe's `mods/` directory, if it exists.
-/// Checks cwd first, then the default config dir.
 pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
     let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods"));
     if cwd.is_dir() {
@@ -235,8 +235,6 @@ pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
         .join("mods");
     cfg.is_dir().then_some(cfg)
 }
-/// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
-/// falling back to cwd for the error).
 pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
     let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
     if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
@@ -248,12 +246,9 @@ pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
             return Ok(toml::from_str(&content)?);
         }
     }
-    // Neither found — let the cwd path produce a clean error.
     let content = tokio::fs::read_to_string(&cwd_path).await?;
     Ok(toml::from_str(&content)?)
 }
-/// OS-appropriate default config dir (e.g. `~/.config/sm121-launcher`),
-/// or `None` for sandboxes with no known config dir.
 pub fn default_config_dir() -> Option<std::path::PathBuf> {
     dirs::config_dir().map(|p| p.join("sm121-launcher"))
 }
@@ -305,6 +300,11 @@ mod tests {
             socket_ifname: "enp1s0f0np0".to_string(),
             ib_gid_index: 3,
             ib_disable: 0,
+            gloo_socket_ifname: None,
+            ucx_net_devices: None,
+            nccl_cross_nic: None,
+            nccl_proto: None,
+            nccl_p2p_level: None,
         }
     }
     #[test]

+ 26 - 24
src/docker/mod.rs

@@ -14,7 +14,6 @@ use std::borrow::Cow;
 use std::collections::HashSet;
 use tokio::time::{Duration, sleep};
 use tracing::info;
-
 /// Options for the serve operation.
 #[derive(Clone, Debug)]
 pub struct ServeOpts {
@@ -23,7 +22,6 @@ pub struct ServeOpts {
     pub dry_run: bool,
     pub auto_power_on: bool,
 }
-
 pub(crate) fn shq(s: &str) -> String {
     let mut out = String::with_capacity(s.len() + 2);
     out.push('\'');
@@ -37,7 +35,6 @@ pub(crate) fn shq(s: &str) -> String {
     out.push('\'');
     out
 }
-
 pub async fn serve(
     config: &Config,
     recipe_name: &str,
@@ -52,7 +49,6 @@ pub async fn serve(
         "{}",
         fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
     );
-
     // Check if cluster switch is off and power on if requested
     if opts.auto_power_on {
         crate::cluster_power::power_on_cluster_if_off(
@@ -64,7 +60,6 @@ pub async fn serve(
         )
         .await?;
     }
-
     let (cluster, nccl): (Cluster, NcclConfig) = if let Some(nn) = node_name {
         let (solo_c, solo_n) = config
             .clusters
@@ -86,6 +81,11 @@ pub async fn serve(
                             socket_ifname: "lo".to_string(),
                             ib_gid_index: 3,
                             ib_disable: 1,
+                            gloo_socket_ifname: None,
+                            ucx_net_devices: None,
+                            nccl_cross_nic: None,
+                            nccl_proto: None,
+                            nccl_p2p_level: None,
                         });
                     (solo, nc)
                 })
@@ -124,7 +124,6 @@ pub async fn serve(
         opts.dry_run,
     )
     .await?;
-
     let mods_remote_path = mods::upload_mods_to_cluster(
         config,
         &cluster,
@@ -134,7 +133,6 @@ pub async fn serve(
         opts.dry_run,
     )
     .await?;
-
     let image = try_repaired_image(config, &cluster, &recipe).await?;
     let mut recipe = recipe;
     if image != recipe.docker_image {
@@ -204,7 +202,6 @@ pub async fn stop_cluster(
     info!("{}", fluent_i18n::t!("cluster-all-stopped"));
     Ok(())
 }
-
 /// Find the cluster that has running containers for a recipe by reading
 /// the `sm121.cluster` label from live containers.  Used by stop/logs
 /// when `--cluster auto`.
@@ -236,7 +233,6 @@ pub async fn resolve_cluster_by_container_label<'c>(
     }
     anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
 }
-
 /// Query running containers across all unique nodes and return the set of
 /// occupied node names.  A node is occupied if any sm121 container is
 /// running on it.
@@ -269,7 +265,6 @@ pub async fn occupied_nodes(config: &Config) -> Result<HashSet<String>> {
     }
     Ok(occupied)
 }
-
 /// Find the cluster that has running containers for a recipe by reading
 /// the `sm121.cluster` label from live containers.  Used by stop/logs
 /// when `--cluster auto`.
@@ -280,7 +275,6 @@ struct StartClusterOpts {
     mods_remote_path: String,
     dry_run: bool,
 }
-
 /// Shared context for [`build_docker_command`] — constructed once, cloned per node.
 #[derive(Clone)]
 struct DockerCommandCtx {
@@ -296,7 +290,6 @@ struct DockerCommandCtx {
     init_port: u16,
     mods_remote_path: String,
 }
-
 async fn start_cluster(
     config: &Config,
     cluster: &Cluster,
@@ -406,7 +399,6 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
         cmd.push(format!("-v {}:/mods:ro", ctx.mods_remote_path));
     }
-
     let mut env = vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
     if ctx.nnodes > 1 {
         let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
@@ -415,25 +407,37 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
         } else {
             env.push(("VLLM_HOST_IP", worker_ip));
         }
-        let nccl_socket_ifname = ctx.nccl.socket_ifname.as_str();
-        env.extend([
+        let sock = ctx.nccl.socket_ifname.as_str();
+        let gloo_if = ctx.nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
+        let ucx_dev = ctx.nccl.ucx_net_devices.as_deref().unwrap_or(sock);
+        let mut ncclext: Vec<(&str, Cow<str>)> = vec![
             ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
             ("NCCL_IB_HCA", ctx.nccl.ib_hca.as_str().into()),
-            ("NCCL_SOCKET_IFNAME", nccl_socket_ifname.into()),
-            ("GLOO_SOCKET_IFNAME", nccl_socket_ifname.into()),
-            ("TP_SOCKET_IFNAME", nccl_socket_ifname.into()),
+            ("NCCL_SOCKET_IFNAME", sock.into()),
+            ("GLOO_SOCKET_IFNAME", gloo_if.into()),
+            ("TP_SOCKET_IFNAME", gloo_if.into()),
             (
                 "NCCL_IB_GID_INDEX",
                 ctx.nccl.ib_gid_index.to_string().into(),
             ),
-            ("MN_IF_NAME", nccl_socket_ifname.into()),
-            ("UCX_NET_DEVICES", nccl_socket_ifname.into()),
-            ("OMPI_MCA_btl_tcp_if_include", nccl_socket_ifname.into()),
+            ("MN_IF_NAME", gloo_if.into()),
+            ("UCX_NET_DEVICES", ucx_dev.into()),
+            ("OMPI_MCA_btl_tcp_if_include", gloo_if.into()),
             ("OMP_NUM_THREADS", "4".into()),
             ("NCCL_IGNORE_CPU_AFFINITY", "1".into()),
             ("NCCL_NET", "IB".into()),
             ("NCCL_DEBUG", "WARN".into()),
-        ]);
+        ];
+        if let Some(v) = ctx.nccl.nccl_cross_nic {
+            ncclext.push(("NCCL_CROSS_NIC", v.to_string().into()));
+        }
+        if let Some(ref v) = ctx.nccl.nccl_proto {
+            ncclext.push(("NCCL_PROTO", v.as_str().into()));
+        }
+        if let Some(ref v) = ctx.nccl.nccl_p2p_level {
+            ncclext.push(("NCCL_P2P_LEVEL", v.as_str().into()));
+        }
+        env.extend(ncclext);
     }
     if is_head {
         env.push(("VLLM_LOGGING_COLOR", "1".into()));
@@ -462,7 +466,6 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     for (k, v) in env {
         cmd.push(format!("-e {k}={v}"));
     }
-
     let tctx = tail::TailCtx {
         recipe: &ctx.recipe,
         tp_size: ctx.tp_size,
@@ -486,6 +489,5 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     }
     cmd
 }
-
 #[cfg(test)]
 mod tests;

+ 5 - 0
src/docker/tests.rs

@@ -34,6 +34,11 @@ fn tnccl() -> NcclConfig {
         socket_ifname: "e".into(),
         ib_gid_index: 3,
         ib_disable: 0,
+        gloo_socket_ifname: None,
+        ucx_net_devices: None,
+        nccl_cross_nic: None,
+        nccl_proto: None,
+        nccl_p2p_level: None,
     }
 }