Accumulate and Average Replica Gradients
Easy · Solved in PyTorch · Machine learning coding practice
Problem
In Megatron's 3D parallelism, every data-parallel replica holds the same parameter shard and a different slice of the global batch (Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)). Each replica runs K microbatches, accumulates their gradients locally, and only then synchronizes with the other replicas — one collective per step instead of one per microbatch.
microbatch_grads holds the per-microbatch gradients of one parameter shard of shape (O, I), arranged as (R, K, O, I) for R replicas and K microbatches. microbatch_sizes holds the integer sample count n_{rk} \geq 1 of every microbatch, of shape (R, K). Each stored gradient G_{rk} is the mean over its own n_{rk} samples, so microbatches of different sizes must not be given equal weight. Do not modify the inputs.
Topics: megatron, data-parallel, gradients, pytorch, numpy
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Accumulate and Average Replica Gradients in the browser and run it against the tests.