Lab 01 - DDP Gradient Buckets

Overview

Lab 01: DDP Gradient Buckets

Annotated code reading lab. Running code is optional.

Related handbook section

DDP Gradient Buckets

This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Open handbook section Open starter file

Concept Goal

DDP Gradient Buckets

Read DDP backward as local gradients plus bucketed all-reduce.

Mental Model

Mechanism to keep in mind

`rank` owns a full model replica.
`bucket` groups gradients before communication.
`all_reduce` makes each replica see the averaged gradient.

rank local backward
  -> gradients become ready
  -> DDP bucket fills
  -> bucket all-reduce
  -> each rank applies the averaged gradient locally

Annotated Code Preview

Starter preview

Excerpt from code/lab-01-ddp-gradient-buckets/ddp_gradient_buckets.py. The linked starter file is the source of truth.

Open starter file

# DDP Gradient Buckets
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Read DDP backward as local gradients plus bucketed all-reduce.

world_size = 4
gradient_bucket = ["layer3.weight.grad", "layer3.bias.grad"]
# Each rank has a local bucket from its data shard.
local_bucket = f"bucket_from_rank_{0}"
synced_bucket = f"all_reduce_mean({local_bucket}, world_size={world_size})"
optimizer_step = f"apply({synced_bucket})"

# What to explain while reading:
# - rank owns a full model replica.
# - bucket groups gradients before communication.
# - all_reduce makes each replica see the averaged gradient.
#
# Common traps:
# - DDP syncs gradients, not input batches.
# - AllReduce happens during backward, not only after it.

Line-by-line Explanation

What each block is doing

Setup / contract: `rank` owns a full model replica.
Main transition: `bucket` groups gradients before communication.
Interview hook: `all_reduce` makes each replica see the averaged gradient.

What to Notice

Reading checkpoints

DDP does not shard parameters.
Dataloader sharding is the user's responsibility.
Bucket overlap hides some communication when configured well.

Common Misunderstandings

What this lab prevents

DDP syncs gradients, not input batches.
AllReduce happens during backward, not only after it.

Interview Explanation

How to say it out loud

Read DDP backward as local gradients plus bucketed all-reduce. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.

External intuition notes

Additional intuition

Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.

InfraLens