# Derive Sharded Matrix Multiplication and Backpropagation Consider `Y = X @ W`, where `X` has shape `(B, D)` and `W` has shape `(D, H)`. Training runs on `P` devices. Assume `H` is divisible by `P`, `W` is sharded by output columns, and device `p` stores `W_p` with shape `(D, H/P)`. Each device initially has the full `X` and receives its matching upstream-gradient shard `dY_p`. Derive the forward computation and backward gradients for `X` and `W`. Identify every collective communication operation required to produce a replicated full `dX` while keeping `dW` column-sharded. Then give NumPy-like pseudocode and a debugging plan for shape, numerical, and collective errors. ### Constraints & Assumptions - Use standard dense matrix multiplication with no bias. - Devices execute synchronously and use the same floating-point dtype. - Local matrix multiplication is available; collective operations are abstract primitives. - The scalar loss supplies `dY` with the same logical shape as `Y`. - Discuss, but do not implement, the alternative where `X` or the contraction dimension is sharded. ### Clarifying Questions to Ask - Which tensor dimensions are sharded and which tensors must be replicated after each phase? - Is the objective to minimize communication volume, peak memory, or latency? - Are gradients accumulated across data-parallel replicas as well as tensor-parallel devices? - What numerical tolerances and reference implementation are available for debugging? ### Hints - Write local tensor shapes beside every expression. - Derive gradients from differentials before deciding on collectives. - Compare a tiny distributed simulation against one unsharded calculation. ### What a Strong Answer Covers - Correct forward and backward equations with consistent local shapes. - A precise explanation of why and where a collective is needed. - Communication and compute complexity, plus memory trade-offs. - Tests for shard ordering, transposes, gradient scaling, and collective semantics. - Awareness that a different sharding axis changes the communication pattern. ### Follow-up Questions 1. What changes if `W` is sharded along its input dimension instead? 2. How would data parallelism compose with this tensor-parallel layout? 3. Where can communication overlap with computation? 4. How would you detect an accidentally duplicated or missing gradient reduction?