Assignment 2 · Systems

What you build

Six pieces, in the order the handout presents them: a benchmarking and profiling harness, an activation-checkpointing study, a FlashAttention-2 Triton kernel, distributed data parallel training, optimizer state sharding, and fully sharded data parallel training. On top of the code there is a substantial written component — every problem has a short deliverable (timings, tables, profiler screenshots, or a few sentences of analysis), and Section 8 is pure pencil-and-paper: deriving when DP, FSDP, TP, and 2D parallelism become communication-bottlenecked.

The repo starts you from an empty cs336_systems module — “there’s no code in here, so you should be able to do whatever you want from scratch.” You are not required to reuse your own Assignment 1 model: the staff solution ships in cs336-basics, so you profile a known-good Transformer (or repoint pyproject.toml at your own). Tests hook in through tests/adapters.py, and you submit writeup.pdf plus code.zip to Gradescope.

Which lectures feed it

This is the systems arc of the course made concrete. Lecture 5 (GPUs) gives you the mental model the profiling section assumes — memory hierarchy, arithmetic intensity, why kernels launch asynchronously and what torch.cuda.synchronize() is actually for. Lecture 6 (kernels, Triton) is the direct preparation for the FlashAttention-2 kernel: tiling, fusion, block pointers, and the online softmax trick all come from there. Lecture 7 (Parallelism I) covers the collectives — all-reduce, all-gather, reduce-scatter — and the DDP variants you implement in Section 5. Lecture 8 (Parallelism II) is the ZeRO/FSDP story behind Sections 6 and 7 and the analysis in Section 8. When stuck on the kernel, rewatch Lecture 6; when stuck on why sharding saves memory, rewatch Lecture 8.

The handout also tells you what to read before the kernel section: the original FlashAttention paper for the online-softmax intuition, and Horace He’s “Making Deep Learning Go Brrr” for how GPUs execute PyTorch code.

Order of attack

Follow the handout’s order — it is deliberately staged. Profiling comes first because “we risk optimizing parts of the model that don’t account for significant time or memory.” Build the benchmarking script with command-line arguments from the start (the handout bolds this): you will rerun it with different model sizes, context lengths, precisions, and compile settings all assignment long. The model-size table (small through xl, plus a 10B config) with vocabulary 10,000, batch size 4, and context length 512 is your standard grid.

The kernel section is also staged for a reason. You first write a pure-PyTorch FlashAttention-2 forward pass as an autograd.Function — slow, but it becomes the reference you debug the Triton kernel against, tile by tile. Only then do you write the Triton forward kernel, then causal masking, then the backward pass (which, thanks to recomputation from the saved logsumexp, can be plain compiled PyTorch — the Triton backward is optional, mainly for leaderboard speed).

For the distributed sections, the handout’s advice is to debug locally with the Gloo backend on CPU, then benchmark with NCCL on GPU. The DDP progression is naïve per-parameter all-reduce, then a single flattened all-reduce, then overlapping communication with the backward pass via gradient hooks — benchmark each against the last on the standard 1-node, 2-GPU, xl-model setup.

What’s hard

The point totals tell you where the effort goes: FlashAttention-2 forward, optimizer state sharding, and FSDP are 15 points each. The Triton kernel is the classic wall — the online softmax bookkeeping (running max, running denominator, rescaling the output accumulator) has to be exactly right, block-pointer offsets are easy to get subtly wrong, and precision matters: on-chip accumulators in fp32, casts before writing out. FSDP is hard for a different reason — it’s about when things happen: all-gathering weights just in time, freeing them after use, reduce-scattering gradients, all without blocking compute. The handout recommends running the DDP, sharded-optimizer, and FSDP test suites about five times each to catch race conditions.

Resource-wise the handout is explicit: the distributed-communication benchmark uses up to 6 GPUs with each run under 5 minutes, and the standard distributed configuration elsewhere is 1 node with 2 GPUs.

The leaderboard

Assignment 2’s leaderboard times one full training step — forward, loss, backward, AdamW update — of an 8B-scale model (d_model 4096, 34 layers, context length 32,768, vocabulary 151,936) at batch size 2 in bf16 with causal masking, on two B200s. It is intentionally difficult to fit in memory. Your benchmarking run must finish within 10 minutes from an empty compile cache, submissions must pass the same tests as the regular implementation, and the naïve baseline to beat is 10 seconds. The handout’s menu of ideas: tune kernel tile sizes with the Triton autotuner, write a fused AdamW, fuse the LM head with the cross-entropy loss instead of materializing full logits, implement the Triton backward pass, skip fully-masked tiles in causal attention, and reach for activation checkpointing only if you need the memory.