clean ups

Author: Saeed Maleki

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)

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()