Note: Tensor Parallelism, Splitting One Layer's Matmul Across GPUs

The problem, from first principles

02-pipeline-parallelism-bubble-fraction split a model between layers, putting different layers on different GPUs. Tensor parallelism goes a level deeper: it splits a SINGLE layer's matmul itself across GPUs, so even one Linear layer too big to fit on one GPU can still run — but naively splitting a matmul can silently produce the wrong answer unless you split it along the right axis, and know when the pieces can be combined by simple concatenation versus when they need to be summed.

From theory to code

Implement column_parallel_linear and row_parallel_linear — the two standard ways (from Megatron-LM) to split a Linear layer's weight matrix across GPUs — plus the split_weight_by_output_features/split_by_input_features helpers that set each scheme up correctly.

Constraints

• column_parallel_linear: each GPU holds a row-slice of the weight matrix (a slice of the OUTPUT features) and sees the SAME full input; the GPUs' outputs are concatenated along the last axis to reconstruct the full output — no communication needed mid-computation.
• row_parallel_linear: each GPU holds a column-slice of the weight matrix (a slice of the INPUT features) and only its matching slice of the input; each GPU computes a partial output, and the partials are summed (an all-reduce, 04-collective-communication-primitives) before the (un-split) bias is added once.
• Both schemes must reconstruct linear(input, weight, bias) (01-hypothesis-function's function) exactly, for any number of GPUs, including uneven splits.

Hints

Open one at a time. Each gives away a little more than the last.