Data Parallel & FSDP Matrix Multiplication · xAI

Problem Statement

This question tests your understanding of distributed systems intuition and NumPy fundamentals. You are given starter code with a Communicator class that simulates inter-device communication using Queues. Your task is to implement two parallel matrix multiplication strategies:

Data Parallel (DP): Split matrix A row-wise across devices. Each device computes its partial result using the full matrix B, then gathers results to rank 0.
Fully Sharded Data Parallel (FSDP): Both A (row-wise) and B (column-wise) are sharded across devices. Each device uses an all-gather pattern — rotating B shards across devices and accumulating partial results.

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 function for Data Parallel strategy.
    Each device computes a_chunk @ b, then sends result to rank 0.
    """
    # YOUR CODE HERE
    pass

def dp_mat_mul(a: np.ndarray, b: np.ndarray, num_devices: int) -> np.ndarray:
    """
    Data Parallel Matrix Multiplication.
    Split A row-wise; each device computes partial result; gather to rank 0.
    """
    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 Matrix Multiplication.
    TODO: Implement from scratch.
    Both A (row-wise) and B (column-wise) are sharded across devices.
    Use all-gather: rotate B shards, each device accumulates partial results.
    """
    # YOUR CODE HERE
    pass

if __name__ == "__main__":
    np.random.seed(42)
    M, K, N, num_devices = 8, 6, 4, 2
    a = np.random.randn(M, K)
    b = np.random.randn(K, N)
    expected = a @ b

    result_dp = dp_mat_mul(a, b, num_devices)
    assert np.allclose(result_dp, expected), "DP result mismatch!"
    print("DP passed!")

    result_fsdp = fsdp_mat_mul(a, b, num_devices)
    assert np.allclose(result_fsdp, expected), "FSDP result mismatch!"
    print("FSDP passed!")

Constraints

Use the provided Communicator class for all inter-device communication
Each device (thread) should only access its own shard of data — no shared memory shortcuts
Results must be numerically equivalent to a @ b (verified via np.allclose)
Your solution should work for any num_devices that evenly divides both M and N

Test Cases

# Test 1: Basic 2-device DP
np.random.seed(42)
a = np.random.randn(8, 6)
b = np.random.randn(6, 4)
assert np.allclose(dp_mat_mul(a, b, 2), a @ b)

# Test 2: Basic 2-device FSDP
assert np.allclose(fsdp_mat_mul(a, b, 2), a @ b)

# Test 3: 4-device DP
a = np.random.randn(16, 8)
b = np.random.randn(8, 12)
assert np.allclose(dp_mat_mul(a, b, 4), a @ b)

# Test 4: 4-device FSDP
assert np.allclose(fsdp_mat_mul(a, b, 4), a @ b)

# Test 5: Single device (degenerate case)
a = np.random.randn(4, 3)
b = np.random.randn(3, 5)
assert np.allclose(dp_mat_mul(a, b, 1), a @ b)
assert np.allclose(fsdp_mat_mul(a, b, 1), a @ b)

Follow-Up Discussion Topics

Why FSDP Over DP?

In real-world LLM training, model parameters can be hundreds of billions. DP replicates the full model on every device, which is infeasible when a single GPU has 80GB memory but the model requires 400GB. FSDP shards both parameters and optimizer states, reducing per-device memory from O(Model) to O(Model/D).

Communication Patterns in Practice

Pattern	Used By	Description
All-Reduce	DP gradient sync	Sum gradients across all devices
All-Gather	FSDP forward pass	Reconstruct full parameters from shards
Reduce-Scatter	FSDP backward pass	Reduce gradients + scatter to shard owners
Ring topology	NCCL backend	Bandwidth-optimal for large tensors
Tree topology	Small messages	Latency-optimal for small tensors

Hybrid Parallelism (Real xAI/Grok Scale)

Training Grok-scale models uses multiple parallelism strategies simultaneously:

FSDP across nodes: shard model parameters across machines
Tensor Parallel (TP) within a node: split individual matrix multiplications across GPUs on the same machine (fast NVLink)
Pipeline Parallel (PP): split model layers across stages, process micro-batches in pipeline fashion
Sequence Parallel (SP): shard along the sequence dimension for long-context training

Problem