Tensor Parallelism Demo (Matrix Multiplication)

In Large Language Models, matrices are too big to fit on one GPU. Tensor Parallelism splits these matrices across multiple GPUs.

In this demo, we simulate Column Parallelism. The Weight Matrix $W$ is split column-wise. GPU 1 calculates the first half of the output, GPU 2 calculates the second half. Finally, they are concatenated (All-Gather).