Bladeren bron

feat: set safe NCCL defaults (IGNORE_CPU_AFFINITY, NET, CROSS_NIC, DEBUG)

Adds four hard-coded env vars to the docker command, alongside the
existing `NCCL_CUMEM_ENABLE` and `OMP_NUM_THREADS`:

- `NCCL_IGNORE_CPU_AFFINITY=1` — GB10 has odd CPU affinity that
  trips NCCL warnings; ignore it (key for DGX Spark).
- `NCCL_NET=IB` — force NCCL to load `libnccl-net.so` instead of
  the default socket fallback.
- `NCCL_CROSS_NIC=1` — allow traffic across multiple RDMA NICs
  (multi-rail).
- `NCCL_DEBUG=WARN` — keep logs lean by default; override via
  recipe `[env]` to `INFO`/`TRACE` while debugging.

All four are overridable from the recipe's `[env]` table since recipe
env is appended after the defaults.

README updated to document each new var's purpose.
Alexander Korolev 4 maanden geleden
bovenliggende
commit
9c0e7680a4
3 gewijzigde bestanden met toevoegingen van 46 en 0 verwijderingen
  1. 4 0
      README.md
  2. 4 0
      src/docker/mod.rs
  3. 38 0
      src/docker/tests.rs

+ 4 - 0
README.md

@@ -240,6 +240,10 @@ The launcher passes the following env vars to every container.  Recipe-level
 | Env var | Value | Purpose |
 |---|---|---|
 | `NCCL_CUMEM_ENABLE` | `0` | Disable CUDA-managed memory in NCCL (use system allocator) |
+| `NCCL_IGNORE_CPU_AFFINITY` | `1` | Don't pin NCCL threads to a specific CPU set (needed on DGX Spark) |
+| `NCCL_NET` | `IB` | Force NCCL to use the InfiniBand plugin (`IB` = `libnccl-net.so`) |
+| `NCCL_CROSS_NIC` | `1` | Allow traffic across multiple RDMA NICs (multi-rail) |
+| `NCCL_DEBUG` | `WARN` | NCCL log level (override with `[env]` for `INFO`/`TRACE` while debugging) |
 | `OMP_NUM_THREADS` | `4` | OpenMP thread limit per process |
 | `HF_HOME` | `/cache/huggingface` | HuggingFace cache root (volume-mounted from `ssh_user`'s host cache) |
 

+ 4 - 0
src/docker/mod.rs

@@ -392,6 +392,10 @@ fn build_docker_command(
     // Default cluster env (recipe env below overrides these if duplicated)
     cmd.push("-e NCCL_CUMEM_ENABLE=0".to_string());
     cmd.push("-e OMP_NUM_THREADS=4".to_string());
+    cmd.push("-e NCCL_IGNORE_CPU_AFFINITY=1".to_string());
+    cmd.push("-e NCCL_NET=IB".to_string());
+    cmd.push("-e NCCL_CROSS_NIC=1".to_string());
+    cmd.push("-e NCCL_DEBUG=WARN".to_string());
 
     for (k, v) in &recipe.env {
         cmd.push(format!("-e {}={}", k, v));

+ 38 - 0
src/docker/tests.rs

@@ -127,6 +127,19 @@ fn mp() {
         cmd.contains("-e OMP_NUM_THREADS=4"),
         "default OMP_NUM_THREADS missing"
     );
+    assert!(
+        cmd.contains("-e NCCL_IGNORE_CPU_AFFINITY=1"),
+        "default NCCL_IGNORE_CPU_AFFINITY missing"
+    );
+    assert!(cmd.contains("-e NCCL_NET=IB"), "default NCCL_NET missing");
+    assert!(
+        cmd.contains("-e NCCL_CROSS_NIC=1"),
+        "default NCCL_CROSS_NIC missing"
+    );
+    assert!(
+        cmd.contains("-e NCCL_DEBUG=WARN"),
+        "default NCCL_DEBUG missing"
+    );
     // Head node: no --headless
     assert!(!cmd.contains("--headless"));
 }
@@ -282,6 +295,10 @@ fn recipe_env_overrides_cluster_defaults() {
     let mut r = trec();
     r.env.insert("NCCL_CUMEM_ENABLE".into(), "1".into());
     r.env.insert("OMP_NUM_THREADS".into(), "16".into());
+    r.env.insert("NCCL_IGNORE_CPU_AFFINITY".into(), "0".into());
+    r.env.insert("NCCL_NET".into(), "SOCKET".into());
+    r.env.insert("NCCL_CROSS_NIC".into(), "0".into());
+    r.env.insert("NCCL_DEBUG".into(), "INFO".into());
     let cmd = run(&r, &tcls().nodes[0], true);
     // Docker uses the last -e for duplicate keys, so recipe values come after.
     // rfind returns the last occurrence, which should be the recipe override.
@@ -296,6 +313,27 @@ fn recipe_env_overrides_cluster_defaults() {
         cmd[omp_pos..].starts_with("OMP_NUM_THREADS=16"),
         "recipe override should be last"
     );
+    // New NCCL defaults can also be overridden.
+    let ignore_pos = cmd.rfind("NCCL_IGNORE_CPU_AFFINITY=").unwrap();
+    assert!(
+        cmd[ignore_pos..].starts_with("NCCL_IGNORE_CPU_AFFINITY=0"),
+        "NCCL_IGNORE_CPU_AFFINITY should be overridable"
+    );
+    let net_pos = cmd.rfind("NCCL_NET=").unwrap();
+    assert!(
+        cmd[net_pos..].starts_with("NCCL_NET=SOCKET"),
+        "NCCL_NET should be overridable"
+    );
+    let cross_pos = cmd.rfind("NCCL_CROSS_NIC=").unwrap();
+    assert!(
+        cmd[cross_pos..].starts_with("NCCL_CROSS_NIC=0"),
+        "NCCL_CROSS_NIC should be overridable"
+    );
+    let debug_pos = cmd.rfind("NCCL_DEBUG=").unwrap();
+    assert!(
+        cmd[debug_pos..].starts_with("NCCL_DEBUG=INFO"),
+        "NCCL_DEBUG should be overridable"
+    );
 }
 
 #[test]