From 705894c744a199b537fe45b6e8d2e4fa303e1ee9 Mon Sep 17 00:00:00 2001 From: Ryo Wakizaka Date: Thu, 9 Jul 2026 11:22:14 +0900 Subject: [PATCH] Fix NCCL warm-up allreduce count mismatch --- include/sbd/chemistry/tpb/sbdiag.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/include/sbd/chemistry/tpb/sbdiag.h b/include/sbd/chemistry/tpb/sbdiag.h index e34d6c3..7f65723 100644 --- a/include/sbd/chemistry/tpb/sbdiag.h +++ b/include/sbd/chemistry/tpb/sbdiag.h @@ -314,17 +314,17 @@ namespace sbd { } if (mpi_size_b > 1) { init_nccl_comm(&b_nccl_comm, b_comm); - thrust::device_vector A(W.size(), 0.0); + thrust::device_vector A(1, 0.0); nccl_allreduce(A, ncclSum, b_nccl_comm); } if (false && mpi_size_t > 1) { init_nccl_comm(&t_nccl_comm, t_comm); - thrust::device_vector A(W.size(), 0.0); + thrust::device_vector A(1, 0.0); nccl_allreduce(A, ncclSum, t_nccl_comm); } if (mpi_size_a > 1) { init_nccl_comm(&a_nccl_comm, a_comm); - thrust::device_vector A(W.size(), 0.0); + thrust::device_vector A(1, 0.0); nccl_allreduce(A, ncclSum, a_nccl_comm); } printf("[%s,%d] NCCL communicators have been created.\n",