Data Parallelism: A Toy Example Splitting a Batch Across Simulated Workers

The problem, from first principles

01-memoization's tricks all made ONE worker's training loop cheaper. Data parallelism is the opposite strategy: instead of making one worker faster, split a single batch across MANY workers, have each compute a gradient on its own slice independently, and combine their answers back into one update — the foundational idea every larger distributed-training scheme (04-collective-communication-primitives through 07-ring-attention-online-softmax) builds on.

From theory to code

Implement local_gradient (the ordinary MSE gradient, computed on just one worker's shard), split_batch_across_workers, average_gradients, and data_parallel_gradient, which chains all three into a full toy data-parallel training step.

Constraints

• split_batch_across_workers assumes the batch size divides evenly by num_workers (a real DDP setup requires every replica to use the same local batch size).
• local_gradient(X_shard, y_shard, w) computes (2 / n_shard) * X_shard.T @ (X_shard @ w - y_shard) — the same MSE gradient formula, applied to only that shard's rows.
• average_gradients is a plain elementwise mean across the list of per-worker gradients.
• data_parallel_gradient(X, y, w, num_workers) must equal the full-batch gradient computed directly on all of X/y at once, whenever the shards are equal-sized.

Hints

Open one at a time. Each gives away a little more than the last.