Testing, Benchmarking, and Optimizing Horovod collectives on Jean-Zay

[EiffL]

This issue is to track the developments needed to finalize and validate the modified version of Horovod we developed. This overarching goal will encapsulate several smaller issues.

Goal

By the end of the hackweek, having a tested code with an associated Pull Request to https://github.com/horovod/horovod which can fully support our needs for Mesh TensorFlow.

Context

With @kimchitsigai and @mypey we worked on some modifications to Horovod that can support multiple communicators. A description of what we did can be found here: DifferentiableUniverseInitiative/horovod#2
In parallel, a different proposal for supporting multiple groups of devices was proposed here horovod/horovod#2839
In the end, probably one of these 2 implementation will be merged, but we can try to find which one works the best for our purposes.

Participants

The main participants to this task are:

• @EiffL
• @andrevitorelli

Tasks

• : Run profiling on 3D demo FFT Benchmarking 3D FFTs with NVIDIA Nsight Systems #2
• : Identify and solve the origin of deadlocks in our modified horovod: multiple_communicators branch gets deadlock on Alltoall horovod#5
• : Test alternative implementation proposed here Concurrently running collective operations on process subsets [TensorFlow] horovod/horovod#2839
• : Test scaling of distributed operations for large tensors
• : Depending on the results above, maybe implement additional specialized collectives for expensive Mesh TensorFlow bottlenecks

[EiffL]

And actually I added a new task, following some weird deadlocks identified by @andrevitorelli and documented there DifferentiableUniverseInitiative/horovod#5
@kimchitsigai this is probably a bug in our horovod modifications somewhere :-|

[andrevitorelli]

Just to leave the error message here:

<...>/horovod/horovod/common/stall_inspector.cc:105] One or more tensors were submitted to be reduced, gathered or 
broadcasted by subset of ranks and are waiting for remainder of ranks for more than 60 seconds. This may indicate that different 
ranks are trying to submit different tensors or that only subset of ranks is submitting tensors, which will cause deadlock.

Missing ranks:
0: [iFFT3D_2/HorovodAlltoall_iFFT3D_2_stack_0]