test_gpu_scripts/modules/nccl_test.py
zulifeng fc97a768cf feat: 按 H100 生产验收标准更新测试指标与判定逻辑
- gpu_specs: H100 新增 compute_pass_thresholds_tflops 字段
  (fp32:54 / tf32:444 / fp16:734 / bf16:745 / fp8:1400),
  与 marketing peak 解耦,作为绝对 TFLOPS PASS 门槛
- benchmark: compute 结果中透出 pass_thresholds_tflops 供 report 使用
- report: compute 判定改用绝对 TFLOPS (PASS ≥门槛 / WARN ≥门槛×90% /
  FAIL <门槛×90%);表头切换为 Threshold 列;Memory D2D verdict
  由 50/30 收紧至 80/60;无阈值配置的 GPU 保留旧 % 效率逻辑
- nccl: _OP_BW_FRACTIONS 收紧至 AllReduce/AllGather/ReduceScatter
  0.45、Broadcast/SendRecv 0.40、AllToAll 0.35,与验收文档 §5 一致
- configs: benchmark 默认 matrix_size 4096→8192、warmup 10→50、
  iterations 100→500、use_compile 改 true;health temp_warning
  80→75、temp_critical 90→85,匹配生产验收稳态温度要求

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-13 14:52:41 +08:00

460 lines
18 KiB
Python

"""NCCL multi-GPU communication test — wraps official nccl-tests."""
import glob
import os
import re
import shutil
import subprocess
from datetime import datetime
from typing import Optional
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from modules.gpu_specs import detect_gpu_type, get_gpu_specs, resolve_tools_dir
TORCH_AVAILABLE = False
try:
import torch
if torch.cuda.is_available():
TORCH_AVAILABLE = True
except ImportError:
pass
# Per-operation bandwidth thresholds, as a fraction of NVLink bidirectional BW.
# Values aligned with the H100 production acceptance criteria (acceptance doc §5).
# AllToAll runs ~10-20% lower than AllReduce on 8-GPU NVSwitch, so its fraction is
# set lower; broadcast/sendrecv sit between.
_OP_BW_FRACTIONS = {
"allreduce": 0.45,
"allgather": 0.45,
"reducescatter": 0.45,
"broadcast": 0.40,
"sendrecv": 0.40,
"alltoall": 0.35,
}
class NCCLTest:
def __init__(self, config: dict):
self.config = config
self.console = Console()
self.nccl_cfg = config.get("nccl", {})
self.tools_dir = resolve_tools_dir(config)
self.gpu_type = detect_gpu_type()
self.specs = get_gpu_specs(self.gpu_type)
def _find_nccl_test(self, name: str) -> Optional[str]:
p = shutil.which(name)
if p:
return p
build_dir = os.path.join(self.tools_dir, "nccl-tests", "build")
local = os.path.join(build_dir, name)
if os.path.isfile(local) and os.access(local, shutil.os.X_OK):
return local
matches = glob.glob(os.path.join(self.tools_dir, "nccl-tests", "**", name), recursive=True)
for m in matches:
if os.access(m, shutil.os.X_OK):
return m
return None
def _find_mpirun(self) -> Optional[str]:
for cmd in ["mpirun", "mpiexec", os.path.join(self.tools_dir, "mpi", "bin", "mpirun")]:
p = shutil.which(cmd)
if p:
return p
return None
def run(self) -> dict:
gpu_count = 0
if TORCH_AVAILABLE:
gpu_count = torch.cuda.device_count()
if gpu_count < 2:
self.console.print(f"[yellow]NCCL test requires at least 2 GPUs (found {gpu_count})[/yellow]")
return {"error": "need_at_least_2_gpus", "gpu_count": gpu_count}
tests = []
if self.nccl_cfg.get("test_allreduce", True):
tests.append(("all_reduce_perf", "AllReduce"))
if self.nccl_cfg.get("test_alltoall", True):
tests.append(("alltoall_perf", "AllToAll"))
if self.nccl_cfg.get("test_broadcast", True):
tests.append(("broadcast_perf", "Broadcast"))
if self.nccl_cfg.get("test_reduce_scatter", False):
tests.append(("reduce_scatter_perf", "ReduceScatter"))
if self.nccl_cfg.get("test_allgather", False):
tests.append(("allgather_perf", "AllGather"))
if self.nccl_cfg.get("test_sendrecv", False):
tests.append(("sendrecv_perf", "SendRecv"))
nvlink_bw = self.specs.get("nvlink_bandwidth_gbps", 0)
# User-provided override applies uniformly across all ops; otherwise
# each op gets its own threshold from _OP_BW_FRACTIONS.
user_override = self.nccl_cfg.get("min_bandwidth_gbps")
def threshold_for(label: str) -> float:
if user_override:
return float(user_override)
if nvlink_bw <= 0:
return 10.0 # conservative floor
frac = _OP_BW_FRACTIONS.get(label.lower(), 0.45)
return round(nvlink_bw * frac)
if self.gpu_type == "unknown":
self.console.print("[yellow]Unknown GPU — using conservative bandwidth thresholds[/yellow]")
# Strategy: try nccl-tests binary directly (single-node, -g N),
# then mpirun, then torchrun fallback
results = {}
any_binary_worked = False
with Progress(
SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
TimeElapsedColumn(), console=self.console,
) as progress:
task = progress.add_task("NCCL tests...", total=len(tests))
for binary, label in tests:
progress.update(task, description=f"NCCL {label}...")
op_min_bw = threshold_for(label)
result = self._run_one_nccl_test_direct(
binary, label, gpu_count, op_min_bw
)
if result.get("status") not in ("SKIP", None) and "error" not in result:
any_binary_worked = True
results[label.lower()] = result
else:
# Try mpirun fallback
mpirun = self._find_mpirun()
if mpirun:
result = self._run_one_nccl_test_mpirun(
binary, label, gpu_count, mpirun, op_min_bw
)
if result.get("status") not in ("SKIP", None) and "error" not in result:
any_binary_worked = True
results[label.lower()] = result
progress.advance(task)
if not any_binary_worked:
self.console.print("[yellow]nccl-tests binaries failed, falling back to torchrun[/yellow]")
return self._run_torchrun_fallback(gpu_count)
all_passed = all(
r.get("status") == "PASS"
for r in results.values()
if isinstance(r, dict) and "status" in r
)
return {
"passed": all_passed,
"source": "nccl-tests",
"min_bandwidth_gbps": {
lbl.lower(): threshold_for(lbl) for _, lbl in tests
},
"tests": results,
"gpu_count": gpu_count,
"timestamp": datetime.now().isoformat(),
"detected_gpu_type": self.gpu_type,
}
def _run_one_nccl_test_direct(self, binary_name: str, label: str,
gpu_count: int, min_bw: float) -> dict:
"""Run nccl-tests binary directly with -g N (no mpirun needed for single-node)."""
binary = self._find_nccl_test(binary_name)
if not binary:
return {"status": "SKIP", "error": f"{binary_name} not found"}
cmd = [
binary,
"-b", "8M",
"-e", "8G",
"-f", "2",
"-g", str(gpu_count),
"-w", "5",
"-n", "20",
]
try:
env = os.environ.copy()
env["NCCL_DEBUG"] = "WARN"
r = subprocess.run(cmd, capture_output=True, text=True, timeout=180, env=env)
combined = r.stdout + r.stderr
# Check for NCCL/CUDA compatibility errors
if "CUDA driver version is insufficient" in combined or \
"Test NCCL failure" in combined:
error_msg = "NCCL/CUDA driver version mismatch" \
if "CUDA driver version" in combined \
else "NCCL test failure (library incompatibility)"
return {"status": "FAIL", "error": error_msg}
if r.returncode != 0:
return {"status": "FAIL", "error": r.stderr[:300]}
return self._parse_nccl_output(r.stdout, min_bw)
except subprocess.TimeoutExpired:
return {"status": "FAIL", "error": "timeout"}
except Exception as e:
return {"status": "FAIL", "error": str(e)}
def _run_one_nccl_test_mpirun(self, binary_name: str, label: str,
gpu_count: int, mpirun: str, min_bw: float) -> dict:
"""Run nccl-tests via mpirun (multi-node or per-GPU-process mode)."""
binary = self._find_nccl_test(binary_name)
if not binary:
return {"status": "SKIP", "error": f"{binary_name} not found"}
cmd = [
mpirun,
"-np", str(gpu_count),
"--allow-run-as-root",
"-x", "NCCL_DEBUG=WARN",
"-x", "CUDA_VISIBLE_DEVICES=" + ",".join(str(i) for i in range(gpu_count)),
binary,
"-b", "8",
"-e", "256M",
"-f", "2",
"-g", "1",
"-w", "5",
"-n", "20",
]
try:
env = os.environ.copy()
env["NCCL_DEBUG"] = "WARN"
r = subprocess.run(cmd, capture_output=True, text=True, timeout=180, env=env)
combined = r.stdout + r.stderr
if "CUDA driver version is insufficient" in combined or \
"Test NCCL failure" in combined:
error_msg = "NCCL/CUDA driver version mismatch" \
if "CUDA driver version" in combined \
else "NCCL test failure (library incompatibility)"
return {"status": "FAIL", "error": error_msg}
if r.returncode != 0:
return {"status": "FAIL", "error": r.stderr[:300]}
return self._parse_nccl_output(r.stdout, min_bw)
except subprocess.TimeoutExpired:
return {"status": "FAIL", "error": "timeout"}
except Exception as e:
return {"status": "FAIL", "error": str(e)}
@staticmethod
def _parse_nccl_output(stdout: str, min_bw: float) -> dict:
"""Parse nccl-tests tabular output and extract bandwidth results."""
best_algbw = 0.0
best_busbw = 0.0
size_results = []
for line in stdout.split("\n"):
line = line.strip()
if not line or line.startswith("#"):
continue
parts = line.split()
# nccl-tests data lines: size count type redop root time algbw busbw #wrong [time algbw busbw #wrong]
if len(parts) >= 9:
try:
size = int(parts[0])
# parts[2] is dtype string ('float'/'int32'/etc.), not a number
# out-of-place columns: time=parts[5], algbw=parts[6], busbw=parts[7]
time_us = float(parts[5])
algbw = float(parts[6])
busbw = float(parts[7])
size_results.append({
"size": size,
"time_us": time_us,
"algbw_gbps": algbw,
"busbw_gbps": busbw,
})
if busbw > best_busbw:
best_busbw = busbw
if algbw > best_algbw:
best_algbw = algbw
except (ValueError, IndexError):
continue
status = "PASS" if best_busbw >= min_bw else "WARN"
return {
"status": status,
"best_algbw_gbps": round(best_algbw, 1),
"best_busbw_gbps": round(best_busbw, 1),
"min_required_gbps": min_bw,
"by_size": size_results[-5:] if size_results else [],
}
def _run_torchrun_fallback(self, gpu_count: int) -> dict:
"""Basic NCCL connectivity test via torchrun — verifies NCCL works but does not benchmark performance."""
self.console.print("[yellow]nccl-tests not available, running basic NCCL connectivity check[/yellow]")
code = f"""
import torch, torch.distributed as dist, os
os.environ.setdefault("MASTER_ADDR","127.0.0.1")
os.environ.setdefault("MASTER_PORT","29500")
rank=int(os.environ.get("LOCAL_RANK",0))
ws={gpu_count}
dist.init_process_group("nccl",rank=rank,world_size=ws)
torch.cuda.set_device(rank)
x=torch.randn(1024*1024,device=f"cuda:{{rank}}",dtype=torch.float32)
# Test AllReduce
try:
dist.all_reduce(x.clone())
if rank==0: print("allreduce:ok")
except Exception as e:
if rank==0: print(f"allreduce:fail:{{e}}")
# Test Broadcast
try:
dist.broadcast(x.clone(),src=0)
if rank==0: print("broadcast:ok")
except Exception as e:
if rank==0: print(f"broadcast:fail:{{e}}")
# Test AllGather
try:
tensor_list=[torch.empty_like(x) for _ in range(ws)]
dist.all_gather(tensor_list,x.clone())
if rank==0: print("allgather:ok")
except Exception as e:
if rank==0: print(f"allgather:fail:{{e}}")
# Test ReduceScatter
try:
chunks=list(x.chunk(ws))
output=torch.empty_like(chunks[0])
dist.reduce_scatter(output,chunks)
if rank==0: print("reducescatter:ok")
except Exception as e:
if rank==0: print(f"reducescatter:fail:{{e}}")
# Test AllToAll
try:
chunks=list(x.chunk(ws))
output_list=[torch.empty_like(c) for c in chunks]
dist.all_to_all(output_list,chunks)
if rank==0: print("alltoall:ok")
except Exception as e:
if rank==0: print(f"alltoall:fail:{{e}}")
dist.destroy_process_group()
"""
import tempfile
tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False)
tmp.write(code)
tmp.close()
try:
# Prefer torchrun from the same venv as the running Python
import sys
venv_torchrun = os.path.join(os.path.dirname(sys.executable), "torchrun")
torchrun_cmd = venv_torchrun if os.path.isfile(venv_torchrun) else "torchrun"
r = subprocess.run(
[torchrun_cmd, f"--nproc_per_node={gpu_count}", tmp.name],
capture_output=True, text=True, timeout=120,
env={**os.environ, "NCCL_DEBUG": "WARN"},
)
os.unlink(tmp.name)
# Parse connectivity results — format: op_name:ok or op_name:fail:error
tests = {}
all_passed = True
for line in r.stdout.split("\n"):
line = line.strip()
if not line:
continue
parts = line.split(":")
op_name = parts[0]
result = parts[1] if len(parts) > 1 else "unknown"
if result == "ok":
status = "PASS"
else:
status = "FAIL"
all_passed = False
tests[op_name] = {
"status": status,
"error": ":".join(parts[2:]) if len(parts) > 2 and result == "fail" else None,
}
return {
"passed": all_passed,
"source": "torchrun_fallback",
"tests": tests,
"gpu_count": gpu_count,
}
except Exception as e:
return {"passed": False, "source": "torchrun_fallback", "error": str(e)}
@staticmethod
def print_results(results: dict, console: Console = None):
c = console or Console()
if "error" in results:
c.print(f"[bold red]Error: {results['error']}[/bold red]")
return
passed = results.get("passed", False)
source = results.get("source", "unknown")
if source == "torchrun_fallback":
# Connectivity check mode
verdict = "[bold green]✓ NCCL Connectivity OK[/bold green]" if passed else "[bold red]✗ NCCL Connectivity FAILED[/bold red]"
c.print(f"{verdict} [dim](basic check via torchrun)[/dim]")
tests = results.get("tests", {})
if tests:
c.print("\n[dim]Operations tested:[/dim]")
for op_name, result in tests.items():
if not isinstance(result, dict):
continue
status = result.get("status", "FAIL")
s_color = "green" if status == "PASS" else "red"
error = result.get("error")
if error:
c.print(f" [{s_color}]{op_name}[/{s_color}] — {error}")
else:
c.print(f" [{s_color}]{op_name}[/{s_color}]")
c.print("\n[yellow]Note: functional connectivity test only (no performance data)[/yellow]")
else:
# nccl-tests mode
verdict = "[bold green]✓ NCCL tests PASSED[/bold green]" if passed else "[bold yellow]⚠ NCCL tests WARNING[/bold yellow]"
c.print(f"{verdict} [dim](via {source})[/dim]")
tests = results.get("tests", {})
for op_name, result in tests.items():
if not isinstance(result, dict):
continue
c.print(f"\n[bold cyan]{op_name.upper()}[/bold cyan]")
status = result.get("status", "FAIL")
s_color = "green" if status == "PASS" else ("yellow" if status == "WARN" else "red")
c.print(f" Status: [{s_color}]{status}[/{s_color}] "
f"Best bus BW: {result.get('best_busbw_gbps', 'N/A')} GB/s "
f"(min: {result.get('min_required_gbps', 'N/A')} GB/s)")
by_size = result.get("by_size", [])
if by_size:
t = Table(box=None, padding=(0, 1))
t.add_column("Size", style="bold", justify="right")
t.add_column("Time (us)", justify="right")
t.add_column("Alg BW (GB/s)", justify="right")
t.add_column("Bus BW (GB/s)", justify="right")
for r in by_size:
sz = r.get("size", 0)
sz_str = f"{sz/1024:.0f}K" if sz < 1048576 else f"{sz/1048576:.0f}M"
t.add_row(sz_str, f"{r.get('time_us',0):.1f}",
f"{r.get('algbw_gbps',0):.1f}", f"{r.get('busbw_gbps',0):.1f}")
c.print(t)