Column and Row Tensor Parallel Linear
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Tensor parallelism splits one linear layer across ranks so each rank stores and multiplies only a slice of the weights (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)). Simulate both decompositions of
Topics: inference, vllm, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Column and Row Tensor Parallel Linear in the browser and run it against the tests.