Add the world size info in NCCL metadata

Summary: This diff adds the world size info in NCCL metadata, as we need the information to calculate the algorithmic bandwidth and bus Bandwidth. Differential Revision: D50439185
facebookresearch · Oct 20, 2023 · 943e87c · 943e87c
1 parent 22e6b27
commit 943e87c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/torch_ucc_tracing.cpp b/src/torch_ucc_tracing.cpp
@@ -176,7 +176,8 @@ void CommTraceLogger::recordComms(
       outSize,
       dtype,
       curInSplitSizes_,
-      curOutSplitSizes_);
+      curOutSplitSizes_,
+      world_size);
 
   ++seqnum;