Megatron-LM: Column-parallel Weight Shards
Beginner
The problem, from first principles
Megatron-LM (Shoeybi et al., 2019) splits a large linear layer across GPUs. In the column-parallel case each device holds a block of the output columns, so each computes part of the output independently.
From theory to code
Implement column_shard(W, world), splitting the weight into one column block per device.
Constraints
- The output dimension is divisible by the number of devices.
Hints
Hint 1
Split along the output axis into equal parts.
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.