Project objective
Design a distributed parallelization scheme for a PyTorch layer dominated by GEMM, and implement at least one communication–computation fusion kernel that overlaps GEMM with All-to-All communication. Preserve numerical correctness and improve end-to-end throughput relative to a non-overlapped baseline.
Justify the decomposition, fine-grained pipeline, scheduling policy, and dependency management. Evaluate on multiple GPUs, including a cross-node configuration, and analyze differences between intra-node and inter-node communication. Quantify exposed communication, the fraction hidden by overlap, layer latency, and end-to-end throughput.
The supplied --overlap kernel entry initially delegates to sequential execution. Existing with-sm and sm-free implementations are comparison points. Selecting one of these options without implementing a new fusion kernel does not complete the project.
Interface, layout, and parameters
Implement research_topic_2/overlap-src/kernel_overlap_plugin.py, or set A2A_FFN_KERNEL_PLUGIN to your importable module:
def forward_ffn(ffn, inputs, *, comp_only=False):
# inputs and output: [B,T,D]
# Starter behavior only; replace with your overlap implementation.
return ffn.forward(inputs, overlap="none", comp_only=comp_only)
ffn provides w1,b1,wout,bout,norm_x,norm_gate,ep_group,ep_index,ep_group_size,fwd_mode,bwd_mode,gemm_backend. Let P=ep_group_size, L=T/P, and H=D*k. Require T%P==0. Rank ordering follows the process group.
| Data | Per-rank shape | Meaning |
|---|---|---|
| inputs | [B,T,D] |
Local batch |
| Dispatched input | [P*B,L,D] |
Split tokens into P slices; concatenate by source rank along batch |
| w1 / b1 | [L,D,2H] / [L,2H] |
Parameters for the rank's token slice |
| wout / bout | [L,H,D] / [L,D] |
Second GEMM and bias |
| norm_x / norm_gate | [H] weights/biases |
LayerNorm; default epsilon 1e-5 |
| output | [B,T,D] |
Local batch with full token dimension restored |
Do not mutate inputs. comp_only=True removes communication for measurement and must preserve _forward_comp_only layout semantics. Report comp-only measurements separately from full-layer results. There is no separate backward hook: autograd must propagate through forward_ffn to the input and all parameters.
Computation and communication example
The dense_v2 + bmm reference computes:
Xlocal = A2A(inputs, split_dim=1, concat_dim=0)
Z[n,l,:] = Xlocal[n,l,:] @ w1[l,:,:] + b1[l,:]
(U,G) = split Z into two H-dimensional halves
LN(z) = ((z-mean(z))/sqrt(mean((z-mean(z))^2)+1e-5))*gamma + beta
Hlocal = LN_x(U) * silu(LN_gate(G))
Ylocal[n,l,:] = Hlocal[n,l,:] @ wout[l,:,:] + bout[l,:]
output = A2A(Ylocal, split_dim=0, concat_dim=1)
LayerNorm includes mean subtraction. Preserve normalization, activation, bias, and the selected precision. Compare against overlap=none with identical dtype/fwd_mode/gemm_backend.
For the layout example P=2,B=1,T=2,D=1, rank 0 has [[[10],[11]]] and rank 1 has [[[20],[21]]]. Dispatch produces [10,20] on rank 0 and [11,21] on rank 1, each shaped [2,1,1]. If local computation is identity, combine reconstructs the original rank inputs. This example checks exchange ordering; the actual benchmark must execute the full layer above.
The benchmark generates inputs using --seed 42 and initializes parameters with its he-uniform function. Keep the generator, precision, and random state fixed. No advertising dataset or model weights are needed.
Environment and quick start
Extract the benchmark archive and enter research_topic_2/. Use Linux x86_64, Python 3.12, at least two Hopper/SM90 GPUs (H20/H100/H800), and a driver supporting CUDA 12.6. Dependencies are PyTorch 2.6.0+cu126, NCCL 2.21.5, Triton 3.2.0, and NumPy 1.26.4. The bmm path does not require a custom Hopper extension.
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
nvidia-smi -L
NPROC=2 bash scripts/run_prod_local.sh --batch-size 4 --tokens 16 --dim 128 --k 2 --gemm-backend bmm --no-compile --overlap none --warmup 5 --num-tests 20
NPROC=2 bash scripts/run_prod_local.sh --batch-size 4 --tokens 16 --dim 128 --k 2 --gemm-backend bmm --no-compile --overlap kernel --check-correctness bitwise --warmup 5 --num-tests 20
# Implement the plugin, then rerun correctness and measure backward.
NPROC=2 bash scripts/run_prod_local.sh --batch-size 4 --tokens 16 --dim 128 --k 2 --gemm-backend bmm --no-compile --overlap kernel --backward --check-correctness approx --warmup 5 --num-tests 20
After correctness passes, measure the larger workload:
NPROC=8 bash scripts/run_prod_local.sh --overlap none --warmup 10 --num-tests 100
NPROC=8 bash scripts/run_prod_local.sh --overlap kernel --warmup 10 --num-tests 100
NPROC=8 bash scripts/run_prod_local.sh --overlap kernel --backward --check-correctness approx --warmup 10 --num-tests 100
Defaults are B=512,T=64,D=1280,k=6,H=7680,P=8. In bf16, one dispatch or combine sends B*T*D*2*(P-1)/P=70 MiB per rank, excluding the local slice. One forward pass performs two exchanges.
Cross-node tests and correctness
Use two nodes with identical software and code. Set MASTER_ADDR to node 0's IP address reachable from node 1; use port 29500. Replace NODE0_IP below before launching both commands concurrently on their respective nodes:
# Node 0
NNODES=2 NPROC=2 NODE_RANK=0 MASTER_ADDR=NODE0_IP bash scripts/run_prod_cross.sh --batch-size 4 --tokens 16 --dim 128 --k 2 --overlap kernel --check-correctness approx
# Node 1
NNODES=2 NPROC=2 NODE_RANK=1 MASTER_ADDR=NODE0_IP bash scripts/run_prod_cross.sh --batch-size 4 --tokens 16 --dim 128 --k 2 --overlap kernel --check-correctness approx
Repeat for --overlap none and --backward. This torchrun launch does not require MPI. Configure NCCL's network interface for the available IB/RDMA network. A submission without cross-node execution is a partial milestone, not complete cross-node validation.
The documented default workloads are single-node P=8 and two-node P=16 with (B,T,D,k)=(512,64,1280,6). The P=2/P=4 small-shape commands above are development configurations, not additional fixed test cases. --comm-scale 2 and --compute-scale 2 separately examine communication-bound and compute-bound behavior. Communication scaling only pads the payload; compute scaling increases H. The full RankMixer block also requires D%T==0.
The current forward checker uses torch.equal even when --check-correctness approx is selected. It checks a batch slice of min(B,2) against overlap=none; a mismatch raises an assertion with maximum absolute error and mismatch count. Do not describe forward approx mode as a tolerance-based check.
Backward checks also use min(B,2). Bitwise mode checks input and parameter gradients with torch.equal. Backward approx uses per-gradient Frobenius relative error at most 1e-5 for float32 or 2e-2 for bf16/fp16, with the denominator clamped at 1e-12. Maximum absolute error and p99/max ULP are diagnostics, not acceptance gates. These checks run only for a non-large block, a non-comp-only run, enabled checking, and overlap in with-sm/sm-free/kernel. Backward requires a supported dense path.
The complete timed output is separately checked for shape. The reduced-batch numerical check does not establish full-batch numerical coverage: report that limitation and supplement it when making full-workload claims. Preserve the supplied checker. The CLI accepts configurable workloads; the full 32-configuration regression matrix remains with the instructor.
The main input/model seed is seed+rank (default seed 42). The backward numerical check uses a separate generator with seed+7; timed backward uses the full-size fixed upstream gradient. Compare the same active parameter-gradient keys in the baseline and overlap paths.
Public checks and instructor coverage
python preflight.py --target cuda
# Or run the public cases directly:
python public_check.py
The public runner uses two GPUs, (B,T,D,k)=(4,16,128,2), seed 7101, bmm, no compilation, two warmups and three measurements. It runs sequential forward, kernel forward with bitwise checking, and kernel forward/backward with approx checking. Forward remains bitwise in all checked runs.
The instructor retains the full 32-configuration regression matrix and adds independent input seeds and shapes. Held-out bmm inputs span 1<=B<=512, 16<=T<=64, 128<=D<=1280, D%T==0, T%P==0, and k in {2,3,6}. Test valid EP sizes supported by the topology, including cross-node operation. Public success does not establish coverage of larger batches or other topologies. Optional fused/FP8/CuTeDSL paths require additional backend dependencies; report their availability explicitly.
CPU environment check
Install requirements-cpu.txt in a separate Python 3.12 environment and run python preflight.py --target cpu. This executes the actual complete block and its gradients at EP=1 in float32/bf16. It verifies local compute only; All-to-All, overlap, CUDA timing and cross-node behavior require GPUs. See ENVIRONMENT.md and environment verification.
Timing and analysis
The actual parser defaults to 50 warmups, 50 measurements, a 10 ms prewarm sleep, L2 flushing enabled, and compilation enabled. The explicit commands above override warmup/count for development or longer runs; record all overrides. The benchmark measures one RankMixer block, not a full six-layer training step or optimizer update. --backward measures its forward and backward with a fixed upstream gradient.
The timing utility flushes L2 and aligns ranks before recording each start event, measures the call with CUDA events, and synchronizes before reading times. It drops the first measurement when multiple measurements exist, then returns mean/min/max in seconds and prints microseconds. Rank 0 prints its local statistics; the runner does not automatically report the slowest-rank global latency. Additional global throughput or percentile analysis must explicitly collect per-rank/raw measurements. Keep CUDA graph, L2, and compilation settings identical between comparisons.
Measure full sequential latency Tserial, overlapped latency Toverlap, isolated communication Tcomm, and isolated compute Tcompute. Report (Tserial-Toverlap)/Tcomm as an estimated hidden fraction, preserving negative values. This estimate can include changes in compute and scheduling: validate actual overlap and exposed communication using device traces and the dependency critical path. Do not use forward-only FLOPs to derive forward-plus-backward TFLOPS.
Required submission
Submit a proposal explaining the parallel strategy, fusion boundaries, pipeline, dependencies, correctness plan, topology, baseline, and methodology; a code ZIP with the new plugin/kernel, configurations, single-node and cross-node launch commands, correctness logs, raw timings, and traces; and a separate report.
The demo must run sequential and fused implementations, check outputs and applicable gradients, and report communication and throughput. Discuss intra-node/inter-node differences, achieved overlap, bottlenecks, scaling, buffer reuse, synchronization, limitations, and unsuccessful approaches.
Evaluation considers numerical correctness, sound parallel and pipeline design, communication hiding, end-to-end throughput, topology-aware analysis, scalability, reproducibility, and engineering evidence.