Megatron-LM: Summing Row-parallel Partial Outputs
Beginner
The problem, from first principles
In a row-parallel layer each device computes a partial product over its share of the input features. Summing those partial results across devices reproduces the full output. The sum is an all-reduce collective.
From theory to code
Implement allreduce_sum(parts), summing the partial results element-wise.
Constraints
- All parts have the same shape.
Hints
Hint 1
Stack the partial results and sum over the device axis.
Keywords
research-paperssystemsparallelismtraining
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.