Tensor-Parallel Forward and Backward Collectives

Hard · Solved in PyTorch · Machine learning coding practice

Problem

Classic tensor parallelism splits one Transformer MLP across P ranks and inserts a single pair of conjugate operators, f and g (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)). On the forward pass f copies its input to every rank and g sums the rank outputs; on the backward pass the roles swap, so f sums the rank input gradients and g copies the output gradient to every rank.

Simulate both operators, and the backward pass between them, for the two-matrix MLP

Topics: megatron, tensor-parallel, backpropagation, pytorch, numpy

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Tensor-Parallel Forward and Backward Collectives in the browser and run it against the tests.

Related problems

• Vocabulary-Parallel Embedding and Cross-Entropy Hard
• Accumulate and Average Replica Gradients Easy
• Schedule Pipeline Microbatches Medium
• Bradley Terry Ranking Medium
• Looped Transformer: Adaptive Early Exit Medium
• Looped Transformer: Shared-Weight Unrolling Easy