Problem Statement

This coding exercise evaluates your grasp of distributed-computation patterns and NumPy operations. A starter Communicator class models device-to-device messaging with queues. Complete two parallel matrix-multiplication approaches:

1. Data Parallel (DP): Partition matrix A by rows. Every device receives the complete matrix B, computes its assigned rows, and sends the resulting block to rank 0 for gathering.
2. Fully Sharded Data Parallel (FSDP): Partition A by rows and B by columns. Have devices exchange B shards in an all-gather-style rotation, computing and placing each column block as it becomes available.

Starter Code

import threading
import numpy as np
from queue import Queue


class Communicator:
    """Models communication between devices with 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: Compute one device's work for the Data Parallel strategy.
    The device should multiply a_chunk by b and send the result to rank 0.
    """
    # IMPLEMENT HERE
    pass


def dp_mat_mul(a: np.ndarray, b: np.ndarray, num_devices: int) -> np.ndarray:
    """
    Multiply matrices with Data Parallelism.
    Divide A across rows, compute each block independently, and gather at 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:
    """
    Multiply matrices with Fully Sharded Data Parallelism.
    Divide A by rows and B by columns across devices.
    TODO: Build this implementation from scratch.
    Use an all-gather rotation in which B shards circulate while devices
    accumulate the corresponding partial products.
    """
    # IMPLEMENT 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

• All inter-device messaging must use the supplied Communicator class.
• A device thread may access only its own data shard; do not bypass communication through shared-memory shortcuts.
• The computed matrices must be numerically equivalent to a @ b, as checked with np.allclose.
• The implementation should support 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 production LLM training, parameter counts can reach hundreds of billions. Data Parallelism duplicates the complete model on every device, which is impractical when one GPU provides 80GB but the model needs 400GB. FSDP instead shards parameters and optimizer state, changing per-device storage from O(Model) to O(Model/D).

Communication Patterns in Practice

Hybrid Parallelism (Real xAI/Grok Scale)

Training models at Grok scale combines several forms of parallelism:

• FSDP across nodes: distribute model-parameter shards among machines
• Tensor Parallel (TP) within a node: divide individual matrix multiplications among GPUs connected by fast NVLink
• Pipeline Parallel (PP): assign model layers to stages and stream micro-batches through them
• Sequence Parallel (SP): partition the sequence dimension for long-context training