Distributed Matrix Multiplication — DP and FSDP

ML/Infra coding round. You are given a `Communicator` class that simulates inter-device send/recv over Python `Queue`s, plus a partial `dp_mat_mul` and a stub `fsdp_mat_mul`. You complete data-parallel matmul (each device computes a row chunk and gathers to rank 0) and then implement fully-sharded data-parallel matmul from scratch using a rotating all-gather over column shards of `B`. Requirements Starter code: import threading import numpy as np from queue import Queue class Communicator: """Simulates inter-device communication using Queues.""" def __init__(self, num_devices: int): self.num_devices = num_devices self.inboxes = [Queue() for _ in range(num_devices)] def send(self, src: int, dst: int, data: np.ndarray) -> None: self.inboxes[dst].put((src, data)) def recv(self, dst: int) -> tuple: return self.inboxes[dst].get() def _compute_fn(comm, rank, a_chunk, b, result): """TODO: per-device compute for Data Parallel.""" pass def dp_mat_mul(a: np.ndarray, b: np.ndarray, num_devices: int) -> np.ndarray: comm = Communicator(num_devices) result = [None] a_chunks = np.array_split(a, num_devices, axis=0) threads = [] for rank in range(num_devices): t = threading.Thread(target=_compute_fn, args=(comm, rank, a_chunks[rank], b, result)) threads.append(t) t.start() for t in threads: t.join() return result[0] def fsdp_mat_mul(a: np.ndarray, b: np.ndarray, num_devices: int) -> np.ndarray: """ Fully Sharded Data Parallel. A is row-sharded, B is column-sharded across devices. Use all-gather: rotate B shards, each device accumulates partial results. """ pass Driver expectations: M, K, N, num_devices = 8, 6, 4, 2 result_dp = dp_mat_mul(a, b, num_devices) assert np.allclose(result_dp, expected) result_fsdp = fsdp_mat_mul(a, b, num_devices) assert np.allclose(result_fsdp, expected) DP: each device computes a_chunk @ b and sends the partial result to rank 0; rank 0 concatenates row-wise. FSDP: both A (rows) and B (columns) are sharded; over num_devices rounds, every device rotates its B shard to its right neighbor and accumulates the partial product into the correct column slice of its local output. Final assertion is np.allclose(result, a @ b) for both functions. Constraint: the solution must work for any num_devices that evenly divides both M and N, including the degenerate num_devices = 1 (single device → just a @ b). Notes The Communicator is intentionally simple — send/recv are blocking queues. Deadlock-by-design is the most common failure: a naïve FSDP loop where every device first sends and then receives will hang because every queue is full at startup. The fix is to alternate send/recv carefully, or to spawn the receive thread before the send. In FSDP, the rotation step requires each rank to know where its current B shard maps in the output column space. Track current_shard_index = (rank - step) mod num_devices and accumulate into output[:, current_shard_index*shard_w:(current_shard_index+1)*shard_w]. A cleaner, split-robust variant: send (owner_index, shard) tuples so each shard's owning column-range rides along with the data, then place it via precomputed col_offsets[owner] / col_sizes[owner] — this stays correct even when np.array_split yields uneven column widths. np.array_split is allowed; use axis=0 for A rows and axis=1 for B columns. Watch the dtype — accumulating into np.zeros_like(expected) matters; an int-dtype accumulator silently truncates float partial sums. The round is paired with a 20-minute research talk on the same day; prepare a 10-slide deck about a project you can defend at depth. The canonical PyTorch FSDP framing is all-gather (forward, to materialize the full unsharded parameter) then reduce-scatter (backward, to shard gradients back across ranks) — "rotating all-gather" is the candidate's mental model for this exercise, but if the interviewer probes the production analog, name FullyShardedDataParallel and the all-gather + reduce-scatter cycle. The FULL_SHARD strategy shards parameters, gradients, and optimizer state. Collective-vocabulary for the follow-up: all-reduce (DP gradient sync), all-gather (FSDP forward, materialize full params), reduce-scatter (FSDP backward, shard grads back); NCCL picks a ring topology for bandwidth-bound large tensors and a tree topology for latency-bound small ones. Grok-scale training composes these — FSDP across nodes, tensor parallel within a node (split one matmul across GPUs over fast NVLink), pipeline parallel across layer stages, and sequence parallel along the sequence dim for long context. Preparation Implement DP from the starter in under 10 minutes; the bulk of the round is FSDP. Practice the FSDP rotation pattern by hand on num_devices=2 and num_devices=4 until you can derive the shard-index formula on the fly. Read up on FSDP / Megatron-LM tensor-parallel mechanics for the underlying intuition; the round rewards candidates who can explain why sharding parameters + gradients minimizes peak memory vs. plain data parallel. Concretely: DP replicates all of B on every device (O(model) memory), so a 400 GB model cannot fit an 80 GB GPU; FSDP shards it to O(model / num_devices) — that memory argument is the whole point of the technique. Write a sequential reference (a @ b) and diff against your DP/FSDP outputs to debug — np.allclose failures usually point to the wrong axis split.