A context manager to optimize communication

[siddharth9820]

If the forward and backward pass are done within this context, axonn overlaps all the tensor parallel communication in the backward pass (reduce scatters for depth tensor parallelism and all reduces for row and column tensor parallelism)

Example of Usage

for batch, label in loader():
     with axonn.intra_layer.optimize_communication():
         out = model(batch)
         out.backward() # enqueues reduce scatters asynchronously
         # gradients are not ready

     # as soon as we exit the context all reduce scatters are synchronized and the gradients 
     # are ready
     optimizer.step() # important to run OUTSIDE context.

[siddharth9820]

[siddharth9820]

ToDo

• CI is broken when it tries to test asynchrony.

Gotchas

• overlap_reduce_scatter=True might cause issues with frameworks that try to attach gradient hooks to weight tensors - like DDP, ZeRO.