ZeRO-Sharded Adam Step

The problem

In plain data parallelism every worker keeps the full Adam state: $m$ and $v$, the two running averages, for every parameter. With ZeRO (Zero Redundancy Optimizer), the parameters are divided into shards, one per worker, and worker $i$ keeps the Adam state of its shard only. You will simulate one training step with this split, and check that it gives exactly the result of ordinary Adam.

The step

1. Reduce-scatter the gradients: worker $i$ receives the mean over all workers of the gradient, but only for the parameters of its shard.
2. Each worker runs Adam on its shard only, with its own shard of $m$ and $v$.
3. All-gather the updated shards, so that every worker has the full, updated parameters again.

The shards

For $n$ parameters and $N$ workers, shard $i$ is np.array_split(np.arange(n), N)[i]: the $i$-th of $N$ nearly equal slices of the parameter indices.

Example. 5 parameters, 2 workers: np.array_split(np.arange(5), 2) is [array([0, 1, 2]), array([3, 4])]. Worker 0 owns parameters 0, 1, 2 and keeps $m$ and $v$ of length 3; worker 1 owns parameters 3, 4 and keeps $m$ and $v$ of length 2.

Example. One step with those shards. All 5 parameters are 1; the two workers computed the local gradients [0.2, 0.4, 0.6, 0.8, 1.0] and [0.4, 0.0, 0.2, -0.8, 3.0].

1. Reduce-scatter: worker 0 gets the mean of its slice, [0.3, 0.2, 0.4]; worker 1 gets [0.0, 2.0].
2. Adam on each shard (first step, lr = 0.1): every parameter with a non-zero gradient moves by lr against it (section 3), so worker 0 gets [0.9, 0.9, 0.9] and worker 1 gets [1.0, 0.9] (parameter 3 has a mean gradient of 0, and does not move).
3. All-gather: everyone has [0.9, 0.9, 0.9, 1.0, 0.9], exactly what Adam on the full vector with the mean gradient gives.

What to implement

• init_states(n_params, n_workers): a list of $N$ pairs (m_shard, v_shard) of zeros, one per worker, each the size of that worker's shard.
• zero_step(param, grads, states, t, lr=1e-2): param is the full parameter vector (identical on every worker), grads is the list of the $N$ workers' local gradients (full length), states comes from init_states or a previous step, and t is the step number. Return (new_param, new_states).

adam_update is given: one Adam step on one array.

Requirements

• The result must equal ordinary Adam applied to the full vector with the mean gradient.
• Each worker's m and v stay the size of its shard: that is the memory saving.
• No Python loop over parameters (a loop over the $N$ workers is fine).