Derive Sharded Matrix Multiplication and Backpropagation
Company: OpenAI
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Onsite
# 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?
Overview: Derive forward and backward computation for column-sharded matrix multiplication across multiple devices. Track every local tensor shape, identify the collective needed for replicated input gradients, and propose tests for shard ordering, scaling, and numerical errors.
Read the full OpenAI Machine Learning Engineer interview experience this question came from