Zero update

[ngc92]

First, clean up the multi-gpu code a bit. We only need to #ifdef away sections that would lead to compile errors, anything else can be handled with a regular plain if, making the code more readable. Also gets rid of the two copies of the update function, let's keep the unified as long as we can.

the current ZERO-1 implementation performs an all-reduce on the gradients, but only one shard is actually needed per rank.
As per the zero paper,

State-of-art implementation of all-reduce uses a two-step approach, where the first step is a reduce-scatter operation, which reduces different part of the data on different process. The next step is an all-gather operation where each process gathers the reduced data on all the process. The result of these two steps is an all-reduce.

The claim that ZERO-1 does not introduce communication overhead hinges on the fact that we just insert the adam update between these two ops. So we go from reduce-scatter -> all-gather -> adam to reduce-scatter -> adam -> all-gather.

[karpathy]

dam, good catch, we were doing a lot more communication than we needed to. merging

[karpathy]

also nice to merge the two updates into one function 👍