Distributed Training Step End-to-End

Implement distributed_step: average linear-regression gradients across supplied worker batches, take one Adam update, then update an exponential moving average (EMA) of the model weights.

This is a simulation in one function. Each worker's data is a slice of the input tensor; do not launch processes or communicate over a network. The exercise models the arithmetic of gradient sharing, not a distributed runtime.

Inputs and output

worker_xs: (W, A, B, d), with W workers, A micro-batches per worker, B examples per micro-batch and d features.
worker_ys: (W, A, B), regression targets.
weights: (d,), the shared model weights before this step.
ema_weights: (d,), the previous moving-average weights.
m and v: (d,), the previous, uncorrected Adam first and second moments.
lr: a finite, nonnegative learning rate.
beta1 and beta2: coefficients in [0, 1).
eps: a positive finite number added outside Adam's square root.
decay: the EMA coefficient in [0, 1].
t: a positive integer, the Adam update number for this call, starting at 1.

W, A, B and d are positive. Data and state tensors are finite, share a floating-point dtype and device, and have matching shapes. Entries of v are nonnegative. Leave all inputs unchanged.

Return a flat tensor of shape (4*d,), concatenating [new_weights, new_ema_weights, new_m, new_v] in that order. Preserve dtype and device. Return raw moments, not their bias-corrected versions. State is supplied explicitly; do not keep hidden state between calls or reset nonzero input moments.

1. Average the gradients

The model predicts X @ weights, with no separate bias term. For one micro-batch, use mean squared error without a factor of one half:

grad = (2/B) * X.T @ (X @ weights - y)

Compute every micro-batch gradient using the same input weights. Average the A gradients within each worker, then average the W worker gradients to obtain g. There is no parameter update between micro-batches or workers.

Because every micro-batch has the same B examples, this equals the gradient of the mean loss over all WAB examples. You may use that equivalence to vectorize the calculation. Do not divide by d or average the result again before Adam. In particular, average gradients before squaring them for the second moment; the average of squared worker gradients is a different quantity.

2. Apply one Adam update

All operations below are elementwise:

new_m = beta1*m + (1-beta1)*g
new_v = beta2*v + (1-beta2)*g*g
m_hat = new_m / (1-beta1**t)
v_hat = new_v / (1-beta2**t)
new_weights = weights - lr*m_hat/(sqrt(v_hat)+eps)

Use the supplied t as written. It counts optimizer updates, not examples, workers or micro-batches. A caller performing the next update supplies the returned raw moments and increments t by one. There is no weight decay, clipping, momentum beyond Adam's moments or AMSGrad maximum.

3. Update the moving-average weights

Use the newly updated model weights:

new_ema_weights = decay*ema_weights + (1-decay)*new_weights

Do this once, after Adam. Do not use the pre-update weights or add EMA bias correction. With decay=0, the EMA equals new_weights; with decay=1, it stays unchanged. With lr=0, model weights stay unchanged, but moments still update and the EMA can still move toward the model weights.

For a one-feature example, let X=[1], y=[1], weights=[0], m=v=[0], lr=0.2, beta1=beta2=0, eps=2, ema_weights=[1], decay=0.5 and t=1. Then g=-2, new_m=-2, new_v=4, new_weights=0.1 and new_ema_weights=0.55. Return [0.1, 0.55, -2, 4].

Build the gradient, Adam and EMA calculations from basic tensor operations. Do not use autograd, optimizer helpers or distributed collective libraries.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Combine gradients before updating shared parameters

Data-parallel workers start from the same model parameters and process different examples. For a synchronized training step, combine their gradients before the update. An all-reduce gives the combined result to every worker. Check whether the operation returns a sum or a mean.

For equal-sized batches with mean gradients 2 and 6, the combined mean is (2 + 6) / 2 = 4. If the first worker has one example and the second has three, the example-weighted mean is (1*2 + 3*6) / 4 = 5. Averaging worker means gives the correct global mean only when their counts match, or by coincidence.

Pass gradient chunks around a ring

Ring all-reduce splits each tensor into chunks. First, reduce-scatter adds contributions until each worker holds a completed sum for its assigned chunk. Then all-gather distributes those completed chunks so every worker has the whole result. With N workers and the usual equal-chunk scheme, each phase takes N - 1 neighbor exchanges.

In a simulation, all sends for a round must read the previous round's buffers. Updating a buffer and then using its new contents for another send in the same round lets information travel farther than intended. A correct local simulation checks the arithmetic; it does not measure real network performance.

Take one consistent optimizer step

Microbatch accumulation uses the same weighting principle. Combine the gradients, update optimizer state and parameters once, and update any exponential moving average of parameters at the specified point.

Workers also need consistent optimizer state. Running independent Adam updates and averaging the resulting weights is generally different from applying Adam once to the combined gradient, because Adam's update is nonlinear in the gradient.

Detecting runtime…