Add the world size info in NCCL metadata

Summary: X-link: pytorch/pytorch#111842 This diff adds the world size info in NCCL metadata, as we need the information to calculate the algorithmic bandwidth and bus Bandwidth. ghstack-source-id: 205210972 exported-using-ghexport Reviewed By: aaronenyeshi Differential Revision: D50439185 fbshipit-source-id: bdeb99c7abf262d037f04b693c6f5ff6c7e6a97d
facebookresearch · Oct 25, 2023 · 13ccc87 · 13ccc87
1 parent 22e6b27
commit 13ccc87
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;