Fix DDP with Zero1

[NouamaneTazi]

Currently we have for a 1B model with

Zero0 (uses DDP)

Zero1 (DDP broken)

Which makes Zero1 slower, it would be great to fix it (or if not possible overlap the grads comms with compute)

Memory gains

for zero0

12/10/2024 16:19:16 [INFO|DP=0|PP=0|TP=0|ip-26-0-164-187]:  Memory usage: 19909.64MiB. Peak allocated 19909.69MiB. Peak reserved: 49104.00MiB
12/10/2024 16:19:16 [INFO|DP=0|PP=0|TP=0|ip-26-0-164-187]:  Memory usage: 19909.64MiB. Peak allocated 47208.43MiB. Peak reserved: 49104.00MiB

for zero1

12/10/2024 16:19:16 [INFO|DP=0|PP=0|TP=0|ip-26-0-168-238]:  Memory usage: 6517.47MiB. Peak allocated 8584.28MiB. Peak reserved: 34706.00MiB
12/10/2024 16:19:16 [INFO|DP=0|PP=0|TP=0|ip-26-0-168-238]:  Memory usage: 6517.47MiB. Peak allocated 33819.58MiB. Peak reserved: 34706.00MiB

12/10/2024 16:19:03 [INFO|DP=0|PP=0|TP=0|ip-26-0-168-238]: [ZeRO sharding] DP Rank 0 has 43.3M out of 1.04G (4.17%) params' optimizer states
12/10/2024 16:19:03 [INFO|DP=0|PP=0|TP=0|ip-26-0-168-238]: [ZeRO sharding] DP Rank 1 has 43.3M out of 1.04G (4.17%) params' optimizer states
...