clean ups

microsoft · Apr 28, 2023 · 030a750 · 030a750
1 parent db171e0
commit 030a750
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/examples/mscclang/allgather_a100_pcie.py b/examples/mscclang/allgather_a100_pcie.py
@@ -6,8 +6,8 @@
 from msccl.topologies import *
 from msccl.language.collectives import AllGather
 
-# Allpairs allgather for A100 
-def allgather_allpairs(gpus, instances, protocol):
+# Hierarchical allgather for A100 
+def allgather_hier(gpus, instances, protocol):
     size = gpus
     chunksperloop = 1
     topology = fully_connected(gpus)
@@ -39,4 +39,4 @@ def allgather_allpairs(gpus, instances, protocol):
 parser.add_argument('--protocol', type=str, default='LL128', choices=['Simple', 'LL', 'LL128'], help ='NCCL protocol. Default: Simple')
 args = parser.parse_args()
 
-allgather_allpairs(args.num_gpus, args.instances, args.protocol)
+allgather_hier(args.num_gpus, args.instances, args.protocol)
diff --git a/examples/mscclang/allgather_allpairs.py b/examples/mscclang/allgather_allpairs.py
@@ -10,7 +10,7 @@
 def allgather_allpairs(gpus, instances, protocol):
     size = gpus
     topology = fully_connected(gpus)
-    collective = AllGather(size, size, True)
+    collective = AllGather(size, 1, True)
 
     with MSCCLProgram(f"allgather_allpairs", topology, collective, instances,
          protocol=protocol, threadblock_policy=ThreadblockPolicy.manual):
@@ -20,7 +20,7 @@ def allgather_allpairs(gpus, instances, protocol):
             for r2 in range(gpus):
                 if r1 != r2:
                     index = 0
-                    c = chunk(r1, Buffer.input, index, size)
+                    c = chunk(r1, Buffer.input, index, 1)
                     c.copy(r2, Buffer.input, index, sendtb=r2, recvtb=r1)
         XML()
         Check()