diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7f380f5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +# OhMyOpenCode 工作状态 +.omo/ diff --git a/src/code/issue4/.gitignore b/src/code/issue4/.gitignore new file mode 100644 index 0000000..ee44239 --- /dev/null +++ b/src/code/issue4/.gitignore @@ -0,0 +1,8 @@ +# 编译产物 +__pycache__/ +*.pyc +*.pyo + +# 测试输出 +*.json +.nccl_diag_checkpoint* diff --git "a/src/code/issue4/Analysis Report/issue4\345\210\206\346\236\220\346\212\245\345\221\212.docx" "b/src/code/issue4/Analysis Report/issue4\345\210\206\346\236\220\346\212\245\345\221\212.docx" deleted file mode 100644 index 6d79557..0000000 Binary files "a/src/code/issue4/Analysis Report/issue4\345\210\206\346\236\220\346\212\245\345\221\212.docx" and /dev/null differ diff --git a/src/code/issue4/README.md b/src/code/issue4/README.md new file mode 100644 index 0000000..312ea1a --- /dev/null +++ b/src/code/issue4/README.md @@ -0,0 +1,113 @@ +# NCCL 集合通信位级可复现性诊断工具 + +> 定位 AllReduce / Reduce-Scatter 中由 NCCL 算法/协议选择引入的逐位非确定性。 + +## 快速开始 + +```bash +cd hpn/src/code + +# 自检 — 验证环境是否能跑 NCCL + 管线是否正常 +PYTHONPATH=. python -m issue4.diagnose --self-test --nranks 4 + +# 扫描 — 遍历多种配置，输出报告 +PYTHONPATH=. python -m issue4.diagnose --mode standard --nranks 4 --json report.json + +# 复现 — 跨算法对比，定位非确定性根因 +PYTHONPATH=. python issue4/reproduce_case.py --algo Tree --dtype float16 --size 128M --nranks 4 +``` + +## CLI 参数 + +**diagnose.py** — 主入口： + +| 参数 | 作用 | 示例 | +|---|---|---| +| `--mode` | 扫描模式：`quick` / `standard` / `exhaustive` | `--mode standard` | +| `--nranks` | 参与通信的 GPU 数量 | `--nranks 4` | +| `--n-calls` | 单进程组内 NCCL 调用次数（>1 启用三级对比） | `--n-calls 5` | +| `--json` | 输出 JSON 报告路径 | `--json report.json` | +| `--list-configs` | 只打印配置矩阵，不执行 | `--list-configs` | +| `--inject-difference` | ULP 注入验证（`auto` 或 `offset=42,magnitude=3`） | `--inject-difference auto` | +| `--force` | 配合 `--inject-difference`：注入后继续完整扫描 | `--force` | +| `--self-test` | 跑一次最小 GPU 冒烟测试并退出 | `--self-test` | +| `--algo/--proto/--dtype/--size` | 限定单个配置（不扫矩阵） | `--algo Ring --size 128M` | + +**reproduce_case.py** — 复现脚本： + +| 参数 | 作用 | +|---|---| +| `--algo/--proto/--dtype/--size/--nranks` | 指定测试配置 | +| `--trials` | 重复运行次数（≥3 启用迭代演化检测） | +| `--collective` | `allreduce`（默认）或 `reducescatter` | + +## 扫描模式 + +| 模式 | 算法 | 协议 | 精度 | 数据规模 | 配置总数 | +|---|---|---|---|---|---| +| `quick` | Ring, Tree | Simple | fp32 | 16M, 128M | 4 | +| `standard` | Ring, Tree, PAT | LL, Simple | fp32, fp16 | 4K~128M（5 档） | 60 | +| `exhaustive` | Ring, Tree, PAT | LL, Simple | fp32, fp16, bf16 | 1K~128M（6 档） | 108（含 LL128 过滤后更少） | + +PAT 和 LL128 受硬件兼容性过滤，不支持时自动跳过。 + +## 硬件环境变量 + +4090D / 5090（无 NVLink / 无 P2P）需要在所有命令前加： + +```bash +NCCL_P2P_DISABLE=1 +``` + +4090 公版、A100、H100 等有 NVLink/P2P 的卡不需要。 + +## 验证结果（4×RTX 5090） + +``` +Case 1: Tree/Simple/float16 — run-to-run BITWISE IDENTICAL +Case 2: Ring vs Tree 跨算法比对: + Diff count: 23751354 / 67108864 (35.39%) + Max abs diff: 3.91e-03 + Distribution: uniform → 算法级根因 +``` + +## 目录结构 + +``` +issue4/ +├── diagnose.py 主入口 8 阶段流水线 +├── config_matrix.py 硬件检测 + 配置矩阵生成 +├── data_generator.py SHA-256 确定性数据生成 +├── runner.py NCCL 执行器 +├── comparator.py 逐位比对 + XOR/ULP + 三级对比 +├── reporter.py 报告输出 + 诊断建议引擎 +├── reproduce_case.py 独立复现脚本 +├── requirements.txt +└── docs/ 5 张 Mermaid 架构图 +``` + +## 硬件兼容性 + +| 算法 | 要求 | +|---|---| +| Ring, Tree, PAT | 所有 GPU | +| NVLS, NVLSTree | NVSwitch（H100 / B200） | +| CollnetDirect | NVSwitch + SHARP | +| CollnetChain | SHARP 网络 | + +| 协议 | 要求 | +|---|---| +| LL, Simple | 所有 GPU | +| LL128 | SM90+（Hopper / Blackwell）。非 Hopper 启用**静默数据损坏** | + +| GPU | P2P 可用 | LL128 | 注意事项 | +|---|---|---|---| +| 4090D / 5090 | ❌ 需 `NCCL_P2P_DISABLE=1` | ✅（仅 5090） | 5090 NCCL ≥ 2.27 | +| 4090 公版 | ✅ | ❌（SM89） | — | +| A100 / H100 | ✅ | ❌（仅 H100） | 有 NVSwitch，可跑 NVLS | + +## 参考 + +- NCCL#1975: AllReduce determinism +- NCCL#1055: Ring vs Tree precision on A100/A800 +- NCCL#157: Chunk partitioning and determinism diff --git a/src/code/issue4/Test Scripts/FlashOverlap/test.py b/src/code/issue4/Test Scripts/FlashOverlap/test.py deleted file mode 100644 index fa10386..0000000 --- a/src/code/issue4/Test Scripts/FlashOverlap/test.py +++ /dev/null @@ -1,374 +0,0 @@ -''' - Using multiprocessing for distributed running, - please specify the GPUs via CUDA_VISIBLE_DEVICES: - CUDA_VISIBLE_DEVICES=0,1 python3 result.py --m 4096 --n 8192 --k 4096 -''' - -import torch -import json -from pathlib import Path -import torch.multiprocessing as mp -import pandas as pd -import argparse -import os - -torch.ops.load_library("../build/lib/libst_pybinding.so") - -WARM_UP=20 -REP=200 - -def div_up(x: int, y: int): - return (x + y - 1) // y - -def reorder_indices(S, hint): - # Generate the original array of indices [0, 1, ..., S-1] - original = list(range(S)) - - # Create an empty list to store the new order of indices - new_order = [-1] * S - - # Place the indices of the hint list in the first positions of the new order - for i, element in enumerate(hint): - new_order[element] = i - - # Place the remaining indices in the new order - remaining_elements = [x for x in original if x not in hint] - for i, element in enumerate(remaining_elements, start=len(hint)): - new_order[element] = i - - return torch.tensor(new_order, dtype=torch.int, device="cuda") - -def generate_row_remap_array( - M, N, BM, BN, S_list, world_size, device="cuda" -): - total_tiles = (M * N) // (BM * BN) - assert sum(S_list) == total_tiles, "sum(S_list) must equal total number of tiles" - - original_row_ids = torch.arange(M * N // BN, dtype=torch.int, device=device) - reordered_row_id = torch.empty_like(original_row_ids) - - current_row = 0 - for S in S_list: - chunk_size = S * BM - chunk_row_ids = original_row_ids[current_row : current_row + chunk_size] - - # Compute row_id % world_size for the current chunk - mod_values = chunk_row_ids % world_size - - # Sort the chunk based on mod_values (stable sort) - _, sorted_indices = torch.sort(mod_values, stable=True) - reordered_chunk = chunk_row_ids[sorted_indices] - - reordered_row_id[current_row : current_row + chunk_size] = reordered_chunk - current_row += chunk_size - - # Compute remap: remap[original_row_id] = new_row_id - remap = torch.empty_like(original_row_ids) - remap[reordered_row_id] = torch.arange(len(reordered_row_id), dtype=torch.int, device=device) - - return remap - -def perf_running_process(rank, world_size, nccl_id, - M: int, N: int, K: int, - BM: int, BN: int, Algo: int, cSeg: list, hint: list, - comm_op: str, - result_dict): - - cSeg_CPU = torch.tensor(cSeg, dtype=torch.int32) - cSeg_GPU = cSeg_CPU.cuda(rank) - - TileNum = div_up(M, BM) * div_up(N, BN) - - torch.cuda.set_device(rank) - - gemm_class = torch.classes.flashoverlap_class.OverlapImpl() - - gemm_class.nccl_init(rank, world_size, nccl_id) - gemm_class.cutlass_init() - gemm_class.overlap_init() - - A = torch.empty((M, K), dtype=torch.float16, device="cuda").normal_(mean=0., std=0.5) - B = torch.empty((N, K), dtype=torch.float16, device="cuda").normal_(mean=0., std=0.5) - C = torch.empty((M, N), dtype=torch.float16, device="cuda") - - MonitoredMatrix = torch.zeros(((N+BN-1)//BN), dtype=torch.int, device="cuda") - ReorderedArray = reorder_indices(TileNum, hint).reshape(((M+BM-1)//BM, (N+BN-1)//BN)) - - if comm_op == "reduce_scatter": - D = torch.empty((M // world_size, N), dtype=torch.float16, device="cuda") - RowArray = generate_row_remap_array(M, N, BM, BN, cSeg, world_size) - - _warm_up = WARM_UP - _freq = REP - - if len(cSeg) == 1: - # No overlapping - if comm_op == "all_reduce": - for _ in range(_warm_up): - gemm_class.gemm_allreduce(A, B, C, Algo) - - gemm_class.gemm_allreduce(A, B, C, Algo) - - start_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - for i in range(_freq): - start_event[i].record() - gemm_class.gemm_allreduce(A, B, C, Algo) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - elif comm_op == "reduce_scatter": - for _ in range(_warm_up): - gemm_class.gemm_reducescatter(A, B, C, D, Algo) - - MonitoredMatrix[0] = 0 - gemm_class.gemm_reducescatter(A, B, C, D, Algo) - - start_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - for i in range(_freq): - start_event[i].record() - gemm_class.gemm_reducescatter(A, B, C, D, Algo) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - - else: - if comm_op == "all_reduce": - for _ in range(_warm_up): - gemm_class.gemm_allreduce_overlap(A, B, C, MonitoredMatrix, ReorderedArray, 1, cSeg_CPU, cSeg_GPU, Algo, False) - - start_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - for i in range(_freq): - start_event[i].record() - gemm_class.gemm_allreduce_overlap(A, B, C, MonitoredMatrix, ReorderedArray, 1, cSeg_CPU, cSeg_GPU, Algo, False) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - elif comm_op == "reduce_scatter": - for _ in range(_warm_up): - gemm_class.gemm_reducescatter_overlap(A, B, C, D, MonitoredMatrix, ReorderedArray, RowArray, 1, cSeg_CPU, cSeg_GPU, Algo, False) - - start_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(_freq)] - for i in range(_freq): - start_event[i].record() - gemm_class.gemm_reducescatter_overlap(A, B, C, D, MonitoredMatrix, ReorderedArray, RowArray, 1, cSeg_CPU, cSeg_GPU, Algo, False) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - else: - dur = torch.zeros((_freq)) - - result_dict[rank] = torch.mean(dur).item() - -def perf_running(M: int, N: int, K: int, - BM: int, BN: int, Algo: int, - cSeg: list, hint: list, comm_op: str): - world_size = torch.cuda.device_count() - if world_size < 2: - raise RuntimeError("At least 2 GPUs are required for this program.") - - nccl_id = torch.ops.flashoverlap_op.generate_nccl_id() - torch.cuda.synchronize() - # print(f"NCCL ID generated: {nccl_id[0]}") - - manager = mp.Manager() - result_dict = manager.dict() - - mp.spawn( - perf_running_process, - args=(world_size, nccl_id, M, N, K, BM, BN, Algo, cSeg, hint, comm_op, result_dict), - nprocs=world_size - ) - - dur = torch.empty((world_size)) - for i in range(world_size): - dur[i] = result_dict[i] - - return dur.max() - -# Function to initialize NCCL in each process -def perf_comm_process(rank, world_size, nccl_id, M, N, comm_type, result_dict): - torch.cuda.set_device(rank) - - comm_class = torch.classes.flashoverlap_class.OverlapImpl() - - comm_class.nccl_init(rank, world_size, nccl_id) - comm_class.cutlass_init() - - C = torch.empty((M, N), dtype=torch.float16, device="cuda") - if comm_type == "reduce_scatter": - D = torch.empty((M // world_size, N), dtype=torch.float16, device="cuda") - - if comm_type == "all_reduce": - for _ in range(WARM_UP): - comm_class.nccl_allreduce(C) - start_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - for i in range(REP): - start_event[i].record() - comm_class.nccl_allreduce(C) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - elif comm_type == "reduce_scatter": - for _ in range(WARM_UP): - comm_class.nccl_reducescatter(C) - start_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - for i in range(REP): - start_event[i].record() - comm_class.nccl_reducescatter(C) - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - else: - dur = torch.zeros((REP)) - - result_dict[rank] = torch.mean(dur).item() - -def perf_comm(M: int, N: int, comm_type: str): - world_size = torch.cuda.device_count() - if world_size < 2: - raise RuntimeError("At least 2 GPUs are required!") - - nccl_id = torch.ops.flashoverlap_op.generate_nccl_id() - torch.cuda.synchronize() - # print(f"NCCL ID generated: {nccl_id[0]}") - - manager = mp.Manager() - result_dict = manager.dict() - - # get the all reduce time - mp.spawn( - perf_comm_process, - args=(world_size, nccl_id, M, N, comm_type, result_dict), - nprocs=world_size - ) - - return result_dict[0] - -# Function to initialize NCCL in each process -def perf_baseline_process(rank, world_size, nccl_id, M, N, K, comm_op, result_dict): - torch.cuda.set_device(rank) - - A = torch.empty((M, K), dtype=torch.float16, device="cuda").normal_(mean=0., std=0.5) - B = torch.empty((N, K), dtype=torch.float16, device="cuda").normal_(mean=0., std=0.5) - C = torch.empty((M, N), dtype=torch.float16, device="cuda") - - if comm_op == "reduce_scatter": - D = torch.empty((M // world_size, N), dtype=torch.float16, device="cuda") - - # **** Init Baseline Class **** # - gemm_comm = torch.classes.flashoverlap_class.BaselineImpl() - gemm_comm.nccl_init(rank, world_size, nccl_id) - gemm_comm.cublas_init() - - if comm_op == "all_reduce": - for _ in range(WARM_UP): - gemm_comm.gemm_allreduce(A, B, C) - start_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - for i in range(REP): - start_event[i].record() - # torch.cuda.cudart().cudaProfilerStart() - gemm_comm.gemm_allreduce(A, B, C) - # torch.cuda.cudart().cudaProfilerStop() - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - elif comm_op == "reduce_scatter": - for _ in range(WARM_UP): - gemm_comm.gemm_reducescatter(A, B, C, D) - start_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(REP)] - for i in range(REP): - start_event[i].record() - # torch.cuda.cudart().cudaProfilerStart() - gemm_comm.gemm_reducescatter(A, B, C, D) - # torch.cuda.cudart().cudaProfilerStop() - end_event[i].record() - torch.cuda.synchronize() - dur = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)], dtype=torch.float) - else: - dur = torch.zeros((REP)) - - result_dict[rank] = torch.mean(dur).item() - -def perf_baseline(M: int, N: int, K: int, comm_op: str): - world_size = torch.cuda.device_count() - if world_size < 2: - raise RuntimeError("At least 2 GPUs are required for this program.") - # Use the custom NCCL initialization wrapper to get a unique NCCL ID - # nccl_id = NcclInit() - nccl_id = torch.ops.flashoverlap_op.generate_nccl_id() - torch.cuda.synchronize() - - manager = mp.Manager() - result_dict = manager.dict() - - # Spawn processes - mp.spawn( - perf_baseline_process, - args=(world_size, nccl_id, M, N, K, comm_op, result_dict), - nprocs=world_size - ) - - dur = torch.empty((world_size)) - for i in range(world_size): - dur[i] = result_dict[i] - - return dur.max() - -def main(): - device = torch.cuda.current_device() - props = torch.cuda.get_device_properties(device) - gpu_name = props.name[7:11].lower() - sm_count = props.multi_processor_count - wave_size = sm_count - 2 - - parser = argparse.ArgumentParser() - parser.add_argument('--m', type=int, default=4096) - parser.add_argument('--k', type=int, default=8192) - parser.add_argument('--n', type=int, default=8192) - parser.add_argument('--comm_op', type=str, default='all_reduce') - args = parser.parse_args() - - comm_op = args.comm_op - - m, n, k = args.m, args.n, args.k - - file_path = f'../configs/m{m}n{n}k{k}_{gpu_name}.json' - - with open(file_path, 'r', encoding='utf-8') as f: - data = json.load(f) - - tile_num = m // data["BM"] * n // data["BN"] - wave_num = (tile_num + wave_size - 1) // wave_size - - gemm_dur = data["dur"] - comm_dur = perf_comm(m, n, comm_op) - overlap_dur = perf_running(m, n, k, - data["BM"], data["BN"], data["Algo"], data["cSeg"], data["hint"], comm_op) - baseline_dur = perf_baseline(m, n, k, comm_op) - - speedup = baseline_dur / overlap_dur - - print(f""" - {'Item':<10} {'Value':>15} - {'-----':<10} {'-----':>15} - {'m':<10} {m:>15} - {'n':<10} {n:>15} - {'k':<10} {k:>15} - {'tile_num':<10} {tile_num:>15} - {'gemm_dur (ms)':<10} {gemm_dur:>15.4f} - {'comm_dur (ms)':<10} {comm_dur:>15.4f} - {'baseline_dur (ms)':<10} {baseline_dur:>15.4f} - {'overlap_dur (ms)':<10} {overlap_dur:>15.4f} - {'speedup':<10} {speedup:>15.4f} - """) - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/src/code/issue4/Test Scripts/Flux/launch.sh b/src/code/issue4/Test Scripts/Flux/launch.sh deleted file mode 100644 index 5dad826..0000000 --- a/src/code/issue4/Test Scripts/Flux/launch.sh +++ /dev/null @@ -1,39 +0,0 @@ -#!/bin/bash -# libflux_cuda.so maybe installed under /usr/local/lib or ~/.local/lib/ by pip3 -export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:~/.local/lib/ -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) -FLUX_SRC_DIR=${SCRIPT_DIR} - -# add flux python package to PYTHONPATH -export NVSHMEM_BOOTSTRAP_MPI_PLUGIN=nvshmem_bootstrap_torch.so -export NVSHMEM_DISABLE_CUDA_VMM=1 # moving from cpp to shell -export CUDA_DEVICE_MAX_CONNECTIONS=1 - -# set default communication env vars -export BYTED_TORCH_BYTECCL=O0 -export NCCL_IB_TIMEOUT=${NCCL_IB_TIMEOUT:=23} - -nproc_per_node=$(nvidia-smi --list-gpus | wc -l) -nnodes=1 -node_rank=0 -master_addr="127.0.0.1" -master_port="23456" -additional_args="--rdzv_endpoint=${master_addr}:${master_port}" -IB_HCA=mlx5 - - -export NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX:=3} -export NVSHMEM_IB_GID_INDEX=3 - - -CMD="torchrun \ - --node_rank=${node_rank} \ - --nproc_per_node=${nproc_per_node} \ - --nnodes=${nnodes} \ - ${FLUX_EXTRA_TORCHRUN_ARGS} ${additional_args} $@" - -echo ${CMD} -${CMD} - -ret=$? -exit $ret \ No newline at end of file diff --git a/src/code/issue4/Test Scripts/Flux/test_gemm_rs.py b/src/code/issue4/Test Scripts/Flux/test_gemm_rs.py deleted file mode 100644 index e6598dd..0000000 --- a/src/code/issue4/Test Scripts/Flux/test_gemm_rs.py +++ /dev/null @@ -1,477 +0,0 @@ -################################################################################ -# -# Copyright 2025 ByteDance Ltd. and/or its affiliates. All rights reserved. -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -################################################################################ - -# usage: torchrun --node_rank=0 --nproc_per_node=8 --nnodes=1 --rdzv_id=none --master_addr=127.0.0.1 --master_port=23456 test/python/gemm_rs/test_gemm_rs.py 2048 10240 40960 -import argparse -import os -import time -from typing import Optional - -import torch -import torch.distributed - -import flux -from flux.cpp_mod import ReduceScatterOption -import flux.testing -from flux.testing import DTYPE_MAP, generate_data, initialize_distributed, matmul_int8 -from flux.testing.perf_db_helper import log_perf, set_global_args, should_log_to_rds - - -class PerfResult: - def __init__( - self, name: str, output: torch.Tensor, gemm_time_ms: float, comm_time_ms: float - ) -> None: - self.name = name - self.output = output - self.gemm_time_ms = gemm_time_ms - self.comm_time_ms = comm_time_ms - self.total_ms = self.gemm_time_ms + self.comm_time_ms - - def __repr__(self) -> str: - return ( - f"{self.name}: gemm {self.gemm_time_ms:.3f} ms, comm {self.comm_time_ms:.3f} ms" - f", total {self.total_ms:.3f} ms" - ) - - -@torch.no_grad() -def perf_torch( - input: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - warmup: int, - iters: int, - transpose_weight: bool = False, - input_scale: Optional[torch.Tensor] = None, - weight_scale: Optional[torch.Tensor] = None, -): - TP_GROUP.barrier() - - is_fp8 = flux.util.is_fp8_dtype(input.dtype) - is_s8_dequant = input.dtype == torch.int8 - warmup_iters = warmup - output_dtype = torch.bfloat16 if is_fp8 or is_s8_dequant else input.dtype - m = input.size(0) - with flux.util.with_torch_deterministic(False): - if transpose_weight: - w = weight.t().contiguous() - n = w.size(1) - else: - n = weight.size(0) - w = weight - - full_output = torch.zeros( - [m, n], - dtype=output_dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - output = torch.zeros( - [m // WORLD_SIZE, n], - dtype=output_dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - - op = ( - flux.GemmOnly( - input_dtype=input.dtype, - output_dtype=output_dtype, - transpose_weight=transpose_weight, - use_fp8_gemm=is_fp8, - ) - if is_fp8 - else None - ) - - total_iters = warmup_iters + iters - start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] - gemm_end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] - torch.distributed.barrier() - - for i in range(total_iters): - start_events[i].record() - if is_fp8: - full_output = op.forward( - input, - w, - bias=bias, - output_buf=full_output, - input_scale=input_scale, - weight_scale=weight_scale, - output_scale=None, - fast_accum=False, - ) - elif is_s8_dequant: - accum = matmul_int8(input, weight.t()).to(torch.float32) - full_output = input_scale * weight_scale * accum - full_output = full_output.to(output_dtype) - else: - full_output = torch.matmul(input, weight.t()) - if bias is not None and not is_fp8: - # only apply bias on rank 0 for s8 gemm - if not is_s8_dequant or (is_s8_dequant and TP_GROUP.rank() == 0): - full_output += bias - gemm_end_events[i].record() - torch.distributed.reduce_scatter_tensor(output, full_output, group=TP_GROUP) - end_events[i].record() - - gemm_times = [] - comm_times = [] - for i in range(total_iters): - gemm_end_events[i].synchronize() - end_events[i].synchronize() - if i >= warmup_iters: - gemm_times.append(start_events[i].elapsed_time(gemm_end_events[i]) / 1000) - comm_times.append(gemm_end_events[i].elapsed_time(end_events[i]) / 1000) - # print(gemm_times) - # print(comm_times) - gemm_time = sum(gemm_times) / iters * 1000 - comm_time = sum(comm_times) / iters * 1000 - return PerfResult( - name=f"torch #{TP_GROUP.rank()}", - output=output, - gemm_time_ms=gemm_time, - comm_time_ms=comm_time, - ) - - -@torch.no_grad() -def perf_flux( - input: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - transpose_weight: bool, - fuse_reduction: bool, - ring_reduction: bool, - warmup: int, - iters: int, - input_scale: Optional[torch.Tensor] = None, - weight_scale: Optional[torch.Tensor] = None, - reduce_scatter_option: flux.ReduceScatterOption = flux.ReduceScatterOption(), -): - is_fp8 = flux.util.is_fp8_dtype(input.dtype) - is_s8_dequant = input.dtype == torch.int8 - M = input.size(0) - # todo: transpose here to avoid TN kernel, which has the worst performence - if transpose_weight: - with flux.util.with_torch_deterministic(False): - w = weight.t().contiguous() - N = w.size(1) - else: - w = weight - N = w.size(0) - - output_dtype = torch.bfloat16 if is_fp8 or is_s8_dequant else input.dtype - gemm_only_op = flux.GemmOnly( - w.dtype, - output_dtype, - transpose_weight=transpose_weight, - use_fp8_gemm=is_fp8, - ) - gemm_rs_op = flux.GemmRS( - TP_GROUP, - NNODES, - (M + 1024 - 1) // 1024 * 1024, - N, - input.dtype, - output_dtype, - transpose_weight=transpose_weight, - fuse_reduction=fuse_reduction, - ring_reduction=ring_reduction, - ) - - warmup_iters = warmup - total_iters = warmup_iters + iters - start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] - with flux.util.with_torch_deterministic(False): - gemm_only_output_buf = torch.empty( - [M, N], dtype=output_dtype, device=input.device, requires_grad=False - ) - - torch.distributed.barrier() - for i in range(total_iters): - start_events[i].record() - _ = gemm_only_op.forward( - input, - w, - bias=bias, - output_buf=gemm_only_output_buf, - input_scale=input_scale, - weight_scale=weight_scale, - output_scale=None, - fast_accum=False, - ) - end_events[i].record() - torch.cuda.current_stream().synchronize() - - gemm_times = [] - for i in range(total_iters): - end_events[i].synchronize() - if i >= warmup_iters: - gemm_times.append(start_events[i].elapsed_time(end_events[i]) / 1000) - gemm_time = sum(gemm_times) - - time.sleep(1) - - torch.distributed.barrier() - for i in range(total_iters): - start_events[i].record() - output = gemm_rs_op.forward( - input, - w, - bias=bias, - input_scale=input_scale, - weight_scale=weight_scale, - output_scale=None, - fast_accum=False, - reduce_scatter_option=reduce_scatter_option, - ) - end_events[i].record() - torch.cuda.current_stream().synchronize() - - gemm_rs_times = [] - for i in range(total_iters): - end_events[i].synchronize() - if i >= warmup_iters: - gemm_rs_times.append(start_events[i].elapsed_time(end_events[i]) / 1000) - gemm_rs_time = sum(gemm_rs_times) - - gemm_time_ms = gemm_time / iters * 1000 - comm_time_ms = (gemm_rs_time - gemm_time) / iters * 1000 - - return PerfResult( - name=f"flux #{TP_GROUP.rank()}", - output=output, - gemm_time_ms=gemm_time_ms, - comm_time_ms=comm_time_ms, - ) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("M", type=int) - parser.add_argument("N", type=int) - parser.add_argument("K", type=int) - parser.add_argument("--warmup", default=5, type=int, help="warmup iterations") - parser.add_argument("--iters", default=100, type=int, help="perf iterations") - parser.add_argument("--dtype", default="bfloat16", type=str, choices=list(DTYPE_MAP.keys())) - parser.add_argument( - "--profile", default=False, action="store_true", help="dump torch.profiler.profile" - ) - parser.add_argument( - "--transpose_weight", default=False, action="store_true", help="whether to transpose weight" - ) - parser.add_argument( - "--fuse_reduction", default=False, action="store_true", help="fuse reduction to gemm" - ) - parser.add_argument( - "--ring_reduction", - default=False, - action="store_true", - help="reduce paritial output with ring order", - ) - parser.add_argument("--has_bias", default=False, action="store_true", help="whether have bias") - parser.add_argument( - "--debug", action="store_true", help="debug mode. use human read input", default=False - ) - parser.add_argument( - "--use_1d_ring", - action=argparse.BooleanOptionalAction, - help="use 1d ring for reduction", - ) - parser.add_argument( - "--use_p2p_read", - action=argparse.BooleanOptionalAction, - help="use 1d ring for reduction", - ) - parser.add_argument( - "--use_cudaMemcpyAsync", - action=argparse.BooleanOptionalAction, - help="use 1d ring for reduction", - ) - parser.add_argument( - "--use_gemmk", - action=argparse.BooleanOptionalAction, - help="use 1d ring for reduction", - ) - parser.add_argument( - "--per_tile_flags", - action=argparse.BooleanOptionalAction, - help="use 1d ring for reduction", - ) - parser.add_argument( - "--reduce_scatter_blocks", - type=int, - help="number of blocks for reduce scatter", - ) - parser.add_argument( - "--ring_mode", - choices=["ring1d", "ring2d"], - help="ring mode. auto for auto detect", - ) - return parser.parse_args() - - -if __name__ == "__main__": - TP_GROUP = initialize_distributed() - RANK, WORLD_SIZE, NNODES = TP_GROUP.rank(), TP_GROUP.size(), flux.testing.NNODES() - - args = parse_args() - - input_dtype = DTYPE_MAP[args.dtype] - is_fp8 = flux.util.is_fp8_dtype(input_dtype) - is_s8_dequant = input_dtype == torch.int8 - - if args.transpose_weight and (is_fp8 or is_s8_dequant): - raise ValueError("FP8/S8 GEMM does not support RRR layout") - - assert args.M % TP_GROUP.size() == 0 - assert args.K % TP_GROUP.size() == 0 - local_K = args.K // TP_GROUP.size() - - # input: [M, K], weight: [N, K] - - scale = TP_GROUP.rank() + 1 - if is_s8_dequant: - data_config = [ - ((args.M, local_K), input_dtype, (127, 0)), # A - ((args.N, local_K), input_dtype, (127, 0)), # B - None if not args.has_bias else ((1, args.N), torch.bfloat16, (24, -12)), # bias - ((args.M, 1), torch.float32, (1 / 1024.0, 0)), # input_scale - ((1, args.N), torch.float32, (1 / (args.K * 4 / 1024.0), 0)), # weight_scale - ] - elif is_fp8: - data_config = [ - ((args.M, local_K), input_dtype, (0.01 * scale, 0)), # A - ((args.N, local_K), input_dtype, (0.01 * scale, 0)), # B - None, # bias. not supported now. ((1, args.N), torch.bfloat16, (0.1 * scale, 0)) - ((1), torch.float32, (1, 0)), # input_scale - ((1), torch.float32, (1, 0)), # weight_scale - ] - else: - data_config = [ - ((args.M, local_K), input_dtype, (0.01 * scale, 0)), # A - ((args.N, local_K), input_dtype, (0.01 * scale, 0)), # B - ( # bias - None if not args.has_bias else ((args.M, args.N), input_dtype, (0.1 * scale, 0)) - ), - None, # input_scale - None, # weight_scale - ] - - assert not (args.has_bias and is_fp8), "FP8 does not support bias" - generator = generate_data(data_config) - input, weight, bias, input_scale, weight_scale = next(generator) - - if args.debug: - input.zero_() - input[:, 0].fill_(TP_GROUP.rank() + 1) - weight.fill_(1) - if input_scale is not None: - input_scale.fill_(1) - if weight_scale is not None: - weight_scale.fill_(1) - if bias is not None: - bias.fill_(TP_GROUP.rank() + 1) - - reduce_scatter_option = ReduceScatterOption() - reduce_scatter_option.use_1d_ring = args.use_1d_ring - reduce_scatter_option.use_p2p_read = args.use_p2p_read - reduce_scatter_option.use_cudaMemcpyAsync = args.use_cudaMemcpyAsync - reduce_scatter_option.use_gemmk = args.use_gemmk - reduce_scatter_option.per_tile_flags = args.per_tile_flags - reduce_scatter_option.num_blocks = args.reduce_scatter_blocks - reduce_scatter_option.ring_mode = { - "ring1d": flux.RingMode.Ring1D, - "ring2d": flux.RingMode.Ring2D, - }.get(args.ring_mode, None) - with flux.util.group_profile( - name="gemm_rs_" + os.environ["TORCHELASTIC_RUN_ID"], do_prof=args.profile, group=TP_GROUP - ): - perf_res_flux = perf_flux( - input, - weight, - bias, - args.transpose_weight, - args.fuse_reduction, - args.ring_reduction, - args.warmup, - args.iters, - input_scale, - weight_scale, - reduce_scatter_option=reduce_scatter_option, - ) - perf_res_torch = perf_torch( - input, - weight, - bias, - args.warmup, - args.iters, - args.transpose_weight, - input_scale, - weight_scale, - ) - - if TP_GROUP.rank() == 0: - flux.testing.print_gemm_sol_time(args.M, args.N, local_K, input_dtype) - if should_log_to_rds(): - set_global_args("gemm_rs", args) - for i in range(TP_GROUP.size()): - if i == TP_GROUP.rank(): - log_perf(perf_res_torch) - TP_GROUP.barrier() - for i in range(TP_GROUP.size()): - if i == TP_GROUP.rank(): - log_perf(perf_res_flux) - TP_GROUP.barrier() - - TP_GROUP.barrier() - - flux_output = perf_res_flux.output - torch_output = perf_res_torch.output - THRESHOLD_MAP = { - torch.float16: 1e-2, - torch.bfloat16: 2e-2, - torch.float8_e4m3fn: 3e-2, - torch.float8_e5m2: 3e-2, - torch.int8: 2e-1, - } - - flux_output = flux_output.reshape(torch_output.size()) - atol, rtol = THRESHOLD_MAP[input_dtype], THRESHOLD_MAP[input_dtype] - try: - flux.torch_allclose(flux_output, torch_output, atol=atol, rtol=rtol) - except Exception as e: - torch.save(flux_output, f"flux_output_{RANK}.pt") - torch.save(torch_output, f"torch_output_{RANK}.pt") - print("❌ flux check failed") - raise e - else: - print("✅ flux check passed") - TP_GROUP.barrier() - - if TP_GROUP.rank() == 0: - if flux.bitwise_check(torch_output, flux_output): - print("✅ flux vs torch bitwise check passed") - else: - print("❌ flux vs torch bitwise check failed") - - TP_GROUP.barrier() - torch.cuda.synchronize() diff --git a/src/code/issue4/Test Scripts/TransformerEngine/te_layer_with_overlap.py b/src/code/issue4/Test Scripts/TransformerEngine/te_layer_with_overlap.py deleted file mode 100644 index 0718894..0000000 --- a/src/code/issue4/Test Scripts/TransformerEngine/te_layer_with_overlap.py +++ /dev/null @@ -1,387 +0,0 @@ -#!/usr/bin/python3 - -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -import os -import sys -import socket -import time -import fcntl -import struct -import argparse -import warnings - -import torch -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel - -import transformer_engine.pytorch as te -import transformer_engine.pytorch.cpp_extensions as tex -from transformer_engine.common.recipe import Format, DelayedScaling - -warnings.filterwarnings("ignore", category=DeprecationWarning) -warnings.filterwarnings("ignore", category=FutureWarning) -warnings.filterwarnings("ignore", category=UserWarning) - -os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" -if not tex.device_supports_multicast(): - os.environ["UB_SKIPMC"] = "1" - - -def _te_layer_argtype(name): - te_layers = [ - te.Linear, - te.LayerNormLinear, - te.LayerNormMLP, - te.MultiheadAttention, - te.TransformerLayer, - ] - layer_map = dict(zip([layer.__name__.lower() for layer in te_layers], te_layers)) - if name.lower() not in layer_map.keys(): - raise argparse.ArgumentTypeError( - f"Invalid TE layer name! Please choose from: {layer_map.keys()}" - ) - return layer_map[name.lower()] - - -def _parse_args(argv=None, namespace=None): - parser = argparse.ArgumentParser( - description="Train a Transformer Engine module with GEMM+comm overlap via Userbuffers." - ) - parser.add_argument( - "-i", "--num-iters", type=int, default=10, help="Number of dummy 'training' iterations." - ) - parser.add_argument("-b", "--batch-size", type=int, default=2, help="Input batch size.") - parser.add_argument("-s", "--seq-length", type=int, default=2048, help="Input sequence length.") - parser.add_argument( - "-n", "--num-heads", type=int, default=96, help="Number of attention heads." - ) - parser.add_argument( - "-d", "--head-dim", type=int, default=128, help="Dimension of each attention head." - ) - parser.add_argument( - "--layer-type", - type=_te_layer_argtype, - default=te.Linear, - help="Transformer Engine layer to train with comm+GEMM overlap.", - ) - parser.add_argument("--seed", type=int, default=1234, help="RNG seed.") - parser.add_argument( - "--fp8", action="store_true", default=False, help="Enables the te.fp8_autocast() context." - ) - parser.add_argument( - "--no-comm-overlap", - action="store_true", - default=False, - help="Disable the comm+GEMM overlap.", - ) - parser.add_argument( - "--num-replicas", - type=int, - default=1, - help="Number of data-parallel model replicas per node.", - ) - parser.add_argument( - "--use-global-replica-count", - action="store_true", - default=False, - help="Treat '--num-replicas' as the total number of replicas.", - ) - parser.add_argument( - "--tcp-init", - action="store_true", - default=False, - help="Initialize torch.distributed with TcpStore.", - ) - parser.add_argument( - "--bind-to-device", - action="store_true", - default=False, - help="Initialize torch.distributed with `device_id` to bind each rank to a single device.", - ) - parser.add_argument( - "--bootstrap-backend", - type=str.lower, - default="nccl", - choices=["gloo", "mpi", "nccl"], - help="Communications backend for host tensor collectives during Userbuffers bootstrapping.", - ) - parser.add_argument( - "-v", - "--verbose", - action="store_true", - default=False, - help="Print out from every rank instead of just the root rank of relevant process groups.", - ) - parser.add_argument( - "--debug", - action="store_true", - default=False, - help="Print out additional debug information.", - ) - args = parser.parse_args(argv, namespace) - if args.bootstrap_backend == "nccl": - args.bind_to_device = True - return args - - -def _get_layer_args(config, tp_group, tp_size, reference=False): - hidden_size = config.num_heads * config.head_dim - input_shape = [config.seq_length, config.batch_size, hidden_size] - args = [hidden_size] - kwargs = { - "params_dtype": torch.float32, - "device": "cuda", - "tp_group": tp_group, - "tp_size": tp_size, - "sequence_parallel": True, - } - kwargs["ub_overlap_ag"] = not config.no_comm_overlap - - if config.layer_type is te.Linear: - input_shape[2] = hidden_size // tp_size - args.append(hidden_size) - kwargs["parallel_mode"] = "row" - kwargs["ub_overlap_rs"] = not config.no_comm_overlap - kwargs["ub_name"] = "proj" - else: - input_shape[0] = config.seq_length // tp_size - kwargs["ub_bulk_wgrad"] = not config.no_comm_overlap - kwargs["ub_bulk_dgrad"] = not config.no_comm_overlap - if config.layer_type is te.LayerNormLinear: - args.append(3 * hidden_size) - kwargs["parallel_mode"] = "column" - kwargs["ub_name"] = "qkv" - else: - kwargs["set_parallel_mode"] = True - kwargs["ub_overlap_rs"] = not config.no_comm_overlap - if config.layer_type in [te.LayerNormMLP, te.TransformerLayer]: - args.append(4 * hidden_size) - kwargs["seq_length"] = config.seq_length - if config.layer_type in [te.MultiheadAttention, te.TransformerLayer]: - args.append(config.num_heads) - kwargs["attention_dropout"] = 0.0 - kwargs["fuse_qkv_params"] = True - if config.layer_type is te.MultiheadAttention: - kwargs["input_layernorm"] = True - else: - kwargs["ub_tp_comm_overlap"] = not config.no_comm_overlap - kwargs["hidden_dropout"] = 0.0 - - return args, kwargs, input_shape - - -def _train(opts): - if "OMPI_COMM_WORLD_SIZE" in os.environ: - # Execution with `mpirun -np N` - WORLD_RANK = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - WORLD_SIZE = int(os.getenv("OMPI_COMM_WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("OMPI_COMM_WORLD_LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("OMPI_COMM_WORLD_LOCAL_SIZE", "1")) - opts.tcp_init = True - opts.bind_to_device = True - opts.bootstrap_backend = "mpi" - else: # TORCHELASTIC, SLURM, etc... - WORLD_RANK = int(os.getenv("RANK", "0")) - WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", str(torch.cuda.device_count()))) - - NUM_NODES = WORLD_SIZE // LOCAL_SIZE - - # Initialize torch.distributed global process group and get DP/TP groups - torch.cuda.set_device(LOCAL_RANK) - dist_init_kwargs = { - "backend": "nccl", - "rank": WORLD_RANK, - "world_size": WORLD_SIZE, - } - if opts.tcp_init or NUM_NODES > 1: - if NUM_NODES > 1: - assert ( - "MASTER_ADDR" in os.environ - ), "Multi-node run requires MASTER_ADDR to be set in the environment." - MASTER_ADDR = os.getenv("MASTER_ADDR", socket.gethostbyname(socket.gethostname())) - MASTER_PORT = os.getenv("MASTER_PORT", "1234") - dist_init_kwargs["init_method"] = f"tcp://{MASTER_ADDR}:{MASTER_PORT}" - if opts.bind_to_device or opts.bootstrap_backend == "nccl": - dist_init_kwargs["device_id"] = torch.device(f"cuda:{LOCAL_RANK}") - assert dist.is_nccl_available() - dist.init_process_group(**dist_init_kwargs) - nccl_world = dist.new_group(backend="nccl") - - def dist_print(msg, end="\n", group=nccl_world, src=0, debug=False, error=False): - if debug and not opts.debug: - return - group_rank = dist.get_rank(group) - stream = sys.stderr if error else sys.stdout - if group_rank == src: - stream.write(f"[rank{WORLD_RANK}] {msg}{end}") - dist.barrier(group) - - dist_print(f"Initialized default NCCL process group with {WORLD_SIZE} GPUs") - - total_replicas = ( - opts.num_replicas if opts.use_global_replica_count else opts.num_replicas * NUM_NODES - ) - tp_size = WORLD_SIZE // total_replicas - - if total_replicas > 1: - ranks_per_replica_list = [ - [i * tp_size + t for t in range(tp_size)] for i in range(total_replicas) - ] - - tp_group, _ = dist.new_subgroups_by_enumeration(ranks_per_replica_list, backend="nccl") - ranks_per_replica_tensor = torch.tensor(ranks_per_replica_list, dtype=torch.int32) - dp_group, _ = dist.new_subgroups_by_enumeration( - ranks_per_replica_tensor.transpose(0, 1).tolist(), backend="nccl" - ) - else: - dp_group = None - tp_group = nccl_world - - tp_rank = dist.get_rank(tp_group) - tp_size = dist.get_world_size(tp_group) - dist_print( - f"Created tensor-parallel group: {dist.get_process_group_ranks(tp_group)}", - group=tp_group, - ) - if dp_group is not None: - dp_rank = dist.get_rank(dp_group) - dist_print( - f"Created data-parallel group: {dist.get_process_group_ranks(dp_group)}", - group=dp_group, - ) - else: - dp_rank = 0 - - # Intialize userbuffers - hidden_size = opts.num_heads * opts.head_dim - batched_size = opts.seq_length * opts.batch_size - if not opts.no_comm_overlap: - te.module.base.initialize_ub( - [batched_size, hidden_size], - tp_size, - use_fp8=opts.fp8, - dtype=torch.bfloat16, - bootstrap_backend=opts.bootstrap_backend, - ) - - # Initialize the fused LayerNorm + Multi-layer Perceptron module - torch.manual_seed(opts.seed + dp_rank) - torch.cuda.manual_seed(opts.seed + tp_rank) - layer_args, layer_kwargs, input_size = _get_layer_args(opts, tp_group, tp_size) - model = opts.layer_type(*layer_args, **layer_kwargs) - if dp_group is not None: - model = DistributedDataParallel(model, dim=1, process_group=dp_group) - - # Initialize optimizer with model parameters - optim = torch.optim.Adam(model.parameters(), lr=0.0001) - - # Fp8 recipe setup - fp8_format = Format.HYBRID - fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") - - ########## - WARMUP_ITERS = 3 - dist_print("Starting warmup iterations...") - for i in range(WARMUP_ITERS): - x = torch.randn(input_size, dtype=torch.float32, device="cuda", requires_grad=True) - with torch.amp.autocast("cuda", dtype=torch.bfloat16): - with te.fp8_autocast(enabled=opts.fp8, fp8_recipe=fp8_recipe, fp8_group=nccl_world): - y = model(x) - if isinstance(y, tuple): - out, *_ = y - else: - out = y - loss = out.sum() - loss.backward() - optim.step() - - # 清除梯度，为正式计时做准备 - optim.zero_grad() - torch.cuda.empty_cache() - torch.cuda.synchronize() - ########## - - # Start dummy "training" iterations - dist_print("Starting training iterations...") - - ########## - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - iter_times = torch.zeros(opts.num_iters, device="cuda") - ########## - - for i in range(opts.num_iters): - dist_print(f" Iter {i+1}", group=tp_group, debug=True) - - ########## - torch.cuda.synchronize() - start_event.record() - ########## - - dist_print(" |-- Generate random input batch", group=tp_group, debug=True) - x = torch.randn(input_size, dtype=torch.float32, device="cuda", requires_grad=True) - - dist_print(" |-- Forward pass", group=tp_group, debug=True) - with torch.amp.autocast("cuda", dtype=torch.bfloat16): - with te.fp8_autocast(enabled=opts.fp8, fp8_recipe=fp8_recipe, fp8_group=nccl_world): - y = model(x) - if isinstance(y, tuple): - out, *_ = y - else: - out = y - dist_print(" |-- Compute loss", group=tp_group, debug=True) - loss = out.sum() - - dist_print(" |-- Backward pass", group=tp_group, debug=True) - loss.backward() - - dist_print(" |-- Optimizer step", group=tp_group, debug=True) - optim.step() - - ########## - end_event.record() - torch.cuda.synchronize() - - # 记录本次迭代时间 - iter_time = start_event.elapsed_time(end_event) - iter_times[i] = iter_time - - dist_print(f" Iter {i+1}/{opts.num_iters} - Time: {iter_time:.2f} ms", - group=tp_group, debug=True) - - total_time = iter_times.sum() - avg_time = total_time / opts.num_iters - - # 收集所有rank的时间（确保时间测量一致） - time_data = torch.tensor([avg_time], device="cuda") - dist.all_reduce(time_data, op=dist.ReduceOp.SUM) - avg_time = time_data.item() / WORLD_SIZE - - # 输出结果（只在rank0打印） - if WORLD_RANK == 0: - print(f"\n{'='*50}") - print(f"Training completed with {opts.num_iters} iterations") - print(f"Average iteration time: {avg_time:.2f} ms") - print(f"Throughput: {1000/avg_time:.2f} iter/s") - print(f"{'='*50}\n") - ########## - - torch.cuda.synchronize() - dist_print("Finished training!") - te.module.base.destroy_ub() - - dist_print("Destroying all process groups...", debug=True) - dist.destroy_process_group() - if opts.debug and WORLD_RANK == 0: - print("Exiting...\n", end="", flush=True) - - return 0 - - -if __name__ == "__main__": - sys.exit(_train(_parse_args())) diff --git a/src/code/issue4/__init__.py b/src/code/issue4/__init__.py new file mode 100644 index 0000000..fef2660 --- /dev/null +++ b/src/code/issue4/__init__.py @@ -0,0 +1,47 @@ +""" +NCCL 集合通信位级可复现性诊断工具。 + +在大规模分布式训练中诊断 NCCL 集合通信操作（AllReduce / Reduce-Scatter） +的逐位非确定性问题。 + +根因：浮点加法不满足结合律。NCCL 的分块策略随数据规模、算法、 +协议和拓扑变化，导致归约累加顺序不同 → 逐位结果不同。 + +提供： + 1. 配置矩阵扫描（算法 × 协议 × 数据规模 × 精度） + 2. SHA-256 确定性逐 rank 差异化数据生成 + 3. 逐位比对与差异统计（XOR / ULP 位级分解） + 4. 差异随规模 / 迭代的演化追踪 + 5. 诊断建议与确定性配置推荐 +""" + +__version__ = "0.1.0" +__author__ = "NCCL Determinism Diagnostic Tool Contributors" + +from .config_matrix import ConfigMatrix, ConfigEntry, HardwareCaps +from .data_generator import DataGenerator +from .runner import NcclRunner, RunResult, MultiRunResult +from .comparator import ( + BitwiseComparator, + DiffReport, + DiffDistribution, + XorDetail, + EvolutionReport, + IterEvolutionReport, + ThreeLevelSummary, + compute_ulp, + analyze_xor_float, +) +from .reporter import Reporter, DiagnosticSummary + +__all__ = [ + "ConfigMatrix", "ConfigEntry", "HardwareCaps", + "DataGenerator", + "NcclRunner", "RunResult", "MultiRunResult", + "BitwiseComparator", + "DiffReport", "DiffDistribution", "XorDetail", + "EvolutionReport", "IterEvolutionReport", + "ThreeLevelSummary", + "compute_ulp", "analyze_xor_float", + "Reporter", "DiagnosticSummary", +] diff --git a/src/code/issue4/comparator.py b/src/code/issue4/comparator.py new file mode 100644 index 0000000..f100deb --- /dev/null +++ b/src/code/issue4/comparator.py @@ -0,0 +1,661 @@ +""" +逐位比对器 — NCCL 确定性诊断。 + +对 RunResult 输出进行逐元素比对，提供: + - 首次位级差异出现的偏移位置 + - 差异量级统计（最大值、均值、标准差） + - 差异元素比例 + - 差异随数据规模 / 迭代次数的演化追踪 + - 差异空间分布分析（前/中/后三区聚类 + 直方图） + - 逐位 XOR 分解 + ULP 距离 + +核心比对使用 numpy.array_equal — 零容忍度。 +这能捕捉到最小的位级非确定性。 +""" + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass, field +from typing import Optional + +import numpy as np + +from .config_matrix import ConfigEntry +from .runner import RunResult + + +# --------------------------------------------------------------------------- +# ---- 报告数据类 ---- +# --------------------------------------------------------------------------- + +@dataclass +class XorDetail: + """Byte-level XOR + ULP analysis for a single differing float32 element. + + Interprets two float32 values as IEEE-754 bit patterns, computes: + - XOR of the raw 32-bit patterns → reveals which bits flipped + - Sign / exponent / mantissa sub-patterns + - ULP (Units in Last Place) distance — the canonical measure of + floating-point difference independent of magnitude. + """ + + offset: int = 0 # element index in tensor + baseline_bits: str = "" # hex representation, e.g. "0x40490fdb" + target_bits: str = "" + xor_bits: int = 0 # raw XOR of the two bit-patterns + sign_diff: bool = False # did the sign bit flip? + exp_diff: int = 0 # exponent field difference + mantissa_diff: int = 0 # mantissa (fraction) field difference + ulp_distance: int = 0 # ULP distance between the two values + n_mantissa_bits_flipped: int = 0 # count of flipped mantissa bits + n_total_bits_flipped: int = 0 # total bit flips in the 32-bit word + + def summary(self) -> str: + return ( + f"XOR @ offset {self.offset}: " + f"baseline={self.baseline_bits} target={self.target_bits}\n" + f" XOR=0x{self.xor_bits:08x} " + f"ULP={self.ulp_distance} " + f"sign={'FLIP' if self.sign_diff else 'ok'} " + f"exp_diff={self.exp_diff} " + f"mantissa_flips={self.n_mantissa_bits_flipped}/{23}" + ) + + +def compute_ulp(a: float, b: float) -> int: + """计算两个 float32 值的 ULP 距离。 + + ULP (Units in Last Place): the number of representable float32 values + between `a` and `b`. For IEEE-754 binary32: + - Reinterpret both as int32 + - Convert to sign-magnitude (handle the two's complement quirk) + - Absolute difference = ULP distance + + Reference: "Comparing Floating Point Numbers, 2012 Edition" — Random ASCII + + Examples: + compute_ulp(1.0, 1.0 + 1e-7) ≈ 1 (adjacent floats = 1 ULP) + compute_ulp(1.0, 2.0) ≈ 2**23 (~8 million ULPs) + compute_ulp(0.0, -0.0) == 0 (IEEE-754: +0 and −0 are equal) + """ + # 通过 struct 将 float32 重新解释为 int32，避免 numpy 依赖 + a_bits = struct.unpack(" int: + # IEEE-754 二进制补码转符号-幅度：负数翻转位序以保证 ULP 距离单调 + if x & 0x80000000: # negative (or -0) + return 0x80000000 - (x & 0x7FFFFFFF) + return x + + a_signed = _to_signed(a_bits) + b_signed = _to_signed(b_bits) + return abs(a_signed - b_signed) + + +def analyze_xor_float(a: float, b: float, offset: int = 0) -> XorDetail: + """对两个 float32 值进行字节级 XOR 分解。 + + IEEE-754 binary32 layout: [sign:1][exponent:8][mantissa:23] + """ + a_bits = struct.unpack("> 23) - ((b_bits & exp_mask) >> 23) + ) + mant_diff = abs( + (a_bits & mantissa_mask) - (b_bits & mantissa_mask) + ) + + # 统计尾数区域翻转的位数 + mant_xor = xor & mantissa_mask + n_mantissa_flips = mant_xor.bit_count() + + # 统计总翻转位数 + n_total_flips = xor.bit_count() + + ulp = compute_ulp(a, b) + + return XorDetail( + offset=offset, + baseline_bits=f"0x{a_bits:08x}", + target_bits=f"0x{b_bits:08x}", + xor_bits=xor, + sign_diff=sign_diff, + exp_diff=exp_diff, + mantissa_diff=mant_diff, + ulp_distance=ulp, + n_mantissa_bits_flipped=n_mantissa_flips, + n_total_bits_flipped=n_total_flips, + ) + + +@dataclass +class DiffReport: + """单个配置对的逐位比对详细报告。""" + + config: ConfigEntry + bitwise_match: bool + first_diff_offset: int = -1 # -1 if no diff + first_diff_baseline: float = 0.0 + first_diff_target: float = 0.0 + diff_count: int = 0 # number of differing elements + diff_ratio: float = 0.0 # diff_count / total_elements + max_abs_diff: float = 0.0 + mean_abs_diff: float = 0.0 + std_abs_diff: float = 0.0 + total_elements: int = 0 + diff_indices: Optional[np.ndarray] = None # indices of all diff positions + distribution: Optional[DiffDistribution] = None # spatial clustering + xor_detail: Optional[XorDetail] = None # XOR/ULP for first-diff element + + def summary(self) -> str: + if self.bitwise_match: + return ( + f"[PASS] {self.config}\n" + f" {self.total_elements} elements — bitwise identical" + ) + + lines = [ + f"[FAIL] {self.config}", + f" Total: {self.total_elements} elements, " + f"diff: {self.diff_count} ({self.diff_ratio * 100:.4f}%)", + f" First diff @ offset {self.first_diff_offset}: " + f"baseline={self.first_diff_baseline:.8e}, " + f"target={self.first_diff_target:.8e}", + f" max_abs_diff={self.max_abs_diff:.8e}, " + f"mean_abs_diff={self.mean_abs_diff:.8e}", + ] + if self.distribution: + lines.append(f" {self.distribution}") + if self.xor_detail: + lines.append(f" {self.xor_detail.summary()}") + return "\n".join(lines) + + +@dataclass +class DiffDistribution: + """Spatial distribution of differences within the output tensor.""" + + front_ratio: float = 0.0 # diff ratio in first 1/3 of tensor + mid_ratio: float = 0.0 # diff ratio in middle 1/3 + back_ratio: float = 0.0 # diff ratio in last 1/3 + diff_concentration: str = "uniform" # "front" / "mid" / "back" / "uniform" / "edges" + + # 差异量级直方图（log10 分桶） + histogram_buckets: dict[str, int] = field(default_factory=dict) + + def __str__(self) -> str: + return ( + f"Dist: front={self.front_ratio * 100:.3f}% " + f"mid={self.mid_ratio * 100:.3f}% " + f"back={self.back_ratio * 100:.3f}% " + f"concentration={self.diff_concentration}" + ) + + def to_dict(self) -> dict: + return { + "front_ratio_pct": f"{self.front_ratio * 100:.4f}%", + "mid_ratio_pct": f"{self.mid_ratio * 100:.4f}%", + "back_ratio_pct": f"{self.back_ratio * 100:.4f}%", + "concentration": self.diff_concentration, + "histogram_buckets": self.histogram_buckets, + } + + +@dataclass +class EvolutionReport: + """追踪差异随规模或迭代的演化。""" + + config_label: str # e.g. "Ring/Simple/float32" + entries: list[tuple[int, DiffReport]] = field(default_factory=list) + # (size_bytes, diff_report) + + def add(self, size_bytes: int, report: DiffReport) -> None: + self.entries.append((size_bytes, report)) + + def diff_ratio_curve(self) -> dict[str, list]: + """Return {size_labels: [...], diff_ratios: [...]} for plotting.""" + sizes: list[str] = [] + ratios: list[float] = [] + for sz, r in self.entries: + sizes.append(_format_size(sz)) + ratios.append(r.diff_ratio) + return {"size_labels": sizes, "diff_ratios": ratios} + + def summary(self) -> str: + lines = [f"Evolution: {self.config_label}"] + for sz, r in self.entries: + status = "IDENTICAL" if r.bitwise_match else f"{r.diff_ratio * 100:.4f}% diff" + lines.append(f" {_format_size(sz):>6s} → {status}") + return "\n".join(lines) + + +@dataclass +class IterEvolutionReport: + """Tracks diff accumulation across iterations within a single run. + + Key for the issue requirement: "差异随迭代的演化". + In real training, differences accumulate step-by-step. This tracks + how the diff between a fixed baseline and each iteration grows. + """ + + config_label: str + iteration_diffs: list[dict] = field(default_factory=list) + # Each entry: {"iter": N, "diff_count": ..., "diff_ratio": ..., "max_abs_diff": ...} + + def add(self, iteration: int, report: DiffReport) -> None: + self.iteration_diffs.append({ + "iter": iteration, + "diff_count": report.diff_count, + "diff_ratio": report.diff_ratio, + "max_abs_diff": report.max_abs_diff, + }) + + def growing(self) -> bool: + """Check if diffs are monotonically growing (accumulation pattern).""" + if len(self.iteration_diffs) < 2: + return False + ratios = [d["diff_ratio"] for d in self.iteration_diffs] + return all(ratios[i] <= ratios[i + 1] for i in range(len(ratios) - 1)) + + def summary(self) -> str: + lines = [f"Iter Evolution: {self.config_label}"] + for d in self.iteration_diffs: + lines.append( + f" iter {d['iter']:3d}: diff={d['diff_count']} " + f"({d['diff_ratio'] * 100:.4f}%) max={d['max_abs_diff']:.6e}" + ) + if self.growing(): + lines.append(" → Diffs are MONOTONICALLY GROWING (accumulation)") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Comparator +# --------------------------------------------------------------------------- + +@dataclass +class BitwiseComparator: + """以逐位精度比对 NCCL 运行输出。 + + Parameters + ---------- + tolerance : str + 'bitwise' — strict element equality (default for diagnostic) + 'relative' — use relative tolerance (for sanity checks) + rtol : float + Relative tolerance for 'relative' mode. + """ + + tolerance: str = "bitwise" + rtol: float = 1e-5 + + # ------------------------------------------------------------------ + # ---- 单次比对 ---- + # ------------------------------------------------------------------ + + def compare( + self, baseline: RunResult, target: RunResult + ) -> DiffReport: + """逐元素比对两个运行结果。""" + b = baseline.output + t = target.output + + if b.shape != t.shape: + raise ValueError( + f"Shape mismatch: baseline {b.shape} vs target {t.shape}" + ) + total = b.size + + if self.tolerance == "bitwise": + diff_mask = (b != t) + elif self.tolerance == "relative": + diff_mask = ~np.isclose(b, t, rtol=self.rtol, atol=0) + else: + raise ValueError(f"Unknown tolerance mode: {self.tolerance}") + + diff_indices = np.where(diff_mask)[0] + diff_count = len(diff_indices) + + report = DiffReport( + config=target.config, + bitwise_match=(diff_count == 0), + total_elements=int(total), + diff_count=diff_count, + diff_ratio=diff_count / total if total > 0 else 0.0, + diff_indices=diff_indices if diff_count > 0 else None, + ) + + if diff_count > 0: + abs_diff = np.abs(b[diff_mask].astype(np.float64) - + t[diff_mask].astype(np.float64)) + first_idx = diff_indices[0] + report.first_diff_offset = int(first_idx) + report.first_diff_baseline = float(b.flat[first_idx]) + report.first_diff_target = float(t.flat[first_idx]) + report.max_abs_diff = float(np.max(abs_diff)) + report.mean_abs_diff = float(np.mean(abs_diff)) + report.std_abs_diff = float(np.std(abs_diff)) + + # --- Spatial distribution analysis --- + report.distribution = self.analyze_distribution( + diff_indices, abs_diff, int(total) + ) + + # --- XOR / ULP analysis of first diff --- + report.xor_detail = analyze_xor_float( + report.first_diff_baseline, + report.first_diff_target, + offset=report.first_diff_offset, + ) + + return report + + # ------------------------------------------------------------------ + # ---- 分布分析 ---- + # ------------------------------------------------------------------ + + def analyze_distribution( + self, diff_indices: np.ndarray, abs_diffs: np.ndarray, total_elements: int + ) -> DiffDistribution: + """分析差异在张量中的空间聚类。 + + This matters because: + - Front-heavy diffs → initialization / first-chunk issue + - Back-heavy diffs → tail-of-chunk rounding (most common in NCCL) + - Uniform diffs → systemic non-determinism (algorithm-level) + - Edges → chunk-boundary effects + """ + n = total_elements + third = n // 3 + + # Partition mask + front_mask = diff_indices < third + mid_mask = (diff_indices >= third) & (diff_indices < 2 * third) + back_mask = diff_indices >= 2 * third + + front_count = int(np.sum(front_mask)) + mid_count = int(np.sum(mid_mask)) + back_count = int(np.sum(back_mask)) + + # Per-segment ratios (normalized by segment size) + front_ratio = front_count / third if third > 0 else 0.0 + mid_ratio = mid_count / third if third > 0 else 0.0 + back_ratio = back_count / (n - 2 * third) if (n - 2 * third) > 0 else 0.0 + + # Concentration heuristic + max_seg = max(front_ratio, mid_ratio, back_ratio) + if max_seg == 0: + concentration = "uniform" + elif max_seg >= 2 * min(f for f in (front_ratio, mid_ratio, back_ratio) if f > 0): + # One segment has 2× more diffs than another → concentrated + concentration = { + 0: "front", 1: "mid", 2: "back" + }[np.argmax([front_ratio, mid_ratio, back_ratio])] + elif front_ratio + back_ratio > 2 * mid_ratio: + concentration = "edges" + else: + concentration = "uniform" + + # --- Histogram of diff magnitudes (log10 buckets) --- + histogram: dict[str, int] = {} + if len(abs_diffs) > 0: + log_abs = np.log10(abs_diffs + 1e-40) + bins = [-12, -10, -8, -6, -4, -2, 0] + for lo, hi in zip(bins, bins[1:]): + count = int(np.sum((log_abs >= lo) & (log_abs < hi))) + if count > 0: + histogram[f"1e{lo}~1e{hi}"] = count + + return DiffDistribution( + front_ratio=front_ratio, + mid_ratio=mid_ratio, + back_ratio=back_ratio, + diff_concentration=concentration, + histogram_buckets=histogram, + ) + + # ------------------------------------------------------------------ + # ---- 多试验 & 跨配置比对 ---- + # ------------------------------------------------------------------ + + def compare_trials(self, results: list[RunResult]) -> list[DiffReport]: + """将多次试验与首次（基准）比对。""" + if len(results) < 2: + return [] + baseline = results[0] + reports: list[DiffReport] = [] + for target in results[1:]: + reports.append(self.compare(baseline, target)) + return reports + + def compare_configs( + self, results: list[list[RunResult]] + ) -> list[DiffReport]: + """对每个配置（含 2+ 试验），比对试验 0 与试验 1。 + + Returns one DiffReport per config. + """ + reports: list[DiffReport] = [] + for trials in results: + if len(trials) < 2: + continue + reports.append(self.compare(trials[0], trials[1])) + return reports + + # ------------------------------------------------------------------ + # ---- 跨规模的演化追踪 ---- + # ------------------------------------------------------------------ + + def track_evolution( + self, results: list[list[RunResult]], matrix: list[ConfigEntry] + ) -> dict[str, EvolutionReport]: + """按 (algo/proto/dtype) 分组，追踪差异率 vs 数据规模。 + + Returns {label: EvolutionReport} keyed by config label. + """ + evolutions: dict[str, EvolutionReport] = {} + + for i, (trials, cfg) in enumerate(zip(results, matrix)): + label = f"{cfg.algo}/{cfg.proto}/{cfg.dtype}" + if label not in evolutions: + evolutions[label] = EvolutionReport(config_label=label) + + if len(trials) >= 2: + report = self.compare(trials[0], trials[1]) + evolutions[label].add(cfg.size_bytes, report) + + return evolutions + + # ------------------------------------------------------------------ + # Iteration-level evolution (NEW — "差异随迭代的演化") + # ------------------------------------------------------------------ + + def track_iter_evolution( + self, results: list[RunResult] + ) -> IterEvolutionReport: + """追踪单配置内跨迭代的差异累积。 + + Each result represents a separate allreduce invocation with the same + input. The baseline is iteration 0; subsequent iterations are compared + against it. + + This reveals whether non-determinism is: + - Accumulating (diff ratio grows with iter) → training drift risk + - Stable (diff ratio constant) → single-precision loss, not compounding + """ + if len(results) < 2: + cfg = results[0].config if results else None + label = f"{cfg.algo}/{cfg.proto}/{cfg.dtype}" if cfg else "unknown" + return IterEvolutionReport(config_label=label) + + baseline = results[0] + label = f"{baseline.config.algo}/{baseline.config.proto}/{baseline.config.dtype}" + + report = IterEvolutionReport(config_label=label) + for i, target in enumerate(results[1:], start=1): + diff = self.compare(baseline, target) + report.add(i, diff) + + return report + + # ------------------------------------------------------------------ + # ---- 三级比对（运行 / 调用 / rank） ---- + # ------------------------------------------------------------------ + + def compare_run_vs_run( + self, run0: "MultiRunResult", run1: "MultiRunResult" + ) -> dict[str, list[DiffReport]]: + """Compare two runs: for each (call_idx, rank), trial0 vs trial1. + + Returns {"call0_rank1": DiffReport, ...} keyed by "call{k}_rank{r}". + """ + reports: dict[str, list[DiffReport]] = {} + n_calls = min(len(run0.call_outputs), len(run1.call_outputs)) + for c in range(n_calls): + ranks0 = run0.call_outputs[c] + ranks1 = run1.call_outputs[c] + for r in sorted(set(ranks0) & set(ranks1)): + key = f"call{c}_rank{r}" + b = _make_runresult(run0.config, ranks0[r], r) + t = _make_runresult(run1.config, ranks1[r], r) + reports.setdefault(key, []).append(self.compare(b, t)) + return reports + + def compare_call_vs_call( + self, run: "MultiRunResult" + ) -> dict[int, list[DiffReport]]: + """Within one run, compare call 0 vs each subsequent call, per rank. + + Returns {rank: [DiffReport(call0_vs_call1), DiffReport(call0_vs_call2), ...]}. + """ + reports: dict[int, list[DiffReport]] = {} + calls = run.call_outputs + if len(calls) < 2: + return reports + baseline_call = calls[0] + for rank in sorted(baseline_call): + baseline = _make_runresult(run.config, baseline_call[rank], rank) + for ci in range(1, len(calls)): + if rank in calls[ci]: + tgt = _make_runresult(run.config, calls[ci][rank], rank) + reports.setdefault(rank, []).append(self.compare(baseline, tgt)) + return reports + + def compare_rank_vs_rank( + self, run: "MultiRunResult", call_idx: int = 0 + ) -> dict[str, list[DiffReport]]: + """Within one run and one call, compare rank 0 vs each other rank. + + NCCL guarantees all ranks receive identical results in a single + collective call. This method VERIFIES that guarantee. + + Returns {"rank0_vs_rank1": DiffReport, ...}. + """ + reports: dict[str, list[DiffReport]] = {} + if call_idx >= len(run.call_outputs): + return reports + per_rank = run.call_outputs[call_idx] + ranks = sorted(per_rank) + if len(ranks) < 2: + return reports + baseline = _make_runresult(run.config, per_rank[ranks[0]], ranks[0]) + for r in ranks[1:]: + key = f"rank{ranks[0]}_vs_rank{r}" + tgt = _make_runresult(run.config, per_rank[r], r) + reports.setdefault(key, []).append(self.compare(baseline, tgt)) + return reports + + def full_three_level_report( + self, results: list["MultiRunResult"], label: str = "" + ) -> "ThreeLevelSummary": + """Run all three comparison levels and return a structured summary.""" + summary = ThreeLevelSummary(label=label or results[0].config.algo) + if len(results) >= 2: + summary.run_vs_run = self.compare_run_vs_run(results[0], results[1]) + if results: + summary.call_vs_call = self.compare_call_vs_call(results[0]) + summary.rank_vs_rank = self.compare_rank_vs_rank(results[0], call_idx=0) + return summary + + +# --------------------------------------------------------------------------- +# ---- 三级摘要 ---- +# --------------------------------------------------------------------------- + +@dataclass +class ThreeLevelSummary: + """Structured report from three-level comparison.""" + + label: str = "" + run_vs_run: dict[str, list[DiffReport]] = field(default_factory=dict) + call_vs_call: dict[int, list[DiffReport]] = field(default_factory=dict) + rank_vs_rank: dict[str, list[DiffReport]] = field(default_factory=dict) + + @property + def all_clean(self) -> bool: + for reports in self.run_vs_run.values(): + if any(not r.bitwise_match for r in reports): + return False + for reports in self.call_vs_call.values(): + if any(not r.bitwise_match for r in reports): + return False + for reports in self.rank_vs_rank.values(): + if any(not r.bitwise_match for r in reports): + return False + return True + + def summary(self) -> str: + lines = [f"Three-Level Comparison: {self.label}"] + lines.append(f" Run-vs-Run : {self._summarize_level(self.run_vs_run)}") + lines.append(f" Call-vs-Call : {self._summarize_level(self.call_vs_call)}") + lines.append(f" Rank-vs-Rank : {self._summarize_level(self.rank_vs_rank)}") + return "\n".join(lines) + + @staticmethod + def _summarize_level(reports: dict) -> str: + if not reports: + return "NO DATA" + total = sum(len(v) for v in reports.values()) + failing = sum( + sum(1 for r in v if not r.bitwise_match) for v in reports.values() + ) + if failing == 0: + return f"all {total} identical" + return f"{failing}/{total} with differences" + + +def _make_runresult(config: "ConfigEntry", output: np.ndarray, + rank: int) -> "RunResult": + """Helper to construct a RunResult on-the-fly for comparison.""" + from .runner import RunResult as RR + return RR(config=config, output=output, rank=rank) + + +# --------------------------------------------------------------------------- + +def _format_size(nbytes: int) -> str: + if nbytes < 1024: + return f"{nbytes}B" + if nbytes < 1024 * 1024: + return f"{nbytes // 1024}K" + return f"{nbytes // (1024 * 1024)}M" diff --git a/src/code/issue4/config_matrix.py b/src/code/issue4/config_matrix.py new file mode 100644 index 0000000..89b8d2b --- /dev/null +++ b/src/code/issue4/config_matrix.py @@ -0,0 +1,398 @@ +""" +配置矩阵生成器 — NCCL 确定性诊断。 + +对影响位级可复现性的 NCCL 参数（算法、协议、数据规模、数据类型） +生成笛卡尔积，用于诊断扫描。 + +参考: + - NCCL_ALGO: Ring, Tree, CollnetDirect, CollnetChain, NVLS, NVLSTree, PAT + - NCCL_PROTO: LL, LL128, Simple + - NCCL 环境变量文档: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass, field +from typing import Optional + +# --------------------------------------------------------------------------- +# 可用算法 / 协议定义（来自 NCCL tuning.cc 和文档） +# --------------------------------------------------------------------------- + +# Algorithm → supported collectives matrix (NCCL 2.19+, Table III from ETH paper) +ALGO_COLLECTIVE_SUPPORT: dict[str, set[str]] = { + "Ring": {"allreduce", "reducescatter", "allgather", "broadcast", "reduce"}, + "Tree": {"allreduce", "reducescatter", "allgather", "broadcast", "reduce"}, + "CollnetDirect": {"allreduce"}, + "CollnetChain": {"allreduce"}, + "NVLS": {"allreduce", "reducescatter", "allgather"}, + "NVLSTree": {"allreduce"}, + "PAT": {"allreduce", "reducescatter"}, +} + +# 协议 → 硬件要求 +# LL128 requires Hopper (SM90+) or later +PROTO_REQUIREMENTS: dict[str, str] = { + "LL": "All platforms", + "LL128": "Hopper+ (SM90+) — enabling on unsupported HW causes silent corruption", + "Simple": "All platforms", +} + +# 协议 → 最低 SM 主版本号（None = 无限制） +PROTO_MIN_SM: dict[str, int | None] = { + "LL": None, + "LL128": 90, # Hopper (SM 9.0) + "Simple": None, +} + +# 算法 → 硬件要求 +# Some algorithms require NVSwitch, SHARP-capable network, or specific topology. +ALGO_HW_REQUIREMENTS: dict[str, str] = { + "Ring": "All platforms (may silently corrupt on A800 with many NVLink channels — see NCCL#1055)", + "Tree": "All platforms", + "CollnetDirect": "Requires NVSwitch + SHARP-capable network", + "CollnetChain": "Requires SHARP-capable network (collnet support)", + "NVLS": "Requires NVSwitch (H100/B200/Blackwell) — single-node with NVSwitch fabric", + "NVLSTree": "Requires NVSwitch (H100/B200/Blackwell)", + "PAT": "NCCL 2.23+ — all platforms", +} + +# 算法 → 需要 NVSwitch？ +ALGO_NEEDS_NVSWITCH: set[str] = {"NVLS", "NVLSTree", "CollnetDirect"} + +# 算法 → 需要 SHARP/collnet？ +ALGO_NEEDS_COLLNET: set[str] = {"CollnetDirect", "CollnetChain"} + +# 数据类型 → 对归约顺序的精度敏感度 +DTYPE_SENSITIVITY: dict[str, str] = { + "float32": "Low — 23-bit mantissa, good tolerance", + "float16": "High — 10-bit mantissa, susceptible to order-dependent rounding", + "bfloat16": "High — 7-bit mantissa + larger exponent range, prone to non-determinism", +} + + +# --------------------------------------------------------------------------- +# Hardware capability record +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class HardwareCaps: + """检测到的硬件能力，用于兼容性过滤。""" + + sm_major: int # e.g. 80 (A100), 90 (H100), 100 (B200) + gpu_name: str # e.g. "NVIDIA A100-SXM4-80GB" + gpu_count: int # total GPU count visible + has_nvswitch: bool # NVSwitch fabric present? + has_collnet: bool # SHARP/collnet-capable network? + nccl_version: str = "" # e.g. "2.21.5" + + def supports_ll128(self) -> bool: + """LL128 requires Hopper (SM90+) or later.""" + return self.sm_major >= 90 + + def supports_nvswitch_algos(self) -> bool: + """NVLS/NVLSTree/CollnetDirect require NVSwitch.""" + return self.has_nvswitch + + def supports_collnet_algos(self) -> bool: + """CollnetChain requires SHARP-capable network.""" + return self.has_collnet + + def compatible_algos(self) -> set[str]: + """Return set of algorithm names supported on this hardware.""" + algos = {"Ring", "Tree", "PAT"} + if self.has_nvswitch: + algos.update({"NVLS", "NVLSTree"}) + if self.has_nvswitch and self.has_collnet: + algos.add("CollnetDirect") + if self.has_collnet: + algos.add("CollnetChain") + return algos + + def compatible_protos(self) -> set[str]: + """Return set of protocol names supported on this hardware.""" + protos = {"LL", "Simple"} + if self.supports_ll128(): + protos.add("LL128") + return protos + + def check_and_warn(self, algo: str, proto: str) -> list[str]: + """Return list of warnings for an algo/proto combo on this hardware.""" + warnings: list[str] = [] + if proto == "LL128" and not self.supports_ll128(): + warnings.append( + f"LL128 requires Hopper+ (SM90+). " + f"Detected SM{self.sm_major} ({self.gpu_name}). " + f"Enabling LL128 on unsupported HW causes SILENT DATA CORRUPTION. " + f"SKIPPED." + ) + if algo in ALGO_NEEDS_NVSWITCH and not self.has_nvswitch: + warnings.append( + f"Algorithm '{algo}' requires NVSwitch ({ALGO_HW_REQUIREMENTS[algo]}). " + f"NVSwitch not detected. SKIPPED." + ) + if algo in ALGO_NEEDS_COLLNET and not self.has_collnet: + warnings.append( + f"Algorithm '{algo}' requires SHARP/collnet-capable network. " + f"Not detected. SKIPPED." + ) + return warnings + + @classmethod + def detect(cls, require_nvswitch: bool = False, + require_collnet: bool = False) -> "HardwareCaps": + """Auto-detect hardware capabilities via PyTorch / pynvml.""" + sm_major = 0 + gpu_name = "unknown" + gpu_count = 0 + + try: + import torch + gpu_count = torch.cuda.device_count() + if gpu_count > 0: + props = torch.cuda.get_device_properties(0) + sm_major = props.major + gpu_name = props.name + except Exception: + pass + + # NVSwitch 检测：数据中心 GPU (H100/B200/A100) 通常有 NVSwitch， + # 消费级 GPU (RTX/GeForce) 一定没有。用 GPU 名称精确判断。 + _datacenter_gpus = {"H100", "H200", "H800", "B100", "B200", "A100", "A800", + "GH200", "GB200"} + _consumer_gpus = {"RTX", "GeForce", "GTX", "TITAN", "Quadro"} + + has_nvswitch = False + if sm_major >= 80 and gpu_count >= 2: + gpu_upper = gpu_name.upper() + is_consumer = any(pat.upper() in gpu_upper for pat in _consumer_gpus) + is_datacenter = any(pat.upper() in gpu_upper for pat in _datacenter_gpus) + if is_datacenter and not is_consumer: + has_nvswitch = True + has_nvswitch = has_nvswitch or require_nvswitch + + # Collnet/SHARP: typically InfiniBand with SHARP-capable switches. + # Default to False unless explicitly enabled. + has_collnet = require_collnet + + return cls( + sm_major=sm_major, + gpu_name=gpu_name, + gpu_count=gpu_count, + has_nvswitch=has_nvswitch, + has_collnet=has_collnet, + ) + + +@dataclass(frozen=True) +class ConfigEntry: + """诊断扫描中的单个配置点。""" + + algo: str # e.g. "Ring" + proto: str # e.g. "Simple" + size_bytes: int # e.g. 1048576 + size_label: str # human-readable, e.g. "1M" + dtype: str # e.g. "float32" + nranks: int = 8 # number of ranks + collective: str = "allreduce" # allreduce or reducescatter + + def env_dict(self) -> dict[str, str]: + """返回此配置的 os.environ 风格字典。""" + return { + "NCCL_ALGO": self.algo, + "NCCL_PROTO": self.proto, + } + + def __str__(self) -> str: + return ( + f"Config(algo={self.algo:>14s}, proto={self.proto:>7s}, " + f"size={self.size_label:>6s}, dtype={self.dtype:>8s}, nranks={self.nranks})" + ) + + +@dataclass +class ConfigMatrix: + """生成 NCCL 确定性相关参数的完整笛卡尔积。""" + + algos: list[str] = field(default_factory=lambda: ["Ring", "Tree", "PAT"]) + protos: list[str] = field(default_factory=lambda: ["LL", "Simple"]) + dtypes: list[str] = field(default_factory=lambda: ["float32", "float16"]) + size_bytes: list[int] = field(default_factory=lambda: [ + 1 * 1024, # 1K — single-chunk region + 4 * 1024, # 4K — chunk transition boundary + 64 * 1024, # 64K + 1 * 1024 * 1024, # 1M — multi-chunk region + 16 * 1024 * 1024, # 16M — nranks chunks + 128 * 1024 * 1024, # 128M — full chunk saturation + ]) + nranks: int = 8 + collective: str = "allreduce" + + # ------------------------------------------------------------------ + # 工厂方法：根据不同场景构建合理的默认配置 + # ------------------------------------------------------------------ + + @classmethod + def quick_sweep(cls, nranks: int = 8) -> "ConfigMatrix": + """Minimal sweep — Ring vs Tree, float32, two sizes. ~2 min.""" + return cls( + algos=["Ring", "Tree"], + protos=["Simple"], + dtypes=["float32"], + size_bytes=[16 * 1024 * 1024, 128 * 1024 * 1024], + nranks=nranks, + ) + + @classmethod + def standard_sweep(cls, nranks: int = 8) -> "ConfigMatrix": + """Standard diagnostic sweep — all algos × protocols × 2 dtypes. ~10 min.""" + return cls( + algos=["Ring", "Tree", "PAT"], + protos=["LL", "Simple"], + dtypes=["float32", "float16"], + size_bytes=[4 * 1024, 64 * 1024, 1 * 1024 * 1024, + 16 * 1024 * 1024, 128 * 1024 * 1024], + nranks=nranks, + ) + + @classmethod + def exhaustive_sweep(cls, nranks: int = 8) -> "ConfigMatrix": + """Exhaustive sweep — all combinations including bfloat16. ~30 min.""" + return cls( + algos=["Ring", "Tree", "PAT"], + protos=["LL", "Simple"], + dtypes=["float32", "float16", "bfloat16"], + size_bytes=[1 * 1024, 4 * 1024, 64 * 1024, + 1 * 1024 * 1024, 16 * 1024 * 1024, + 128 * 1024 * 1024], + nranks=nranks, + ) + + @classmethod + def from_cli(cls, algo: Optional[str] = None, proto: Optional[str] = None, + dtype: Optional[str] = None, size: Optional[str] = None, + nranks: int = 8) -> "ConfigMatrix": + """Build matrix from CLI overrides.""" + algos = [algo] if algo else ["Ring", "Tree", "PAT"] + protos = [proto] if proto else ["LL", "Simple"] + dtypes = [dtype] if dtype else ["float32", "float16"] + + if size: + sizes = [_parse_size(size)] + else: + sizes = [4 * 1024, 64 * 1024, 1 * 1024 * 1024, + 16 * 1024 * 1024, 128 * 1024 * 1024] + + return cls(algos=algos, protos=protos, dtypes=dtypes, + size_bytes=sizes, nranks=nranks) + + # ------------------------------------------------------------------ + # Iteration + # ------------------------------------------------------------------ + + def generate(self, hw: HardwareCaps | None = None) -> list[ConfigEntry]: + """Generate all config combinations, filtered by collective support + and optionally by hardware compatibility. + + Parameters + ---------- + hw : HardwareCaps, optional + If provided, unsupported algo/proto combos are silently skipped. + """ + entries: list[ConfigEntry] = [] + for algo, proto, dtype, sz in itertools.product( + self.algos, self.protos, self.dtypes, self.size_bytes + ): + # Skip combinations unsupported by this collective + if self.collective not in ALGO_COLLECTIVE_SUPPORT.get(algo, set()): + continue + + # Hardware compatibility filter + if hw is not None: + warnings = hw.check_and_warn(algo, proto) + if warnings: + # Skipping incompatible config + continue + + entries.append(ConfigEntry( + algo=algo, + proto=proto, + size_bytes=sz, + size_label=_format_size(sz), + dtype=dtype, + nranks=self.nranks, + collective=self.collective, + )) + return entries + + def generate_with_warnings( + self, hw: HardwareCaps + ) -> tuple[list[ConfigEntry], list[str]]: + """Like generate(), but also returns a list of skipped-config warnings.""" + entries: list[ConfigEntry] = [] + skipped_warnings: list[str] = [] + for algo, proto, dtype, sz in itertools.product( + self.algos, self.protos, self.dtypes, self.size_bytes + ): + if self.collective not in ALGO_COLLECTIVE_SUPPORT.get(algo, set()): + continue + + hw_warnings = hw.check_and_warn(algo, proto) + if hw_warnings: + skipped_warnings.extend( + f" Skipped {algo}/{proto}/{dtype}/{_format_size(sz)}: {w}" + for w in hw_warnings + ) + continue + + entries.append(ConfigEntry( + algo=algo, proto=proto, size_bytes=sz, + size_label=_format_size(sz), dtype=dtype, + nranks=self.nranks, collective=self.collective, + )) + return entries, skipped_warnings + + def __len__(self) -> int: + return len(self.generate()) + + def __iter__(self): + return iter(self.generate()) + + def summary(self) -> str: + """Human-readable summary of the sweep space.""" + entries = self.generate() + return ( + f"Config sweep: {len(entries)} combinations\n" + f" Algorithms : {self.algos}\n" + f" Protocols : {self.protos}\n" + f" Data types : {self.dtypes}\n" + f" Size range : {_format_size(min(self.size_bytes))} ~ " + f"{_format_size(max(self.size_bytes))} " + f"({len(self.size_bytes)} points)\n" + f" N ranks : {self.nranks}\n" + f" Collective : {self.collective}" + ) + + +# --------------------------------------------------------------------------- +# ---- 工具函数 ---- +# --------------------------------------------------------------------------- + +def _format_size(nbytes: int) -> str: + """Format byte count as human-readable string.""" + if nbytes < 1024: + return f"{nbytes}B" + if nbytes < 1024 * 1024: + return f"{nbytes // 1024}K" + return f"{nbytes // (1024 * 1024)}M" + + +def _parse_size(s: str) -> int: + """Parse human-readable size string to bytes. e.g. '128M' → 134217728.""" + s = s.strip().upper() + multipliers = {"B": 1, "K": 1024, "M": 1024**2, "G": 1024**3} + for suffix, mult in multipliers.items(): + if s.endswith(suffix): + return int(s[:-1]) * mult + return int(s) diff --git a/src/code/issue4/data_generator.py b/src/code/issue4/data_generator.py new file mode 100644 index 0000000..7a040ef --- /dev/null +++ b/src/code/issue4/data_generator.py @@ -0,0 +1,184 @@ +""" +确定性数据生成器 — NCCL 确定性诊断。 + +使用固定 seed 按 rank 生成输入张量，确保: + - 每个 rank 有不同但确定的输入（暴露归约顺序敏感性） + - 相同 seed 始终生成相同张量（支持跨运行比对） + +核心设计决策: 每个 rank 使用不同 seed，使得归约顺序影响结果。 +若所有 rank 输入相同，任何累加顺序都会得到相同结果—— +诊断工具将永远捕捉不到非确定性。 +""" + +from __future__ import annotations + +import hashlib +import struct +from dataclasses import dataclass +from typing import Optional + +import numpy as np + + +# --------------------------------------------------------------------------- +# Data layout constants (matching nccl-tests buffer semantics) +# --------------------------------------------------------------------------- + +# For AllReduce: each rank contributes `count` elements, output is `count` elements. +# For ReduceScatter: each rank contributes `nranks * count` elements, +# output per rank is `count` elements. +# nccl-tests "size" = nranks × count × sizeof(dtype) for ReduceScatter. +# nccl-tests "size" = count × sizeof(dtype) for AllReduce. + +DTYPE_NP_MAP: dict[str, type] = { + "float32": np.float32, + "float16": np.float16, + "bfloat16": np.float16, # numpy doesn't have bf16 natively; use fp16 as proxy for shape +} + +DTYPE_BYTES: dict[str, int] = { + "float32": 4, + "float16": 2, + "bfloat16": 2, +} + + +@dataclass +class DataGenerator: + """为 NCCL 诊断运行生成确定性输入张量。 + + Parameters + ---------- + dtype : str + One of 'float32', 'float16', 'bfloat16'. + base_seed : int + Root seed. Per-rank seeds are derived as base_seed + rank. + collective : str + 'allreduce' or 'reducescatter' — affects element count calculation. + """ + + dtype: str = "float32" + base_seed: int = 42 + collective: str = "allreduce" + + # ------------------------------------------------------------------ + # ---- 元素计数（匹配 nccl-tests 约定） ---- + # ------------------------------------------------------------------ + + def elem_count(self, size_bytes: int, nranks: int) -> int: + """根据给定总字节数计算每 rank 的元素数。 + + Matching nccl-tests convention: + - AllReduce: count = size_bytes / sizeof(dtype) + - ReduceScatter: count = size_bytes / (sizeof(dtype) * nranks) + """ + type_bytes = DTYPE_BYTES[self.dtype] + if self.collective == "allreduce": + count = size_bytes // type_bytes + elif self.collective == "reducescatter": + count = size_bytes // (type_bytes * nranks) + else: + raise ValueError(f"Unknown collective: {self.collective}") + + if count <= 0: + raise ValueError( + f"size_bytes={size_bytes} too small for dtype={self.dtype} " + f"({type_bytes} bytes/elem) × nranks={nranks} " + f"(collective={self.collective})" + ) + return count + + # ------------------------------------------------------------------ + # ---- 数据生成 ---- + # ------------------------------------------------------------------ + + def generate(self, rank: int, size_bytes: int, nranks: int) -> np.ndarray: + """为单个 rank 生成确定性数据。 + + Uses SHA-256 seeded PRNG: seed → deterministic bits → normalized to [−1, +1] range. + This avoids any platform-dependent RNG behavior across runs. + + For AllReduce: produces `count` elements (the per-rank contribution). + For ReduceScatter: produces `nranks * count` elements (the per-rank send buffer, + of which `count` elements are scattered to this rank after reduction). + """ + count = self.elem_count(size_bytes, nranks) + # ReduceScatter: each rank contributes nranks * recvcount elements as send buffer + if self.collective == "reducescatter": + count = count * nranks + seed = self.base_seed + rank + data = _deterministic_rand(count, seed) + + np_dtype = DTYPE_NP_MAP[self.dtype] + return data.astype(np_dtype) + + def generate_all(self, size_bytes: int, nranks: int) -> list[np.ndarray]: + """为所有 rank 生成输入。 Returns list[np.ndarray] indexed by rank.""" + return [self.generate(r, size_bytes, nranks) for r in range(nranks)] + + # ------------------------------------------------------------------ + # Reference computation (CPU ground truth for validation) + # ------------------------------------------------------------------ + + def reference_allreduce(self, inputs: list[np.ndarray]) -> np.ndarray: + """Compute the reference AllReduce result on CPU in float64. + + This is the "ideal" result — a single-precision floating sum of all + rank inputs. Not bitwise-identical to any GPU result but serves as a + sanity bound for diff magnitude. + """ + acc = np.zeros_like(inputs[0], dtype=np.float64) + for arr in inputs: + acc += arr.astype(np.float64) + return acc + + +# --------------------------------------------------------------------------- +# Internal: deterministic pseudo-random generator using SHA-256 +# --------------------------------------------------------------------------- + +def _deterministic_rand(count: int, seed: int) -> np.ndarray: + """Generate count floats in [−1, 1] using SHA-256 as the entropy source. + + This is fully deterministic across all platforms and Python versions + because SHA-256 is a standardized cryptographic hash. + """ + data = np.empty(count, dtype=np.float32) + generated = 0 + block = 0 + + while generated < count: + # Hash (seed || block) → 32 bytes → 8× float32 + h = hashlib.sha256() + h.update(struct.pack("= count: + break + # Interpret 4 bytes as uint32, scale to [−1, 1] + val = struct.unpack(" bool: + """检查各 rank 输入是否不同（有意义诊断的必要条件）。 + + Returns True if at least one rank has different data from another. + """ + if len(inputs) < 2: + return False + ref = inputs[0] + for arr in inputs[1:]: + if not np.array_equal(ref, arr): + return True + return False diff --git a/src/code/issue4/diagnose.py b/src/code/issue4/diagnose.py new file mode 100644 index 0000000..2b552b0 --- /dev/null +++ b/src/code/issue4/diagnose.py @@ -0,0 +1,723 @@ +#!/usr/bin/env python3 +""" +NCCL Bitwise 可复现性诊断工具 — 主入口 + +完整诊断流水线编排: + 1. 环境预检（GPU 数量、NCCL、PyTorch） + 2. 生成配置扫描矩阵（ALGO × PROTO × SIZE × DTYPE） + 3. 数据生成器验证 + 4. 每个配置运行 N 次（PyTorch NCCL），支持断点恢复 + 5. 逐位比对 + 差异分布分析 + 6. 差异随规模和迭代的演化追踪 + 7. 生成报告 + 确定性配置建议 + +用法: + python -m issue4.diagnose --mode quick --nranks 8 + python -m issue4.diagnose --mode quick --list-configs + python -m issue4.diagnose --algo Ring --proto Simple --dtype float32 --size 128M + python -m issue4.diagnose --json report.json + python -m issue4.diagnose --self-test + python -m issue4.diagnose --inject-difference auto --nranks 4 + +依赖: PyTorch >= 1.12（含 NCCL 后端）、numpy、CUDA GPU >= 2 +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from pathlib import Path +from typing import Optional + +from .config_matrix import ConfigMatrix, ConfigEntry, HardwareCaps, DTYPE_SENSITIVITY +from .data_generator import DataGenerator, validate_data_divergence +from .runner import NcclRunner, RunResult, MultiRunResult +from .comparator import ( + BitwiseComparator, DiffReport, EvolutionReport, IterEvolutionReport, + ThreeLevelSummary, + compute_ulp, analyze_xor_float, +) +from .reporter import Reporter, DiagnosticSummary + + +def main(argv: Optional[list[str]] = None) -> int: + """Main entry point. Returns 0 on success, 1 on pre-flight failure.""" + import multiprocessing as _mp + _mp.set_start_method("spawn", force=True) + + args = _parse_args(argv) + + # ------------------------------------------------------------------ + # Self-test mode: run GPU integration test and exit + # ------------------------------------------------------------------ + if args.self_test: + return _run_self_test(args.nranks) + + # ------------------------------------------------------------------ + # ---- 阶段 0：环境预检 ---- + # ------------------------------------------------------------------ + print("Phase 0: Environment pre-flight check...") + env_ok, env_msg = _check_environment(args.nranks) + print(env_msg) + if not env_ok: + print(" ABORT: Environment check failed. Fix the issues above and retry.") + return 1 + + # Detect hardware capabilities for compatibility filtering + hw = HardwareCaps.detect() + print(f" SM compute : {hw.sm_major} ({hw.gpu_name})") + print(f" NVSwitch : {'yes' if hw.has_nvswitch else 'no'}") + print(f" Collnet/SHARP: {'yes' if hw.has_collnet else 'no'}") + print(f" LL128 support: {'yes' if hw.supports_ll128() else 'no — requires Hopper (SM90+)'}") + print() + + # ------------------------------------------------------------------ + # ---- 阶段 1：构建配置矩阵 ---- + # ------------------------------------------------------------------ + print("Phase 1: Building configuration sweep matrix...") + matrix = _build_matrix(args) + configs, hw_warnings = matrix.generate_with_warnings(hw) + + # Print hardware-incompatible skipped configs + if hw_warnings: + print(" Hardware-incompatible configs skipped:") + for w in hw_warnings: + print(w) + print() + + if not configs: + print(" ERROR: No valid configs generated (all filtered by collective/hardware).") + return 1 + print(matrix.summary()) + print() + + if args.list_configs: + print("Configs to be executed:") + for i, cfg in enumerate(configs): + print(f" [{i:3d}] {cfg}") + print(f"\n ({len(configs)} configs total, ~{len(configs)*2} trials)") + return 0 + + # ------------------------------------------------------------------ + # ---- 阶段 1.5：ULP 注入验证 ---- (if --inject-difference) + # ------------------------------------------------------------------ + if args.inject_difference is not None: + print("Phase 1.5: ULP injection pipeline validation...") + ok = _run_injection_test(args, hw) + if ok: + print(" ✓ Injection pipeline verified — detector correctly identifies injected diffs.") + else: + print(" ✗ Injection pipeline FAILED — detector missed an injected difference!") + print() + if not args.force: + print(" Exiting (injection-only mode). Use --force to continue to full sweep.") + return 0 if ok else 1 + + # ------------------------------------------------------------------ + # ---- 阶段 2：验证数据生成 ---- + # ------------------------------------------------------------------ + print("Phase 2: Validating data generator...") + gen = DataGenerator( + dtype=configs[0].dtype, + base_seed=42, + collective=matrix.collective, + ) + test_inputs = gen.generate_all(configs[0].size_bytes, matrix.nranks) + if not validate_data_divergence(test_inputs): + print(" WARNING: All ranks have identical input data. " + "Non-determinism may not be detectable.") + print(" Consider using per-rank seeds (default behavior).") + else: + print(f" OK: {matrix.nranks} ranks, each with distinct deterministic input.") + print() + + # ------------------------------------------------------------------ + # ---- 阶段 3：运行扫描 ---- + # ------------------------------------------------------------------ + all_results: list[list[RunResult]] = [] + all_full_results: list[list[MultiRunResult]] = [] + start_idx = 0 + + use_three_level = args.n_calls > 1 + mode_label = f" Backend: pytorch, {args.n_calls} calls per run, " + mode_label += "three-level comparison (run / call / rank)" if use_three_level else "rank-0 only" + + if start_idx > 0: + print(f"Phase 3: Resuming from checkpoint ({start_idx}/{len(configs)} done)...") + else: + print(f"Phase 3: Running {len(configs)} configs × {args.trials} trials each...") + print(mode_label) + print() + + runner = NcclRunner(nranks=matrix.nranks, backend="pytorch") + from .runner import MultiRunResult as MRR + + t_start = time.perf_counter() + i = start_idx + for i in range(start_idx, len(configs)): + cfg = configs[i] + print(f" [{i + 1}/{len(configs)}] {cfg}") + try: + if use_three_level: + full = runner.run_full(cfg, n_calls=args.n_calls, trials=args.trials) + all_full_results.append(full) + if len(full) >= 2: + checksums = [ + [float(np.sum(o.astype(np.float64))) + for o in c0.values()] + for c0 in full[0].call_outputs + ] + print(f" {args.n_calls} calls × {matrix.nranks} ranks captured") + else: + results = runner.run(cfg, trials=args.trials) + all_results.append(results) + if len(results) >= 2: + c1, c2 = results[0].checksum(), results[1].checksum() + if abs(c1 - c2) > 0: + print(f" checksum diff detected: {abs(c1 - c2):.6e}") + except Exception as exc: + print(f" ERROR: {exc}") + if use_three_level: + all_full_results.append([]) + else: + all_results.append([]) + + t_end = time.perf_counter() + elapsed = t_end - t_start + effective = max(i - start_idx + 1, 1) + print(f"\n Sweep complete in {elapsed:.1f}s " + f"({elapsed / effective:.1f}s per config)") + print() + + # ------------------------------------------------------------------ + # ---- 阶段 4-6：比对 + 报告（分支：二级 vs 三级） ---- + # ------------------------------------------------------------------ + comp = BitwiseComparator(tolerance="bitwise") + + if use_three_level: + print("Phase 4: Three-level comparison (run-vs-run / call-vs-call / rank-vs-rank)...") + all_summaries: list[ThreeLevelSummary] = [] + for i, (full_trials, cfg) in enumerate(zip(all_full_results, configs)): + if len(full_trials) >= 2: + label = f"{cfg.algo}/{cfg.proto}/{cfg.dtype}/{cfg.size_label}" + summary = comp.full_three_level_report(full_trials, label=label) + all_summaries.append(summary) + status = "CLEAN" if summary.all_clean else "DIFFS" + print(f" [{i+1}] {label}: {status}") + + # Print per-level details for first failing config + for s in all_summaries: + if not s.all_clean: + print(f"\n Detailed breakdown: {s.label}") + print(s.summary()) + # Show first failing cell + for level_name, level_data in [ + ("run-vs-run", s.run_vs_run), + ("call-vs-call", s.call_vs_call), + ("rank-vs-rank", s.rank_vs_rank), + ]: + for key, reports in level_data.items(): + for r in reports: + if not r.bitwise_match: + print(f" {level_name} {key}: {r.diff_count} diffs " + f"({r.diff_ratio*100:.4f}%), " + f"ULP={r.xor_detail.ulp_distance if r.xor_detail else '?'}") + print(f" offset={r.first_diff_offset} " + f"max_abs={r.max_abs_diff:.6e}") + break + + # ---- JSON 输出 ---- + if args.json: + import json as _json + data = { + "config_count": len(configs), + "full_configs": [], + "summary": { + "clean": sum(1 for s in all_summaries if s.all_clean), + "with_diffs": sum(1 for s in all_summaries if not s.all_clean), + }, + } + for s in all_summaries: + cfg_entry = { + "label": s.label, + "all_clean": s.all_clean, + "run_vs_run": [], + "call_vs_call": {}, + "rank_vs_rank": [], + } + for key, reports in s.run_vs_run.items(): + for r in reports: + cfg_entry["run_vs_run"].append({ + "key": key, + "bitwise_match": r.bitwise_match, + "diff_count": r.diff_count, + "diff_ratio": f"{r.diff_ratio*100:.4f}%", + "ulp": r.xor_detail.ulp_distance if r.xor_detail else -1, + }) + for rank, reports in s.call_vs_call.items(): + diffs = [] + for r in reports: + diffs.append({ + "diff_count": r.diff_count, + "diff_ratio": f"{r.diff_ratio*100:.4f}%", + }) + cfg_entry["call_vs_call"][str(rank)] = diffs + for key, reports in s.rank_vs_rank.items(): + for r in reports: + cfg_entry["rank_vs_rank"].append({ + "key": key, + "bitwise_match": r.bitwise_match, + "diff_count": r.diff_count, + "diff_ratio": f"{r.diff_ratio*100:.4f}%", + }) + data["full_configs"].append(cfg_entry) + with open(args.json, "w", encoding="utf-8") as f: + _json.dump(data, f, indent=2) + print(f"\n JSON report written to: {args.json}") + + else: + # Legacy two-level path (unchanged) + print("Phase 4: Bitwise comparison + distribution analysis...") + reports: list[DiffReport] = [] + for i, (trials, cfg) in enumerate(zip(all_results, configs)): + if len(trials) >= 2: + report = comp.compare(trials[0], trials[1]) + reports.append(report) + else: + # 未产生有效结果（worker 崩溃等）→ 标记为未运行，不是 PASS + reports.append(DiffReport(config=cfg, bitwise_match=False, total_elements=0)) + + print("Phase 5: Tracking diff evolution across sizes and iterations...") + evolutions = comp.track_evolution(all_results, configs) + iter_evo: Optional[IterEvolutionReport] = None + worst_cfg_idx = _find_worst_config(all_results) + if worst_cfg_idx is not None and len(all_results[worst_cfg_idx]) >= 2: + iter_evo = comp.track_iter_evolution(all_results[worst_cfg_idx]) + print() + + print("Phase 6: Generating report...") + reporter = Reporter(output_json=args.json is not None, + output_console=True, + json_path=args.json or "") + summary = reporter.report(reports, evolutions, args.json) + if iter_evo and iter_evo.iteration_diffs: + print() + print(iter_evo.summary()) + if summary.diff_count > 0: + print("Diagnosis: Non-determinism detected! See recommendations above.") + else: + print("Diagnosis: All configurations bitwise-deterministic on this hardware.") + + # ------------------------------------------------------------------ + # ---- 清理 + 返回 ---- + # ------------------------------------------------------------------ + return 0 + + +# --------------------------------------------------------------------------- +# ---- 环境预检 ---- +# --------------------------------------------------------------------------- + +def _check_environment(nranks: int) -> tuple[bool, str]: + """Check that the environment is capable of running the diagnostic. + + Returns (ok: bool, message: str). + """ + lines: list[str] = [] + + # 1. Python version + py_ver = f"{sys.version_info.major}.{sys.version_info.minor}" + lines.append(f" Python : {py_ver}") + + # 2. PyTorch + CUDA + try: + import torch + lines.append(f" PyTorch : {torch.__version__}") + cuda_avail = torch.cuda.is_available() + lines.append(f" CUDA avail : {cuda_avail}") + if cuda_avail: + gpu_count = torch.cuda.device_count() + lines.append(f" GPU count : {gpu_count}") + for g in range(min(gpu_count, 4)): + lines.append(f" GPU {g}: {torch.cuda.get_device_name(g)}") + if gpu_count > 4: + lines.append(f" ... and {gpu_count - 4} more") + else: + return False, "\n".join(lines + [" FAIL: CUDA not available."]) + except ImportError: + return False, "\n".join(lines + [" FAIL: PyTorch not installed."]) + + # 3. NCCL backend + try: + import torch.distributed as dist + lines.append(f" NCCL avail : {dist.is_nccl_available()}") + if not dist.is_nccl_available(): + return False, "\n".join(lines + [" FAIL: NCCL backend not available. " + "Rebuild PyTorch with NCCL support."]) + except Exception: + return False, "\n".join(lines + [" FAIL: torch.distributed unavailable."]) + + # 4. GPU count vs requested ranks + gpu_count = torch.cuda.device_count() + if nranks > gpu_count: + return False, "\n".join(lines + [ + f" FAIL: Requested {nranks} ranks but only {gpu_count} GPUs available.\n" + f" Use --nranks {gpu_count} or reduce the rank count." + ]) + if nranks < 2: + return False, "\n".join(lines + [ + " FAIL: At least 2 ranks required for meaningful diagnostic." + ]) + + # 5. NumPy + import numpy as np + lines.append(f" NumPy : {np.__version__}") + + lines.append(" ✓ Environment OK") + return True, "\n".join(lines) + + +# --------------------------------------------------------------------------- +# ---- 杂项 ---- +# --------------------------------------------------------------------------- + +def _find_worst_config(results: list[list[RunResult]]) -> Optional[int]: + """Find index of config with largest checksum delta between trials.""" + worst_idx = None + worst_delta = -1.0 + for i, trials in enumerate(results): + if len(trials) >= 2: + delta = abs(trials[0].checksum() - trials[1].checksum()) + if delta > worst_delta: + worst_delta = delta + worst_idx = i + return worst_idx + + +# --------------------------------------------------------------------------- +# ---- 命令行参数 ---- +# --------------------------------------------------------------------------- + +def _parse_args(argv: Optional[list[str]]) -> argparse.Namespace: + p = argparse.ArgumentParser( + description="NCCL Bitwise Reproducibility Diagnostic Tool", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python -m issue4.diagnose --mode quick + python -m issue4.diagnose --mode quick --list-configs + python -m issue4.diagnose --mode standard --json report.json + python -m issue4.diagnose --algo Ring --size 128M + python -m issue4.diagnose --inject-difference auto --nranks 4 + """, + ) + p.add_argument("--mode", choices=["quick", "standard", "exhaustive"], + default="quick", + help="Sweep mode: quick (~2 min), standard (~10 min), exhaustive (~30 min)") + p.add_argument("--nranks", type=int, default=8, + help="Number of GPU ranks (default: 8)") + p.add_argument("--trials", type=int, default=2, + help="Number of run-to-run trials per config (default: 2)") + p.add_argument("--self-test", action="store_true", + help="Run minimal GPU integration test and exit") + p.add_argument("--n-calls", type=int, default=1, + help="Number of NCCL calls per process group (enables " + "three-level run/call/rank comparison)") + p.add_argument("--algo", type=str, default=None) + p.add_argument("--proto", type=str, default=None) + p.add_argument("--dtype", type=str, default=None) + p.add_argument("--size", type=str, default=None) + p.add_argument("--collective", choices=["allreduce", "reducescatter"], + default="allreduce") + p.add_argument("--json", type=str, default=None, + help="Write JSON report to this file path") + p.add_argument("--list-configs", action="store_true", + help="Print the config sweep matrix and exit (no GPU run)") + p.add_argument("--inject-difference", type=str, default=None, metavar="SPEC", + help="Run ULP injection validation. SPEC: 'auto' or 'offset=42,magnitude=3'") + p.add_argument("--force", action="store_true", + help="With --inject-difference: continue to sweep after validation") + return p.parse_args(argv) + + +# --------------------------------------------------------------------------- +# ---- ULP 注入管线 ---- +# --------------------------------------------------------------------------- + +def _run_injection_test( + args: argparse.Namespace, hw: HardwareCaps +) -> bool: + """Validate the diagnostic pipeline by injecting a known bit-level difference. + + How it works: + 1. Generate deterministic input for a single config + 2. Run the collective once → baseline output + 3. Inject a known ULP difference into the baseline output (simulates a "diff") + 4. Run the comparator between baseline and injected → should DETECT the diff + 5. Verify the XOR/ULP analysis matches the injection + + This end-to-end test confirms that the entire pipeline correctly + detects bit-level non-determinism. + """ + import struct + + # Build a minimal single-config matrix for injection test + algo = args.algo or "Ring" + proto = args.proto or "Simple" + dtype = args.dtype or "float32" + size_str = args.size or "128M" + nranks = args.nranks + + matrix = ConfigMatrix.from_cli( + algo=algo, proto=proto, dtype=dtype, size=size_str, nranks=nranks + ) + matrix.collective = args.collective + configs = matrix.generate(hw) + if not configs: + print(" SKIP: No valid config for injection test (hardware-incompatible).") + return True # not a failure — hardware limitation + + cfg = configs[0] + print(f" Config: {cfg}") + + # Phase A: Run baseline + runner = NcclRunner(nranks=nranks, backend="pytorch") + results = runner.run(cfg, trials=1) + if not results: + print(" FAIL: Runner produced no results.") + return False + baseline = results[0] + baseline_arr = baseline.output.copy() + print(f" Baseline checksum: {baseline.checksum():.6f}") + + # Phase B: Parse injection spec + inj_offset, inj_magnitude_ulp = _parse_injection_spec( + args.inject_difference or "auto", len(baseline_arr) + ) + + # Phase C: Inject difference into baseline output copy + injected_arr = _inject_ulp( + baseline_arr, offset=inj_offset, ulp_magnitude=inj_magnitude_ulp + ) + + # Phase D: Compare baseline vs injected + from .runner import RunResult as RR + injected_result = RunResult( + config=cfg, output=injected_arr, elapsed_ms=0.0, rank=0, + ) + + comp = BitwiseComparator(tolerance="bitwise") + report = comp.compare(baseline, injected_result) + + # Phase E: Validate + if report.bitwise_match: + print(f" FAIL: Comparator reported BITWISE MATCH despite injected " + f"{inj_magnitude_ulp} ULP diff at offset {inj_offset}!") + return False + + print(f" Comparator detected {report.diff_count} differing elements " + f"(expected 1 at offset {inj_offset}).") + + if report.xor_detail: + xd = report.xor_detail + expected_xor = compute_ulp( + float(baseline_arr[inj_offset]), + float(injected_arr[inj_offset]), + ) + print(f" Injection : offset={inj_offset}, magnitude={inj_magnitude_ulp} ULP") + print(f" Detected : offset={xd.offset}, ULP={xd.ulp_distance}") + print(f" XOR detail : {xd.baseline_bits} → {xd.target_bits} " + f"(xor=0x{xd.xor_bits:08x}, flips={xd.n_total_bits_flipped})") + + if xd.offset != inj_offset: + print(f" FAIL: Offset mismatch — injected {inj_offset}, detected {xd.offset}") + return False + if xd.ulp_distance != expected_xor: + print(f" WARN: ULP mismatch — expected {expected_xor}, detected {xd.ulp_distance}") + # Not a hard failure — ULP can differ due to float32 internal rounding + else: + print(" FAIL: XOR detail not computed!") + return False + + return True + + +def _parse_injection_spec(spec: str, max_offset: int) -> tuple[int, int]: + """Parse injection specification string. + + Formats: + "auto" → offset=midpoint, magnitude=1 + "offset=42,magnitude=3" → explicit offset and ULP magnitude + "offset=42" → explicit offset, magnitude=1 + """ + offset = max_offset // 2 + magnitude = 1 + + if spec == "auto": + return (offset, magnitude) + + for part in spec.split(","): + part = part.strip() + if part.startswith("offset="): + offset = min(int(part.split("=")[1]), max_offset - 1) + elif part.startswith("magnitude="): + magnitude = max(1, int(part.split("=")[1])) + + return (offset, magnitude) + + +def _inject_ulp(arr: np.ndarray, offset: int, ulp_magnitude: int) -> np.ndarray: + """Inject a known ULP difference into an array at a specific offset. + + Returns a COPY of arr with the modification applied. + """ + import struct + + result = arr.copy() + original = float(result.flat[offset]) + + # Reinterpret float32 as int32 (signed); use unsigned for bit arithmetic + bits_signed = struct.unpack("= 0: + new_bits = min(bits_unsigned + ulp_magnitude, 0x7F7FFFFF) + else: + # 负浮点数：减小 bit pattern（走向更负） + new_bits = max(bits_unsigned - ulp_magnitude, 0x00800001) # min pos norm + + # Reinterpret int32 → float32 + injected = struct.unpack(" ConfigMatrix: + """Build ConfigMatrix from parsed CLI args.""" + if any([args.algo, args.proto, args.dtype, args.size]): + matrix = ConfigMatrix.from_cli( + algo=args.algo, proto=args.proto, + dtype=args.dtype, size=args.size, + nranks=args.nranks, + ) + else: + builders = { + "quick": ConfigMatrix.quick_sweep, + "standard": ConfigMatrix.standard_sweep, + "exhaustive": ConfigMatrix.exhaustive_sweep, + } + matrix = builders[args.mode](nranks=args.nranks) + + matrix.collective = args.collective + return matrix + + +# --------------------------------------------------------------------------- +# ---- 自检（GPU 集成冒烟测试） ---- +# --------------------------------------------------------------------------- + +def _run_self_test(nranks: int) -> int: + """Run a minimal GPU integration test and return 0 on pass, 1 on fail.""" + import struct + + print("=" * 60) + print(" NCCL Diagnostic Tool — Self-Test (GPU Integration)") + print("=" * 60) + + # 1. Environment check + try: + import torch + import torch.distributed as dist + except ImportError: + print(" FAIL: torch not installed") + return 1 + + gpu_count = torch.cuda.device_count() + print(f" GPUs: {gpu_count}") + if gpu_count < 2: + print(" SKIP: need >= 2 GPUs for meaningful test") + return 0 + + nranks = min(nranks, gpu_count) + print(f" Using {nranks} ranks") + + # 2. Hardware detection + hw = HardwareCaps.detect() + print(f" SM: {hw.sm_major} ({hw.gpu_name})") + if hw.supports_ll128(): + print(" LL128: supported") + else: + print(" LL128: NOT supported (requires Hopper+)") + + # 3. Run one config × 2 trials + from .config_matrix import ConfigEntry + from .data_generator import DataGenerator + from .runner import NcclRunner + from .comparator import BitwiseComparator + + size_bytes = 64 * 1024 # 64K + cfg = ConfigEntry(algo="Tree", proto="Simple", size_bytes=size_bytes, + size_label="64K", dtype="float32", nranks=nranks) + + print(f"\n Config: {cfg}") + print(" Running 2 trials...", end=" ", flush=True) + + try: + runner = NcclRunner(nranks=nranks, backend="pytorch") + results = runner.run(cfg, trials=2) + except Exception as e: + print(f"\n FAIL: {e}") + return 1 + + if len(results) < 2: + print("\n FAIL: not enough results") + return 1 + + print(f"done ({results[0].elapsed_ms:.2f}ms, {results[1].elapsed_ms:.2f}ms)") + + # 4. Compare + comp = BitwiseComparator() + report = comp.compare(results[0], results[1]) + + if report.bitwise_match: + print(" Bitwise comparison: IDENTICAL") + else: + print(f" Bitwise comparison: {report.diff_count} DIFFERENCES " + f"({report.diff_ratio*100:.4f}%)") + if report.xor_detail: + print(f" First diff ULP: {report.xor_detail.ulp_distance}") + + # 5. ULP injection validation + print("\n ULP injection test:", end=" ", flush=True) + baseline = results[0].output.copy() + offset = len(baseline) // 2 + bits = struct.unpack("= 0 else max(bits - 1, -0x7F7FFFFF) + injected = baseline.copy() + injected[offset] = struct.unpack("8 阶段流水线入口"] + subgraph DP["流水线阶段"] + direction LR + P0["P0
环境预检"] + P1["P1
配置矩阵"] + P15["P1.5
ULP 注入"] + P2["P2
数据生成"] + P3["P3
NCCL 扫描"] + P4["P4
逐位比对"] + P5["P5
演化追踪"] + P6["P6
报告输出"] + end + subgraph DU["内部函数"] + ENV["_check_environment()
GPU/NCCL/PyTorch 就绪检查"] + BUILD["_build_matrix()
CLI 参数 → ConfigMatrix"] + SELF["_run_self_test()
最小 GPU 冒烟测试"] + INJECT["_run_injection_test()
运行 → 注入 1 ULP →
比对 → 验证检出"] + INJULP["_inject_ulp()
在数组指定位置
注入 N 个 ULP 差异"] + INJPARSE["_parse_injection_spec()
解析注入参数字符串"] + end + MAIN --> P0 --> P1 --> P15 --> P2 --> P3 --> P4 --> P5 --> P6 + MAIN --> ENV + MAIN --> BUILD + MAIN --> SELF + MAIN --> INJECT + INJECT --> INJULP + INJECT --> INJPARSE + P15 -.-> INJECT + end + + %% ═══════════════ config_matrix.py ═══════════════ + subgraph C["config_matrix.py"] + HW["HardwareCaps
.detect() → SM / NVSwitch
.supports_ll128()
.check_and_warn(algo, proto)
.compatible_algos()
.compatible_protos()"] + CM["ConfigMatrix
.generate(hw)
.generate_with_warnings(hw)
.quick_sweep()
.standard_sweep()
.exhaustive_sweep()
.from_cli()"] + CE["ConfigEntry
algo proto dtype
size_bytes nranks
collective"] + HW -.->|"过滤用"| CM + CM -->|"生成"| CE + end + + %% ═══════════════ data_generator.py ═══════════════ + subgraph G["data_generator.py"] + DG["DataGenerator
.elem_count(size, nranks)
.generate(rank, size, nranks)
.generate_all(size, nranks)
.reference_allreduce(inputs)
内核: SHA-256 PRNG"] + VAL["validate_data_divergence()
验证各 rank 输入是否不同"] + DG --> VAL + end + + %% ═══════════════ runner.py ═══════════════ + subgraph R["runner.py"] + RNR["NcclRunner
.run(cfg, trials) → list[RunResult]
.run_full(cfg, n_calls, trials)
→ list[MultiRunResult]
.sweep(configs, trials)
mp.Process() × N"] + RR["RunResult
.output (rank-0)
.checksum()
.save() / .load()"] + MRR["MultiRunResult
.call_outputs[call][rank]
.elapsed_ms"] + RNR --> RR + RNR --> MRR + end + + %% ═══════════════ comparator.py ═══════════════ + subgraph M["comparator.py"] + BC["BitwiseComparator
.compare(a, b) → DiffReport
.compare_run_vs_run(r0, r1)
.compare_call_vs_call(run)
.compare_rank_vs_rank(run, c)
.full_three_level_report()
→ ThreeLevelSummary
.track_evolution()
.track_iter_evolution()"] + subgraph MR["报告结构"] + DR["DiffReport
offset, max/mean/std
分布, XOR 详情"] + DD["DiffDistribution
前/中/后三区聚类
差异直方图"] + XD["XorDetail
sign/exp/mantissa
ULP 距离"] + EV["EvolutionReport
差异率 vs 数据规模"] + IV["IterEvolutionReport
差异率 vs 迭代
单调增长检测"] + T3["ThreeLevelSummary
run 间 + call 间
+ rank 间"] + end + subgraph MU["位运算工具"] + ULP["compute_ulp(a, b)
符号-幅度 → ULP"] + XOR["analyze_xor_float(a, b)
struct 打包 →
uint32 XOR 分解"] + end + BC --> DR --> DD + DR --> XD + BC --> T3 + BC --> EV + BC --> IV + BC -.-> ULP + BC -.-> XOR + end + + %% ═══════════════ reporter.py ═══════════════ + subgraph P["reporter.py"] + RP["Reporter
.report(reports, evolutions)
→ DiagnosticSummary
._console_report()
._json_report()
._generate_recommendations()
._find_deterministic_config()"] + DS["DiagnosticSummary
配置通过/失败统计
最差差异
确定性配置建议"] + RP --> DS + end + + %% ═══════════════ 跨模块依赖 ═══════════════ + MAIN ==>|"导入"| HW + MAIN ==>|"导入"| CM + MAIN ==>|"导入"| DG + MAIN ==>|"导入"| RNR + MAIN ==>|"导入"| BC + MAIN ==>|"导入"| RP + RNR -.->|import| C + RNR -.->|import| G + BC -.->|import| C + BC -.->|import| R + RP -.->|import| M + RP -.->|import| C + +``` + +--- + +# 模块依赖关系 (简化版) + +```mermaid +graph LR + D["diagnose.py"] -->|import| C["config_matrix.py"] + D -->|import| G["data_generator.py"] + D -->|import| R["runner.py"] + D -->|import| M["comparator.py"] + D -->|import| P["reporter.py"] + R --> C + R --> G + M --> C + M --> R + P --> M + P --> C +``` diff --git a/src/code/issue4/docs/02-data-flow.md b/src/code/issue4/docs/02-data-flow.md new file mode 100644 index 0000000..ead139b --- /dev/null +++ b/src/code/issue4/docs/02-data-flow.md @@ -0,0 +1,73 @@ +# 数据流图 — 全链路流水线 + +```mermaid +--- +title: 端到端流水线（--n-calls=1 传统模式 vs --n-calls>1 三级对比模式） +--- +flowchart TD + CLI["命令行入口
--mode quick --nranks 8
[--n-calls 5] [--json report.json]"] + + CLI --> P0["阶段 0: 环境预检
_check_environment() → GPU/NCCL/PyTorch
HardwareCaps.detect() → SM/NVSwitch/Collnet
⎯ GPU < 2 或无 NCCL 时终止"] + P0 --> P1["阶段 1: 配置矩阵
ConfigMatrix.generate_with_warnings(hw)
硬件过滤: 非 Hopper 剔除 LL128
无 NVSwitch 剔除 NVLS"] + P1 --> LIST{"--list-configs?"} + LIST -->|是| EXIT0["打印配置列表并退出"] + LIST -->|否| P1_5{"--inject-difference?"} + + P1_5 -->|是| INJ["阶段 1.5: ULP 注入验证
基准 = .run() → 注入 1 ULP → .compare()
验证检出位置是否正确"] + INJ --> FORCE{"--force?"} + FORCE -->|否| EXIT1["退出（仅验证模式）"] + FORCE -->|是| P2 + + P1_5 -->|否| P2["阶段 2: 数据生成
DataGenerator(base_seed=42).generate_all()
SHA-256(seed+rank) → fp32 范围 [−1,+1]
每 rank 不同 seed → 归约顺序敏感"] + + P2 --> P3_BRANCH{"--n-calls?"} + + P3_BRANCH -->|"=1 (传统)"| P3A["阶段 3 — 路径 A
NcclRunner.run(cfg, trials=2)
2 个独立 mp.Process 进程组
每次: init NCCL → 预热 → all_reduce
→ 仅 rank-0 输出
→ list[RunResult]"] + + P3_BRANCH -->|">1 (三级对比)"| P3B["阶段 3 — 路径 B
NcclRunner.run_full(cfg, n_calls, trials=2)
1 个进程组，内部 N 次调用
所有 rank 记录每次调用输出
→ list[MultiRunResult]
MultiRunResult.call_outputs[call][rank]"] + + P3A --> P4A["阶段 4 — 路径 A
BitwiseComparator.compare(trial0, trial1)
baseline != target → DiffReport
+ DiffDistribution + XorDetail"] + + P3B --> P4B["阶段 4 — 路径 B
BitwiseComparator.full_three_level_report()
run 间 | call 间 | rank 间
→ ThreeLevelSummary"] + + P4A --> P5["阶段 5-6: 报告与诊断
Reporter.report() → DiagnosticSummary
控制台表格 + JSON + 诊断引擎
_generate_recommendations()
_find_deterministic_config()"] + P4B --> P5 +``` + +--- + +# 三级对比矩阵 + +```mermaid +--- +title: 阶段 4 路径 B — run 间 / call 间 / rank 间 +--- +graph LR + subgraph T0["第 0 次运行 (MultiRunResult)"] + direction TB + M0["c0: [r0,r1,...,r7]"] + M1["c1: [r0,r1,...,r7]"] + M2["c2: [r0,r1,...,r7]"] + M3["..."] + end + + subgraph T1["第 1 次运行 (MultiRunResult)"] + direction TB + N0["c0: [r0,r1,...,r7]"] + N1["c1: [r0,r1,...,r7]"] + N2["c2: [r0,r1,...,r7]"] + N3["..."] + end + + RUN["compare_run_vs_run()
运行 0[call_k][rank_r] vs
运行 1[call_k][rank_r]"] + CALL["compare_call_vs_call()
运行 0[call_0][rank_r] vs
运行 0[call_i][rank_r]"] + RANK["compare_rank_vs_rank()
运行 0[call_k][rank_0] vs
运行 0[call_k][rank_i]"] + + T0 --> RUN + T1 --> RUN + T0 --> CALL + T0 --> RANK + RUN --> T3["ThreeLevelSummary
{run间, call间, rank间}"] + CALL --> T3 + RANK --> T3 +``` diff --git a/src/code/issue4/docs/03-runner-xor.md b/src/code/issue4/docs/03-runner-xor.md new file mode 100644 index 0000000..5d21fcd --- /dev/null +++ b/src/code/issue4/docs/03-runner-xor.md @@ -0,0 +1,63 @@ +# Runner 进程模型 — _pytorch_worker_full() + +```mermaid +sequenceDiagram + actor Parent as NcclRunner.run_full() + participant Manager as mp.Manager + participant R0 as Process (rank=0) + participant R1 as Process (rank=1) + participant RN as Process (rank=N-1) + + Parent->>Manager: 创建共享字典 + + loop 每个 rank + Parent->>R0: spawn Process(target=_pytorch_worker_full) + activate R0 + R0->>R0: os.environ[NCCL_ALGO] = algo + R0->>R0: os.environ[NCCL_PROTO] = proto + R0->>R0: os.environ[CUDA_VISIBLE_DEVICES] = str(rank) + R0->>R0: torch.cuda.set_device(0) + R0->>R1: dist.init_process_group(backend='nccl', device_id=torch.device('cuda:0')) + R1-->>R0: 已连接 + R0->>R1: dist.init_process_group(backend='nccl') + RN-->>R0: 已连接 + end + + Note over R0,RN: 预热: 克隆输入 → all_reduce → cuda.sync() + + loop call_idx = 0 .. n_calls-1 + R0->>R0: tensor = clone(input_data) + R0->>R0: cuda.synchronize() + R0->>R1: dist.all_reduce(tensor) + RN->>R0: all_reduce 完成 + R0->>R0: cuda.synchronize() + R0->>Manager: result_dict["{call}_{rank}"] = tensor.cpu().numpy() + end + + R0->>R0: dist.destroy_process_group() + deactivate R0 + + Parent->>Manager: 收集所有 result_dict["{call}_{rank}"] + Manager-->>Parent: MultiRunResult +``` + +--- + +# XOR / ULP 逐位分解 + +```mermaid +flowchart TD + A["基准值
3.14159 (float32)"] --> A1["struct.pack('<f')"] + B["目标值
3.14160 (float32)"] --> B1["struct.pack('<f')"] + A1 --> ABITS["基准位模式
0x40490FDB (uint32)"] + B1 --> BBITS["目标位模式
0x40490FDA (uint32)"] + ABITS --> XOR["xor_bits = a_bits ⊕ b_bits
0x00000001"] + BBITS --> XOR + XOR --> SIGN["符号位 (bit 31)
sign_diff = xor & 0x80000000
→ 翻转 or 正常"] + XOR --> EXP["指数位 (bits 23−30)
exp_diff = |exp_a − exp_b|"] + XOR --> MANT["尾数位 (bits 0−22)
n_flips = (xor & 0x007FFFFF).bit_count()"] + MANT --> ULP["compute_ulp(a, b)
符号-幅度转换 → ULP 距离
ULP = 1 → 相邻浮点数"] + SIGN --> INTERP["诊断结论:
符号未变 + 指数未变
+ 1 个尾数 LSB 翻转
→ chunk 尾端舍入（NCCL 典型行为）"] + EXP --> INTERP + ULP --> INTERP +``` diff --git a/src/code/issue4/docs/04-state-machine.md b/src/code/issue4/docs/04-state-machine.md new file mode 100644 index 0000000..a4aaab8 --- /dev/null +++ b/src/code/issue4/docs/04-state-machine.md @@ -0,0 +1,74 @@ +# 状态机 — diagnose.py main() 完整分支 + +```mermaid +stateDiagram-v2 + [*] --> 环境预检 + 环境预检: _check_environment() + 环境预检: HardwareCaps.detect() + + 环境预检 --> 硬件就绪: GPU ≥ 2, NCCL 可用 + 环境预检 --> 终止: GPU < 2 或无 NCCL + + 硬件就绪: HardwareCaps OK + LL128/NVSwitch 过滤 + 硬件就绪 --> 配置矩阵: ConfigMatrix.generate_with_warnings(hw) + + 配置矩阵 --> 配置非空: 有可用配置? + 配置非空 --> 终止: 无（全部被过滤） + 配置非空 --> 注入检查: 有 + + 注入检查: --inject-difference? + 注入检查 --> ULP注入验证: 是 + 注入检查 --> 数据生成: 否 + + ULP注入验证: _run_injection_test() + ULP注入验证 --> 注入通过: 比较器检出差异 + ULP注入验证 --> 终止: 比较器未检出差异 + + 注入通过: --force? + 注入通过 --> 数据生成: 是 + 注入通过 --> [*]: 否（仅验证模式退出） + + 数据生成: DataGenerator.generate_all() + 数据生成: SHA-256 每 rank 独立输入 + + 数据生成 --> NCCL扫描 + + state NCCL扫描 { + [*] --> 传统模式: --n-calls = 1 + [*] --> 三级对比模式: --n-calls > 1 + 传统模式: NcclRunner.run(trials=2) + 传统模式: → list[RunResult] + 三级对比模式: NcclRunner.run_full(n_calls, trials=2) + 三级对比模式: → list[MultiRunResult] + } + + NCCL扫描 --> 比较器 + + state 比较器 { + [*] --> 传统比较: --n-calls = 1 + [*] --> 三级比较: --n-calls > 1 + 传统比较: .compare(t0, t1) → DiffReport + 传统比较: .track_evolution() → EvolutionReport + 传统比较: .track_iter_evolution() → IterEvolutionReport + 三级比较: .compare_run_vs_run() + 三级比较: .compare_call_vs_call() + 三级比较: .compare_rank_vs_rank() + 三级比较: → ThreeLevelSummary + } + + 比较器 --> 报告器 + + 报告器: Reporter.report() → DiagnosticSummary + 报告器: 控制台表格 + JSON + 报告器: _generate_recommendations() + 报告器: _find_deterministic_config() + + 报告器 --> 差异判断 + 差异判断: 任何配置存在非确定性? + 差异判断 --> 报告修复建议: 是 + 差异判断 --> 报告全通过: 否 + 报告修复建议 --> [*] + 报告全通过 --> [*] + + 终止 --> [*] +``` diff --git a/src/code/issue4/docs/05-module-deps.md b/src/code/issue4/docs/05-module-deps.md new file mode 100644 index 0000000..ea3540f --- /dev/null +++ b/src/code/issue4/docs/05-module-deps.md @@ -0,0 +1,21 @@ +# 模块依赖图 + +```mermaid +graph TD + diagnose["diagnose.py"] -->|导入| config_matrix["config_matrix.py
HardwareCaps, ConfigMatrix, ConfigEntry"] + diagnose -->|导入| data_generator["data_generator.py
DataGenerator"] + diagnose -->|导入| runner["runner.py
NcclRunner, RunResult, MultiRunResult"] + diagnose -->|导入| comparator["comparator.py
BitwiseComparator, DiffReport,
XorDetail, ThreeLevelSummary,
compute_ulp, analyze_xor_float"] + diagnose -->|导入| reporter["reporter.py
Reporter, DiagnosticSummary"] + + runner -->|导入| config_matrix + runner -->|导入| data_generator + + comparator -->|导入| config_matrix + comparator -->|导入| runner + + reporter -->|导入| comparator + reporter -->|导入| config_matrix + + reproduce_case["reproduce_case.py
(独立运行)"] -.->|可导入| diagnose +``` diff --git a/src/code/issue4/reporter.py b/src/code/issue4/reporter.py new file mode 100644 index 0000000..858dc76 --- /dev/null +++ b/src/code/issue4/reporter.py @@ -0,0 +1,327 @@ +""" +结果报告器 — NCCL 确定性诊断。 + +以多种格式输出诊断结果: + - 控制台表格（人类可读） + - JSON 报告（机器可读，适用于 CI / 自动化分析） + - 基于差异模式的配置建议 + +诊断逻辑: + Ring → 非确定性风险最高（串行累加） + Tree → 风险较低（平衡二叉树归约） + PAT → 风险最低（并行聚合树） + float16/bf16 → 高敏感度（低尾数精度） + float32 → 中等敏感度 +""" + +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass, field +from typing import Optional + +from .comparator import DiffReport, EvolutionReport +from .config_matrix import ConfigEntry + + +@dataclass +class DiagnosticSummary: + """Top-level diagnostic summary with recommendations.""" + + total_configs: int = 0 + bitwise_match_count: int = 0 + diff_count: int = 0 + + # Key findings + failing_configs: list[DiffReport] = field(default_factory=list) + passing_configs: list[DiffReport] = field(default_factory=list) + + # Worst case + worst_diff: Optional[DiffReport] = None + + # Recommendations + recommendations: list[str] = field(default_factory=list) + deterministic_config: Optional[str] = None + + def to_dict(self) -> dict: + return { + "total_configs": self.total_configs, + "bitwise_match_count": self.bitwise_match_count, + "diff_count": self.diff_count, + "worst_diff": _diff_to_dict(self.worst_diff) if self.worst_diff else None, + "failing_configs": [_diff_to_dict(r) for r in self.failing_configs], + "passing_configs": [_diff_to_dict(r) for r in self.passing_configs], + "recommendations": self.recommendations, + "deterministic_config": self.deterministic_config, + } + + +@dataclass +class Reporter: + """格式化并输出诊断结果。""" + + output_json: bool = True + output_console: bool = True + json_path: str = "" + + # ------------------------------------------------------------------ + # Main entry: generate report from comparison results + # ------------------------------------------------------------------ + + def report( + self, + reports: list[DiffReport], + evolutions: Optional[dict[str, EvolutionReport]] = None, + json_path: Optional[str] = None, + ) -> DiagnosticSummary: + """生成并输出完整诊断报告。""" + summary = self._build_summary(reports) + + if self.output_console: + self._console_report(summary, evolutions) + + if self.output_json: + path = json_path or self.json_path or "diagnostic_report.json" + self._json_report(summary, evolutions, path) + + return summary + + # ------------------------------------------------------------------ + # Summary construction + # ------------------------------------------------------------------ + + def _build_summary(self, reports: list[DiffReport]) -> DiagnosticSummary: + s = DiagnosticSummary() + s.total_configs = len(reports) + s.passing_configs = [r for r in reports if r.bitwise_match] + s.failing_configs = [r for r in reports if not r.bitwise_match] + s.bitwise_match_count = len(s.passing_configs) + s.diff_count = len(s.failing_configs) + + # Worst case: highest diff_ratio + if s.failing_configs: + s.worst_diff = max(s.failing_configs, key=lambda r: r.diff_ratio) + + # Generate recommendations + s.recommendations = self._generate_recommendations(reports) + s.deterministic_config = self._find_deterministic_config(reports) + + return s + + # ------------------------------------------------------------------ + # ---- 控制台输出 ---- + # ------------------------------------------------------------------ + + def _console_report( + self, summary: DiagnosticSummary, + evolutions: Optional[dict[str, EvolutionReport]] = None, + ) -> None: + self._print_header("NCCL Bitwise Reproducibility Diagnostic Report") + self._print_section("Overview") + print(f" Configs tested : {summary.total_configs}") + print(f" Bitwise match : {summary.bitwise_match_count}") + print(f" Bitwise diff : {summary.diff_count}") + + if summary.failing_configs: + self._print_section("Failing Configurations (non-bitwise-deterministic)") + for r in sorted(summary.failing_configs, key=lambda x: -x.diff_ratio): + print(f" {_format_diff_line(r)}") + + if summary.passing_configs: + self._print_section("Passing Configurations (bitwise-deterministic)") + for r in summary.passing_configs: + print(f" [PASS] {_format_config_line(r.config)}") + + if summary.worst_diff: + self._print_section("Worst Non-Determinism Case") + print(f" Config: {summary.worst_diff.config}") + print(f" Diff ratio: {summary.worst_diff.diff_ratio * 100:.4f}%") + print(f" Max abs diff: {summary.worst_diff.max_abs_diff:.8e}") + if summary.worst_diff.first_diff_offset >= 0: + print(f" First diff @ offset {summary.worst_diff.first_diff_offset}: " + f"baseline={summary.worst_diff.first_diff_baseline:.8e}, " + f"target={summary.worst_diff.first_diff_target:.8e}") + + if evolutions: + self._print_section("Diff Evolution Across Data Sizes") + for label, evo in evolutions.items(): + print(f" {label}:") + for sz, r in evo.entries: + status = "IDENTICAL" if r.bitwise_match else f"{r.diff_ratio * 100:6.2f}% diff" + print(f" {_format_size_static(sz):>6s} → {status}") + + self._print_section("Recommendations for Bitwise Determinism") + for i, rec in enumerate(summary.recommendations, 1): + print(f" {i}. {rec}") + + if summary.deterministic_config: + self._print_section("Deterministic Configuration") + print(f" {summary.deterministic_config}") + + self._print_footer() + + # ------------------------------------------------------------------ + # ---- JSON 输出 ---- + # ------------------------------------------------------------------ + + def _json_report( + self, summary: DiagnosticSummary, + evolutions: Optional[dict[str, EvolutionReport]], + path: str, + ) -> None: + data = summary.to_dict() + + if evolutions: + data["evolution"] = {} + for label, evo in evolutions.items(): + data["evolution"][label] = evo.diff_ratio_curve() + + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + print(f"\n JSON report written to: {path}") + + # ------------------------------------------------------------------ + # ---- 诊断建议引擎 ---- + # ------------------------------------------------------------------ + + def _generate_recommendations(self, reports: list[DiffReport]) -> list[str]: + recs: list[str] = [] + + # Analyze failing configs by algo + algo_fails: dict[str, list[DiffReport]] = {} + for r in reports: + if not r.bitwise_match: + algo_fails.setdefault(r.config.algo, []).append(r) + + if "Ring" in algo_fails and len(algo_fails["Ring"]) >= len([r for r in reports if r.config.algo == "Ring"]): + recs.append( + "NCCL_ALGO=Ring consistently produces non-bitwise-deterministic results. " + "Use NCCL_ALGO=Tree or NCCL_ALGO=PAT instead. " + "Ring's serial accumulation order is inherently sensitive to chunk partitioning." + ) + + if "Tree" in algo_fails: + recs.append( + "Even Tree algorithm showed non-determinism. Try increasing precision: " + "use float32 instead of float16/bf16, or enable float64 reduction if supported." + ) + + # Analyze by dtype + dtype_fails: dict[str, list[DiffReport]] = {} + for r in reports: + if not r.bitwise_match: + dtype_fails.setdefault(r.config.dtype, []).append(r) + + if "float16" in dtype_fails or "bfloat16" in dtype_fails: + recs.append( + "float16 / bfloat16 amplify non-determinism due to low mantissa precision " + "(10-bit / 7-bit). Use float32 for reduction buffers when bitwise reproducibility " + "is required." + ) + + # Passing configs → suggest them + passing = [r for r in reports if r.bitwise_match] + if passing: + best = min(passing, key=lambda r: r.config.size_bytes) # smallest size that works + recs.append( + f"At least one configuration achieved bitwise determinism: " + f"NCCL_ALGO={best.config.algo}, NCCL_PROTO={best.config.proto}, " + f"dtype={best.config.dtype}." + ) + + # General guidance + recs.append( + "For production determinism: set CUBLAS_WORKSPACE_CONFIG=:4096:8, " + "torch.use_deterministic_algorithms(True), torch.backends.cudnn.benchmark=False, " + "and torch.backends.cudnn.deterministic=True." + ) + + recs.append( + "To guarantee bitwise reproducibility: use Reduce+Broadcast (fixed root) " + "instead of AllReduce. NCCL_ALGO=PAT (Parallel Aggregated Trees, NCCL 2.23+) " + "also provides better determinism than Ring." + ) + + return recs + + def _find_deterministic_config(self, reports: list[DiffReport]) -> Optional[str]: + """Find the most robust deterministic config.""" + passing = [r for r in reports if r.bitwise_match] + if not passing: + return None + + # Tree/PAT 优先，float32 优于 float16 + scored = [] + for r in passing: + score = 0 + if r.config.algo in ("PAT", "Tree"): + score += 2 + if r.config.dtype == "float32": + score += 1 + scored.append((score, r)) + + best = max(scored, key=lambda x: x[0]) + return ( + f"NCCL_ALGO={best[1].config.algo} " + f"NCCL_PROTO={best[1].config.proto} " + f"dtype={best[1].config.dtype}" + ) + + # ------------------------------------------------------------------ + # ---- 格式化工具 ---- + # ------------------------------------------------------------------ + + @staticmethod + def _print_header(title: str) -> None: + print(f"\n{'=' * 68}") + print(f" {title}") + print(f"{'=' * 68}") + + @staticmethod + def _print_section(title: str) -> None: + print(f"\n--- {title} ---") + + @staticmethod + def _print_footer() -> None: + print(f"\n{'=' * 68}\n") + + +def _format_config_line(cfg: ConfigEntry) -> str: + return f"algo={cfg.algo:>14s} proto={cfg.proto:>7s} dtype={cfg.dtype:>8s} size={cfg.size_label:>6s}" + + +def _format_diff_line(r: DiffReport) -> str: + cfg = r.config + return ( + f"[DIFF] {_format_config_line(cfg)} " + f"diff={r.diff_count}/{r.total_elements} ({r.diff_ratio * 100:.4f}%) " + f"max_abs={r.max_abs_diff:.6e}" + ) + + +def _diff_to_dict(r: DiffReport) -> dict: + return { + "algo": r.config.algo, + "proto": r.config.proto, + "dtype": r.config.dtype, + "size_bytes": r.config.size_bytes, + "size_label": r.config.size_label, + "bitwise_match": r.bitwise_match, + "total_elements": r.total_elements, + "diff_count": r.diff_count, + "diff_ratio": f"{r.diff_ratio * 100:.6f}%", + "max_abs_diff": f"{r.max_abs_diff:.8e}" if r.max_abs_diff > 0 else "0", + "mean_abs_diff": f"{r.mean_abs_diff:.8e}" if r.mean_abs_diff > 0 else "0", + "first_diff_offset": r.first_diff_offset, + "first_diff_baseline": f"{r.first_diff_baseline:.8e}", + "first_diff_target": f"{r.first_diff_target:.8e}", + } + + +def _format_size_static(nbytes: int) -> str: + if nbytes < 1024: + return f"{nbytes}B" + if nbytes < 1024 * 1024: + return f"{nbytes // 1024}K" + return f"{nbytes // (1024 * 1024)}M" diff --git a/src/code/issue4/reproduce_case.py b/src/code/issue4/reproduce_case.py new file mode 100644 index 0000000..d015cf2 --- /dev/null +++ b/src/code/issue4/reproduce_case.py @@ -0,0 +1,445 @@ +#!/usr/bin/env python3 +""" +独立非确定性复现案例 + +复现已验证的 NCCL 非确定性场景: + Ring AllReduce + float16 数据 + N GPU，跨运行逐位比对。 + +根因: Ring 算法的串行累加顺序依赖于 chunk 划分，而 chunk 划分 +随缓冲区大小、协议和 NVLink 通道数微调。浮点非结合性导致 +不同累加顺序产生不同的舍入误差 → 位级结果不同。 + +测试用例: + Case 1: 同一配置重复运行 → 应为确定性的（NCCL 保证） + Case 2: Ring vs Tree 同一数据 → 预期有差异（不同归约顺序） + Case 3: float16 vs float32 → 展示精度对确定性的影响 + +用法: + python reproduce_case.py [--nranks 8] [--size 128M] [--dtype float16] + python reproduce_case.py --algo Ring --trials 5 # 迭代演化 + +硬件要求: >= 2 个 NVIDIA GPU，支持 NCCL。 + +参考: + - NCCL#1055: Ring vs Tree 精度对比 (A100/A800) + - PyTorch#138811: H20 allreduce 非确定性 + - NCCL#1975: Ring 算法精度讨论 + - NCCL#157: chunk 划分与确定性 +""" + +from __future__ import annotations + +import argparse +import hashlib +import multiprocessing as mp +import os +import socket +import struct +import sys +import time +from dataclasses import dataclass +from typing import Optional + +import numpy as np + + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + +DTYPE_NP_MAP = {"float32": np.float32, "float16": np.float16, + "bfloat16": np.float16} +DTYPE_BYTES = {"float32": 4, "float16": 2, "bfloat16": 2} + + +@dataclass +class ReproduceResult: + trial: int + algo: str + proto: str + dtype: str + size_bytes: int + output: np.ndarray + elapsed_ms: float + + +# --------------------------------------------------------------------------- +# Deterministic data generation +# --------------------------------------------------------------------------- + +def _deterministic_rand(count: int, seed: int) -> np.ndarray: + """Generate count floats in [-1, 1] via SHA-256 PRNG.""" + data = np.empty(count, dtype=np.float32) + generated = 0 + block = 0 + while generated < count: + h = hashlib.sha256() + h.update(struct.pack("= count: + break + val = struct.unpack(" list[np.ndarray]: + type_bytes = DTYPE_BYTES[dtype] + count = size_bytes // type_bytes + inputs = [] + for rank in range(nranks): + arr = _deterministic_rand(count, seed=42 + rank) + inputs.append(arr.astype(DTYPE_NP_MAP[dtype])) + return inputs + + +# --------------------------------------------------------------------------- +# ---- PyTorch NCCL worker ---- +# --------------------------------------------------------------------------- + +def _torch_dtype(dtype: str): + import torch + return {"float32": torch.float32, "float16": torch.float16, + "bfloat16": torch.bfloat16}[dtype] + + +def _allreduce_worker( + rank: int, world_size: int, port: int, + algo: str, proto: str, dtype: str, input_data: np.ndarray, + collective: str, result_dict: dict, +) -> None: + import torch + import torch.distributed as dist + + os.environ["NCCL_ALGO"] = algo + os.environ["NCCL_PROTO"] = proto + + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, world_size=world_size, + device_id=device, + ) + + dt = _torch_dtype(dtype) + tensor = torch.from_numpy(input_data).to(device=device, dtype=dt) + + # Warmup + warm = tensor.clone() + if collective == "allreduce": + dist.all_reduce(warm) + torch.cuda.synchronize() + + # Timed run + tensor = torch.from_numpy(input_data).to(device=device, dtype=dt) + torch.cuda.synchronize() + t0 = time.perf_counter() + + if collective == "allreduce": + dist.all_reduce(tensor) + elif collective == "reducescatter": + count = input_data.size // world_size + dist.reduce_scatter_tensor( + tensor.view(-1), + torch.zeros(count, dtype=dt, device=device), + ) + + torch.cuda.synchronize() + t1 = time.perf_counter() + + if rank == 0: + result_dict[0] = tensor.cpu().float().numpy() + result_dict["elapsed"] = (t1 - t0) * 1000.0 + + dist.destroy_process_group() + + +def _run_trial( + algo: str, proto: str, dtype: str, size_bytes: int, + nranks: int, collective: str, trial_id: int, +) -> ReproduceResult: + port = _find_free_port() + inputs = generate_inputs(size_bytes, nranks, dtype) + count = size_bytes // DTYPE_BYTES[dtype] + + manager = mp.Manager() + result_dict = manager.dict() + + processes = [] + for rank in range(nranks): + p = mp.Process( + target=_allreduce_worker, + args=(rank, nranks, port, algo, proto, dtype, + inputs[rank], collective, result_dict), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + output = np.array(result_dict.get(0, np.zeros(count, dtype=np.float32))) + elapsed = result_dict.get("elapsed", 0.0) + + return ReproduceResult( + trial=trial_id, algo=algo, proto=proto, dtype=dtype, + size_bytes=size_bytes, output=output, elapsed_ms=elapsed, + ) + + +# --------------------------------------------------------------------------- +# Comparison (with distribution analysis) +# --------------------------------------------------------------------------- + +def compare_bitwise(r1: ReproduceResult, r2: ReproduceResult) -> dict: + b = r1.output + t = r2.output + + if b.shape != t.shape: + return {"error": f"Shape mismatch: {b.shape} vs {t.shape}"} + + diff_mask = (b != t) + diff_idx = np.where(diff_mask)[0] + diff_count = len(diff_idx) + total = b.size + + if diff_count == 0: + return {"bitwise_match": True, "total_elements": int(total)} + + abs_diff = np.abs(b[diff_mask].astype(np.float64) - + t[diff_mask].astype(np.float64)) + + # Spatial distribution analysis + n = total + third = n // 3 + fc = int(np.sum(diff_idx < third)) + mc = int(np.sum((diff_idx >= third) & (diff_idx < 2 * third))) + bc = int(np.sum(diff_idx >= 2 * third)) + + fr = fc / third if third > 0 else 0.0 + mr = mc / third if third > 0 else 0.0 + br = bc / (n - 2 * third) if (n - 2 * third) > 0 else 0.0 + + max_seg = max(fr, mr, br) + if max_seg == 0: + conc = "uniform" + elif max_seg >= 2 * min(f for f in (fr, mr, br) if f > 0): + conc = {0: "front", 1: "mid", 2: "back"}[np.argmax([fr, mr, br])] + elif fr + br > 2 * mr: + conc = "edges" + else: + conc = "uniform" + + return { + "bitwise_match": False, + "total_elements": int(total), + "diff_count": diff_count, + "diff_ratio_pct": diff_count / total * 100, + "first_diff_offset": int(diff_idx[0]) if diff_count > 0 else -1, + "first_diff_baseline": float(b.flat[diff_idx[0]]), + "first_diff_target": float(t.flat[diff_idx[0]]), + "max_abs_diff": float(np.max(abs_diff)), + "mean_abs_diff": float(np.mean(abs_diff)), + "std_abs_diff": float(np.std(abs_diff)), + "distribution": { + "front_ratio_pct": f"{fr * 100:.4f}%", + "mid_ratio_pct": f"{mr * 100:.4f}%", + "back_ratio_pct": f"{br * 100:.4f}%", + "concentration": conc, + }, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(argv: Optional[list[str]] = None) -> int: + import multiprocessing as _mp + _mp.set_start_method("spawn", force=True) + + args = _parse_args(argv) + + print("=" * 70) + print(" NCCL Non-Determinism Reproduction Case") + print("=" * 70) + print(f" Algorithm : {args.algo}") + print(f" Protocol : {args.proto}") + print(f" Data type : {args.dtype}") + print(f" Data size : {args.size}") + print(f" N ranks : {args.nranks}") + print(f" Collective: {args.collective}") + print(f" Trials : {args.trials}") + print() + + available = _count_gpus() + if available < args.nranks: + print(f" WARNING: Only {available} GPUs available, using {available}") + args.nranks = available + if args.nranks < 2: + print(" ERROR: Need at least 2 GPUs. Aborting.") + return 1 + + size_bytes = _parse_size(args.size) + + # 案例 1：同配置多次试验 —— 迭代级演化 + print("--- Case 1: Same config, run-to-run determinism ---") + print(" (NCCL guarantee: same input + same topology = bitwise identical)") + print() + trials: list[ReproduceResult] = [] + for i in range(args.trials): + print(f" Trial {i + 1}/{args.trials}...", end=" ", flush=True) + r = _run_trial( + algo=args.algo, proto=args.proto, dtype=args.dtype, + size_bytes=size_bytes, nranks=args.nranks, + collective=args.collective, trial_id=i + 1, + ) + trials.append(r) + print(f"done ({r.elapsed_ms:.2f}ms, checksum={r.output.sum():.6f})") + + print() + if len(trials) >= 2: + report = compare_bitwise(trials[0], trials[1]) + _print_report(report) + + if len(trials) > 2: + print(" Iteration evolution (trial_0 baseline):") + growing = True + prev_ratio = 0.0 + for i in range(2, len(trials)): + r = compare_bitwise(trials[0], trials[i]) + ratio = r.get("diff_ratio_pct", 0) + growing = growing and (ratio >= prev_ratio) + prev_ratio = ratio + status = "IDENTICAL" if r["bitwise_match"] else f"{ratio:.4f}% diff" + print(f" trial_0 vs trial_{i}: {status}") + if growing and not report["bitwise_match"]: + print(" → Diffs MONOTONICALLY GROWING — accumulation pattern!") + print() + + # 案例 2：Ring vs Tree 跨算法比对 + if args.algo != "Ring": + print("--- Case 2: Ring vs Tree cross-algorithm comparison ---") + print(" Running Ring trial...", end=" ", flush=True) + r_ring = _run_trial("Ring", args.proto, args.dtype, size_bytes, + args.nranks, args.collective, 0) + print(f"done ({r_ring.elapsed_ms:.2f}ms)") + + print(" Running Tree trial...", end=" ", flush=True) + r_tree = _run_trial("Tree", args.proto, args.dtype, size_bytes, + args.nranks, args.collective, 0) + print(f"done ({r_tree.elapsed_ms:.2f}ms)") + + print("\n Ring vs Tree comparison:") + _print_report(compare_bitwise(r_ring, r_tree)) + print() + + # 案例 3：float16 vs float32 精度影响 + if args.dtype == "float16": + print("--- Case 3: float16 vs float32 precision impact ---") + print(" Running float32 trial (same element count)...", end=" ", flush=True) + r_f32 = _run_trial(args.algo, args.proto, "float32", + size_bytes * 2, args.nranks, args.collective, 0) + print(f"done ({r_f32.elapsed_ms:.2f}ms)") + print(" Note: float32 has 23-bit mantissa vs float16's 10-bit — " + "better precision, less non-determinism.") + print() + + print("=" * 70) + print(" Reproduction complete.") + print() + print(" Root cause: NCCL Ring algorithm serial accumulation order depends") + print(" on chunk partitioning (buffer size, NVLink channels).") + print(" Floating-point non-associativity → different rounding errors.") + print() + print(" Recommendations for bitwise determinism:") + print(" 1. Use NCCL_ALGO=Tree or NCCL_ALGO=PAT instead of Ring") + print(" 2. Use float32 for reduction buffers (not float16/bf16)") + print(" 3. Fix all NCCL env vars: NCCL_ALGO, NCCL_PROTO, NCCL_NCHANNELS") + print(" 4. Set torch.use_deterministic_algorithms(True)") + print(" 5. Consider Reduce+Broadcast with fixed root") + print("=" * 70) + + return 0 + + +# --------------------------------------------------------------------------- +# ---- 工具函数 ---- +# --------------------------------------------------------------------------- + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _count_gpus() -> int: + try: + import torch + return torch.cuda.device_count() + except Exception: + return 0 + + +def _parse_size(s: str) -> int: + s = s.strip().upper() + m = {"B": 1, "K": 1024, "M": 1024**2, "G": 1024**3} + for suffix, mult in m.items(): + if s.endswith(suffix): + return int(s[:-1]) * mult + return int(s) + + +def _parse_args(argv: Optional[list[str]]) -> argparse.Namespace: + p = argparse.ArgumentParser(description="NCCL Non-Determinism Reproduction Case") + p.add_argument("--algo", default="Ring", + choices=["Ring", "Tree", "PAT", "CollnetDirect", + "CollnetChain", "NVLS", "NVLSTree"]) + p.add_argument("--proto", default="Simple", choices=["LL", "Simple", "LL128"]) + p.add_argument("--dtype", default="float16", + choices=["float32", "float16", "bfloat16"]) + p.add_argument("--size", default="128M") + p.add_argument("--nranks", type=int, default=8) + p.add_argument("--trials", type=int, default=2, + help="Number of repeated trials (>=3 enables iteration evolution)") + p.add_argument("--collective", default="allreduce", + choices=["allreduce", "reducescatter"]) + return p.parse_args(argv) + + +def _print_report(report: dict) -> None: + if report.get("error"): + print(f" ERROR: {report['error']}") + return + if report["bitwise_match"]: + print(f" RESULT: BITWISE IDENTICAL ({report['total_elements']} elements)") + return + + print(f" RESULT: BITWISE DIFFERENCE DETECTED!") + print(f" Total elements : {report['total_elements']}") + print(f" Diff count : {report['diff_count']} " + f"({report['diff_ratio_pct']:.4f}%)") + print(f" First diff @ {report['first_diff_offset']}: " + f"baseline={report['first_diff_baseline']:.8e}, " + f"target={report['first_diff_target']:.8e}") + print(f" Max abs diff : {report['max_abs_diff']:.8e}") + print(f" Mean abs diff : {report['mean_abs_diff']:.8e}") + print(f" Std abs diff : {report['std_abs_diff']:.8e}") + + dist = report.get("distribution", {}) + if dist: + print(f" Distribution : " + f"front={dist['front_ratio_pct']}, " + f"mid={dist['mid_ratio_pct']}, " + f"back={dist['back_ratio_pct']} " + f"({dist['concentration']})") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/code/issue4/requirements.txt b/src/code/issue4/requirements.txt new file mode 100644 index 0000000..b1465ec --- /dev/null +++ b/src/code/issue4/requirements.txt @@ -0,0 +1,11 @@ +# NCCL Bitwise Reproducibility Diagnostic Tool +# Minimum requirements for the PyTorch backend (default) + +numpy>=1.21.0 +torch>=1.12.0 + +# Note: PyTorch must be built with NCCL backend support. +# Verify with: python -c "import torch.distributed; print(torch.distributed.is_nccl_available())" + +# For nccl-tests backend (optional, requires buffer-dump patch): +# git clone https://github.com/NVIDIA/nccl-tests && cd nccl-tests && make MPI=1 diff --git a/src/code/issue4/runner.py b/src/code/issue4/runner.py new file mode 100644 index 0000000..961f07c --- /dev/null +++ b/src/code/issue4/runner.py @@ -0,0 +1,504 @@ +""" +NCCL 集合通信执行器。 + +提供两种后端: + 1. PyTorch NCCL — 使用 torch.distributed（推荐，无需额外依赖） + 2. nccl-tests — 包装 all_reduce_perf / reduce_scatter_perf 二进制（原始 NCCL） + +执行器在集合操作完成后捕获输出张量，用于后续逐位比对。 +每次运行通过独立进程（每 rank 一个）来避免 NCCL 环境变量缓存问题。 + +设计说明: + nccl-tests 不导出原始缓冲区数据，因此 nccl-tests 后端需打补丁或在 + CUDA 侧保存缓冲区。在此之前，PyTorch 后端是主要路径。 +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import socket +import subprocess +import sys +import time +from dataclasses import dataclass, field +from typing import Optional + +import numpy as np + +from .config_matrix import ConfigEntry +from .data_generator import DataGenerator, DTYPE_BYTES + + +@dataclass +class RunResult: + """单次 NCCL 集合通信调用的输出。""" + + config: ConfigEntry + output: np.ndarray # result tensor (shape depends on collective) + elapsed_ms: float = 0.0 # GPU-side wall time + rank: int = 0 # which rank's output this is + + def checksum(self) -> float: + """浮点校验和，用于快速漂移检测。""" + return float(np.sum(self.output.astype(np.float64))) + + def save(self, path: str) -> None: + """Save output as raw binary (float32) for cross-run comparison.""" + self.output.astype(np.float32).tofile(path) + + @classmethod + def load(cls, path: str, config: ConfigEntry) -> "RunResult": + """从原始二进制文件加载输出。""" + count = config.size_bytes // DTYPE_BYTES.get(config.dtype, 4) + if config.collective == "reducescatter": + count = count // config.nranks + data = np.fromfile(path, dtype=np.float32, count=count) + return cls(config=config, output=data) + + +@dataclass +class MultiRunResult: + """Output of N calls × R ranks within one process-group run. + + call_outputs[call_idx][rank] = np.ndarray (rank's output for that call). + """ + + config: ConfigEntry + call_outputs: list[dict[int, np.ndarray]] # [call_idx] → {rank: output} + elapsed_ms: list[float] # per-call latency + + +@dataclass +class NcclRunner: + """执行 NCCL 集合通信操作并捕获输出张量。 + + Parameters + ---------- + nranks : int + Number of GPU ranks to use. + backend : str + 'pytorch' (default) — torch.distributed + 'nccl-tests' — raw NCCL via nccl-tests binaries (needs buffer-dump patch) + nccl_tests_dir : str, optional + Path to nccl-tests build directory. Required if backend='nccl-tests'. + scratch_dir : str, optional + Directory for temporary output files. Uses system temp if None. + """ + + nranks: int = 8 + backend: str = "pytorch" + nccl_tests_dir: str = "" + scratch_dir: str = "" + + _data_gen: Optional[DataGenerator] = field(default=None, init=False) + + # ------------------------------------------------------------------ + # ---- 公共 API ---- + # ------------------------------------------------------------------ + + def run(self, config: ConfigEntry, trials: int = 2) -> list[RunResult]: + """Run the collective `trials` times under the same config. + + Each trial spawns a fresh process group to guarantee clean NCCL state. + Returns list of RunResult, one per trial (rank-0 only, backward-compat). + """ + results: list[RunResult] = [] + for _ in range(trials): + result = self._run_once(config) + results.append(result) + return results + + def run_full(self, config: ConfigEntry, n_calls: int, + trials: int = 2) -> list[MultiRunResult]: + """Run `n_calls` collective calls inside each of `trials` process groups. + + Returns one MultiRunResult per trial. Each MultiRunResult contains + all call outputs for all ranks — suitable for three-level comparison: + - run-vs-run: trial0[call_k][rank_r] vs trial1[call_k][rank_r] + - call-vs-call: trial_t[call_i][rank_r] vs trial_t[call_j][rank_r] + - rank-vs-rank: trial_t[call_k][rank_a] vs trial_t[call_k][rank_b] + """ + results: list[MultiRunResult] = [] + for _ in range(trials): + mrr = self._run_full_once(config, n_calls) + results.append(mrr) + return results + + def sweep( + self, configs: list[ConfigEntry], trials: int = 2 + ) -> list[list[RunResult]]: + """运行所有配置。 Returns [[trial_0, trial_1, ...], ...] per config.""" + all_results: list[list[RunResult]] = [] + for cfg in configs: + results = self.run(cfg, trials=trials) + all_results.append(results) + return all_results + + # ------------------------------------------------------------------ + # ---- 完整模式：进程组内多调用 / 全部 rank ---- + # ------------------------------------------------------------------ + + def _run_full_once(self, config: ConfigEntry, n_calls: int) -> MultiRunResult: + """Execute `n_calls` collective invocations within one process group.""" + self._data_gen = DataGenerator( + dtype=config.dtype, base_seed=42, collective=config.collective, + ) + port = _find_free_port() + count = self._data_gen.elem_count(config.size_bytes, config.nranks) + inputs = self._data_gen.generate_all(config.size_bytes, config.nranks) + + with mp.Manager() as manager: + result_dict = manager.dict() + processes = [] + for rank in range(config.nranks): + p = mp.Process( + target=self._pytorch_worker_full, + args=(rank, config.nranks, port, config, count, n_calls, + inputs[rank], result_dict), + ) + p.start() + processes.append(p) + for p in processes: + p.join() + + # Collect: result_dict[f"{call_idx}_{rank}"] = np.ndarray + call_outputs: list[dict[int, np.ndarray]] = [] + elapsed_ms: list[float] = [] + for call_idx in range(n_calls): + per_rank: dict[int, np.ndarray] = {} + for rank in range(config.nranks): + key = f"{call_idx}_{rank}" + if key not in result_dict: + raise RuntimeError( + f"Worker did not produce output for call={call_idx} " + f"rank={rank}. Check GPU memory / NCCL." + ) + per_rank[rank] = np.array(result_dict[key]) + call_outputs.append(per_rank) + elapsed_ms.append(float(result_dict.get(f"_elapsed_{call_idx}", 0.0))) + + return MultiRunResult( + config=config, call_outputs=call_outputs, elapsed_ms=elapsed_ms, + ) + + @staticmethod + def _pytorch_worker_full( + rank: int, world_size: int, port: int, config: ConfigEntry, + count: int, n_calls: int, input_data: np.ndarray, result_dict: dict, + ) -> None: + """Worker: init NCCL once, run collective `n_calls` times.""" + import torch + import torch.distributed as dist + + os.environ.update(config.env_dict()) + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, world_size=world_size, + device_id=device, + ) + + dt = _torch_dtype(config.dtype) + + # Warmup + warm = torch.from_numpy(input_data).to(device=device, dtype=dt) + if config.collective == "allreduce": + dist.all_reduce(warm) + elif config.collective == "reducescatter": + wc = input_data.size // world_size + dist.reduce_scatter_tensor( + torch.zeros(wc, dtype=dt, device=device), warm, + ) + torch.cuda.synchronize() + + # Timed calls + for call_idx in range(n_calls): + tensor = torch.from_numpy(input_data).to(device=device, dtype=dt) + torch.cuda.synchronize() + t0 = time.perf_counter() + + if config.collective == "allreduce": + dist.all_reduce(tensor) + out = tensor + elif config.collective == "reducescatter": + recv_count = input_data.size // world_size + out = torch.zeros(recv_count, dtype=dt, device=device) + dist.reduce_scatter_tensor(out, tensor) + else: + raise ValueError(f"Unknown collective: {config.collective}") + + torch.cuda.synchronize() + t1 = time.perf_counter() + + result_dict[f"{call_idx}_{rank}"] = out.cpu().float().numpy() + result_dict[f"_elapsed_{call_idx}"] = (t1 - t0) * 1000.0 + + dist.destroy_process_group() + + # ------------------------------------------------------------------ + # ---- 内部：单次运行编排 ---- (backward-compat) + # ------------------------------------------------------------------ + + def _run_once(self, config: ConfigEntry) -> RunResult: + """执行一次集合通信调用并返回输出。""" + self._data_gen = DataGenerator( + dtype=config.dtype, + base_seed=42, + collective=config.collective, + ) + + if self.backend == "pytorch": + return self._run_pytorch(config) + elif self.backend == "nccl-tests": + return self._run_nccl_tests(config) + else: + raise ValueError(f"Unknown backend: {self.backend}") + + # ------------------------------------------------------------------ + # ---- PyTorch NCCL 后端 ---- + # ------------------------------------------------------------------ + + def _run_pytorch(self, config: ConfigEntry) -> RunResult: + """通过独立进程中的 torch.distributed 运行。""" + port = _find_free_port() + + # 提前生成各 rank 输入（CPU 侧） + count = self._data_gen.elem_count(config.size_bytes, config.nranks) + inputs = self._data_gen.generate_all(config.size_bytes, config.nranks) + + with mp.Manager() as manager: + result_dict = manager.dict() + + processes = [] + for rank in range(config.nranks): + p = mp.Process( + target=self._pytorch_worker, + args=(rank, config.nranks, port, config, count, + inputs[rank], result_dict), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + # 提取 rank-0 输出 (all ranks produce identical result in a + # single AllReduce invocation — NCCL guarantee) + if 0 not in result_dict: + raise RuntimeError( + f"Rank-0 worker did not produce output. " + f"Check GPU memory and NCCL availability for config: {config}" + ) + output = np.array(result_dict[0]) + elapsed = result_dict.get("elapsed", 0.0) + + return RunResult(config=config, output=output, elapsed_ms=elapsed, rank=0) + + @staticmethod + def _pytorch_worker( + rank: int, world_size: int, port: int, config: ConfigEntry, + count: int, input_data: np.ndarray, result_dict: dict, + ) -> None: + """逐 rank 的工作函数（在独立进程中运行）。 + + Uses Gloo backend for barrier/control-plane operations (so NCCL_ALGO + doesn't affect rank synchronization) and NCCL exclusively for the + collective operation being tested. + """ + import torch + import torch.distributed as dist + + os.environ.update(config.env_dict()) + + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + # --- Primary NCCL group (for collectives only) --- + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + device_id=device, + ) + + # --- Gloo subgroup for barrier / control-plane --- + # This isolates NCCL_ALGO changes from synchronization ops. + # NCCL's barrier is a no-op reduction but Gloo is TCP-based + # and completely unaffected by NCCL tuning variables. + all_ranks = list(range(world_size)) + gloo_group: dist.ProcessGroup | None = None + try: + gloo_group = dist.new_group( + ranks=all_ranks, + backend="gloo", + ) + except Exception: + # Gloo may not be available in some builds; fall back to NCCL barrier + gloo_group = None + + def _barrier() -> None: + """Use Gloo if available, else fall back to NCCL barrier.""" + if gloo_group is not None: + dist.barrier(group=gloo_group) + else: + dist.barrier() + + # --- Warmup (on NCCL) --- + dt = _torch_dtype(config.dtype) + tensor = torch.from_numpy(input_data).to(device=device, dtype=dt) + + warm = tensor.clone() + if config.collective == "allreduce": + dist.all_reduce(warm) + elif config.collective == "reducescatter": + # ReduceScatter: sendbuf has nranks*count elements, recvbuf has count + recv_count = input_data.size // world_size + dist.reduce_scatter_tensor( + torch.zeros(recv_count, dtype=dt, device=device), + warm, + ) + torch.cuda.synchronize() + _barrier() + + # --- Timed run (NCCL collective only) --- + tensor = torch.from_numpy(input_data).to(device=device, dtype=dt) + torch.cuda.synchronize() + t0 = time.perf_counter() + + if config.collective == "allreduce": + dist.all_reduce(tensor) + elif config.collective == "reducescatter": + dist.reduce_scatter_tensor( + torch.zeros(count, dtype=dt, device=device), + tensor, + ) + else: + raise ValueError(f"Unknown collective: {config.collective}") + + torch.cuda.synchronize() + t1 = time.perf_counter() + + # --- Store result --- + if rank == 0: + result_dict[0] = tensor.cpu().float().numpy() + result_dict["elapsed"] = (t1 - t0) * 1000.0 + + _barrier() + if gloo_group is not None: + dist.destroy_process_group(gloo_group) + dist.destroy_process_group() + + # ------------------------------------------------------------------ + # ---- nccl-tests 后端（需 buffer-dump 补丁） ---- + # ------------------------------------------------------------------ + + def _run_nccl_tests(self, config: ConfigEntry) -> RunResult: + """通过 nccl-tests 的 all_reduce_perf 二进制运行。 + + WARNING: nccl-tests does NOT export raw tensor data from GPU buffers. + This backend requires one of: + a) A patched nccl-tests binary that dumps sendbuff/recvbuff to disk + b) A CUDA-side LD_PRELOAD interposer that captures buffer contents + c) Using nvprof/nsys to capture memory states + + Until one of these is in place, use the 'pytorch' backend instead. + """ + if not self.nccl_tests_dir: + raise RuntimeError( + "nccl_tests_dir required for nccl-tests backend.\n" + "Use the 'pytorch' backend (default) for bitwise comparison.\n" + "Example: NcclRunner(nranks=8, backend='pytorch')" + ) + + binary_map = { + "allreduce": "all_reduce_perf", + "reducescatter": "reduce_scatter_perf", + } + binary = os.path.join(self.nccl_tests_dir, binary_map[config.collective]) + + if not os.path.isfile(binary): + raise FileNotFoundError( + f"nccl-tests binary '{binary}' not found.\n" + f"Build nccl-tests first:\n" + f" git clone https://github.com/NVIDIA/nccl-tests\n" + f" cd nccl-tests && make MPI=1\n" + f" Then pass --nccl-tests-dir ./build to the diagnostic tool." + ) + + dtype_map = {"float32": "float", "float16": "half", "bfloat16": "bfloat16"} + cmd = [ + binary, + "-b", str(config.size_bytes), + "-e", str(config.size_bytes), + "-g", str(config.nranks), + "-n", "1", + "-w", "1", + "-d", dtype_map.get(config.dtype, "float"), + "-c", "0", + "--blocking", "1", + ] + + env = os.environ.copy() + env.update(config.env_dict()) + + try: + result = subprocess.run( + cmd, env=env, capture_output=True, text=True, timeout=120, + cwd=self.nccl_tests_dir, + ) + except subprocess.TimeoutExpired: + raise RuntimeError(f"nccl-tests timed out for config: {config}") + + if result.returncode != 0: + raise RuntimeError( + f"nccl-tests failed (rc={result.returncode}):\n{result.stderr}" + ) + + raise NotImplementedError( + "nccl-tests backend requires a buffer-dump patch to capture raw " + "output data for bitwise comparison. The nccl-tests binary does " + "not expose GPU buffer contents.\n\n" + "Workaround: use the 'pytorch' backend instead:\n" + " NcclRunner(nranks=8, backend='pytorch')\n\n" + "Future: implement a CUDA LD_PRELOAD interposer or contribute a " + "--dump-buffers flag to nccl-tests." + ) + + +# --------------------------------------------------------------------------- +# ---- 工具函数 ---- +# --------------------------------------------------------------------------- + +def _find_free_port() -> int: + """在本地查找空闲 TCP 端口，同时确保 port+1 也空闲（NCCL 内部占用）。 + + 快速连续创建进程组时，上次的 port+1 可能还没释放——这里做显式检查。 + """ + for _ in range(20): # 最多试 20 次 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + # NCCL 内部也占 port+1 — 检查相邻端口是否空闲 + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2: + s2.bind(("127.0.0.1", port + 1)) + s2.close() + return port + except OSError: + continue + raise RuntimeError("无法找到连续两个空闲端口 (NCCL 需要 port 和 port+1)") + + +def _torch_dtype(dtype: str): + """将 dtype 字符串转换为 torch.dtype。""" + import torch + return { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[dtype]