Megatron-LM: Tensor-parallel Matrix Multiply
Intermediate
The problem, from first principles
Column-parallel matrix multiplication gives the same result as the dense product. Each device computes its own output columns from the replicated input, and concatenating those columns reproduces the full output without communication.
From theory to code
Implement tensor_parallel_matmul(x, shards), multiplying by each column shard and concatenating.
Constraints
- The input is the same on every device.
Hints
Hint 1
Multiply the input by each shard, then concatenate the outputs along the feature 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.