Training Frameworks & Profiling

A training framework is the combination of parallel strategy + memory optimization + observability. PyTorch native offers two paths: DDP (data parallelism; every rank holds a full model replica and runs one gradient All-Reduce after backprop — equivalent to ZeRO-0, solving compute scaling only, no memory win) and FSDP (Fully Sharded Data Parallel, the official implementation of ZeRO-3): every weight tensor is sharded by rank, parameters for each layer are All-Gathered before forward and discarded after use, and backward ends with a Reduce-Scatter of gradient shards, so optimizer states naturally live sharded — per-rank memory $16\Psi/N_d$ at ~1.5× DDP communication (3Φ vs 2Φ); layer-level sharding plus parameter prefetch hide communication inside compute for near-linear scaling. DeepSpeed offers ZeRO-1/2/3 with CPU/NVMe offload, mechanically equivalent to FSDP: pick FSDP for single-node multi-GPU with minimal code, DeepSpeed for cross-node, very large scale, offload, or its mature ecosystem. Profiling for bottlenecks (torch.profiler + Chrome trace): split a step into four parts — ① compute (CUDA kernel time and SM occupancy; too many tiny kernels → operator fusion / fewer syncs); ② communication (NCCL share of time; is the gradient All-Reduce overlapped with backprop — if not, tune bucket size / enable async); ③ IO (are DataLoader workers saturated, is disk/network the limiter — tune worker count and prefetch_factor); ④ idle (GPU idling on sync points or serial kernels). Debug order: check the CUDA time share first, then split forward/backward/optimizer, then drill down to the specific kernel or comm primitive.

interview favorites are the DDP/FSDP/DeepSpeed comparison and selection, FSDP sharding mechanics, and how to debug slow training or low GPU utilization.

bare DDP is helpless when one GPU cannot hold $16\Psi$; FSDP/ZeRO cut per-rank memory to $16\Psi/N_d$ so N GPUs can train a model N times bigger at only ~1.5× communication and near-linear scaling; the profiler turns 'training is slow' from gut feeling into a quantified time ledger (compute/comm/IO/idle shares), shrinking bottleneck hunting from hours to minutes.

Compare the memory budgets: DDP holds $16\Psi$ per rank (full replica, ZeRO-0), FSDP/ZeRO-3 hold $16\Psi/N_d$, DeepSpeed adds CPU/NVMe offload; state which solves compute scaling and which solves memory scaling.

Explain FSDP sharding: weight tensors sharded by rank → All-Gather the layer's params before forward → discard after use → Reduce-Scatter gradient shards in backward; why communication is ~1.5× DDP (3Φ vs 2Φ) and why scaling stays near-linear.

Selection: FSDP for single-node multi-GPU with minimal code; DeepSpeed for cross-node, offload, very large scale, or its ecosystem; what differs between FSDP's FULL_SHARD and SHARD_GRAD_OP.

Use torch.profiler to find bottlenecks: split a step into compute/comm/IO/idle; how to fix un-overlapped communication (bucket size, async gradients) and IO stalls (worker count, prefetch_factor, data format); give a debug order.

Worked scenario: why DDP cannot fit 70B on 8×80 GB ($16\Psi = 1120$ GB > 640 GB); FSDP/ZeRO-3 gives $1120/8 = 140$ GB per rank — still over 80 GB, so 16 ranks or offload ($2\Psi/N_d = 17.5$ GB on GPU) is required; give the full two-level accounting.