Distributed Matrix Multiplication with Data Parallel and FSDP Simulation

Problem: Implement Two Distributed Matrix Multiplication Strategies (DP and FSDP) Using a Communication Simulator You need to implement distributed matrix multiplication C = A @ B in a simulated multi-rank environment. The starter code provides a Communicator (using queues/message passing to mimic cross-device communication) and partially implemented functions. You must complete/implement: dp_mat_mul(...): a Data Parallel (DP) version. The starter code already contains a thread/concurrency structure, but the compute_fn is empty; you need to fill in the computation and any required communication/synchronization. fsdp_mat_mul(...): a Fully Sharded Data Parallel (FSDP) version. This function is largely blank and must be implemented from scratch. Requirements For multiple ranks, compute the correct matrix multiplication result, numerically matching single-machine A @ B. Use the Communicator communication primitives correctly to exchange required shards/intermediate data between ranks. Your implementations should reflect the conceptual differences: DP: each rank processes a different shard of the input data; model parameters are not sharded (or are equivalently available on every rank). FSDP: model parameters are sharded across ranks; the computation may require parameter exchange (e.g., temporary gather) to perform the multiplication. I/O (abstract) Inputs: matrices A, B (NumPy arrays), rank info (e.g., rank, world_size), and a Communicator instance. Output: each rank should return its local output shard, or the full C if that is what the starter code expects. What to Specify/Handle The sharding dimension you choose for A/B/C under DP/FSDP. Where communication happens (input distribution, parameter gather, output gather, etc.). Avoid deadlocks by using consistent send/recv ordering and proper synchronization. Example Test Cases The original post does not include exact function signatures or the Communicator API, so the tests are expressed in terms of correctness: after simulating multiple ranks and merging their outputs, the result must match A @ B. A is 4x3, B is 3x5, world_size=2, shard A by rows; concatenating per-rank outputs should equal A@B. A is 8x8, B is 8x8, world_size=4, shard B (or as required by the prompt); merging outputs should equal A@B. Edge: world_size=1 should degenerate to single-machine matmul for both implementations. Edge: when A.shape[0] is not divisible by world_size, handle the remainder per prompt (e.g., uneven last shard). Multiple random trials; validate via np.allclose. Example Input (conceptual) A=rand(4,3), B=rand(3,5), world_size=2, shard A by rows; run dp_mat_mul per rank; gather rows Output gathered C equals A@B (allclose)