Distributed Training Interview Questions

Sample Answer

Yes, the linear scaling rule says you should multiply the base learning rate by the number of workers $k$, so $\eta_{\text{new}} = k \cdot \eta_{\text{base}}$, because the effective batch size grows by $k$ and you want each step to cover comparable ground in the loss landscape. However, this breaks down early in training when the loss surface has high curvature and large learning rates cause divergence. The standard mitigation, popularized by the Goyal et al. paper at Meta, is a gradual warmup: you linearly ramp the learning rate from $\eta_{\text{base}}$ to $k \cdot \eta_{\text{base}}$ over the first several epochs. Beyond roughly 8k effective batch size, even warmup may not suffice, and you may need LARS or LAMB optimizers that adapt per-layer.

Sample Answer

You could do AllReduce or a parameter server. AllReduce wins here because with homogeneous GPUs on a single node, you have symmetric high-bandwidth interconnects (like NVLink), and AllReduce algorithms such as ring-allreduce distribute the communication load evenly across all workers with bandwidth cost $\frac{2(N-1)}{N} \cdot M$ for message size $M$. A parameter server creates an asymmetric bottleneck: the server must receive gradients from all workers and broadcast updated parameters back, making it the throughput ceiling. Parameter servers are more useful in heterogeneous or wide-area settings where you need asynchronous updates, but for your single-node, homogeneous setup, AllReduce gives you near-optimal bandwidth utilization and keeps all GPUs as equal peers.

Sample Answer

Let's reason through this step by step. First, in synchronous data parallelism, every step has a communication phase where gradients are allreduced, and this cost grows with the number of workers: ring-allreduce takes $2(N-1)/N \cdot M$ time, so going from 16 to 32 workers increases the latency term even though per-worker bandwidth cost is roughly constant. Second, the computation per worker shrinks as you add GPUs (each processes a smaller micro-batch), which means the ratio of communication to computation worsens, reducing parallel efficiency. Third, with 32 GPUs your effective batch size doubles to a point where you may need to reduce the learning rate or add warmup, potentially requiring more total steps to converge, which offsets the wall-clock gain per step. You should profile the communication-to-computation ratio and consider gradient compression or overlapping allreduce with backward computation to recover the lost scaling efficiency.

Sample Answer

You could use ring-based AllReduce or tree-based (recursive halving/doubling) AllReduce. Tree-based wins here because small tensors are latency-bound, not bandwidth-bound. Ring AllReduce requires $2(n-1)$ sequential communication steps, so latency scales as $O(n \cdot \alpha)$ where $\alpha$ is per-message latency. A tree-based algorithm like recursive halving-doubling completes in $O(\log n)$ steps, drastically reducing latency for small messages at the cost of slightly higher bandwidth usage. In practice, libraries like NCCL automatically select the algorithm based on message size, but you should know this tradeoff cold.

Sample Answer

Let's reason through this step by step. In asynchronous SGD, each worker computes gradients on a potentially stale copy of the parameters, so by the time a gradient is applied, the model may have moved several steps ahead. This staleness introduces implicit noise that biases the effective gradient, which can prevent convergence to sharp minima and degrade final accuracy. The staleness problem worsens as you scale to more workers because the expected delay $\tau$ grows. To fix this without going fully synchronous, you can apply staleness-aware corrections: scale each gradient by $\frac{1}{1 + \tau}$ or use bounded staleness (as in SSP, Stale Synchronous Parallel) where you allow workers to drift at most $s$ steps apart before forcing a barrier. This gives you most of the throughput benefit of async while bounding the convergence degradation.

Sample Answer

Let's reason through this step by step. You start with input $X$ and a weight matrix $A$ split column-wise into $[A_1, A_2]$ across two GPUs, so GPU $i$ computes $Y_i = XA_i$ independently since each GPU has the full $X$. After the GeLU activation, each GPU holds $\text{GeLU}(Y_i)$, which is already the correct partial input for a row-wise split of the second weight matrix $B$, where GPU $i$ holds $B_i$. Each GPU computes $Z_i = \text{GeLU}(Y_i) B_i$, and only then do you need an all-reduce to sum $Z = Z_1 + Z_2$, because $XAB = XA_1B_1 + XA_2B_2$. The key insight is that column-then-row partitioning lets the intermediate activation stay local, so you need only one all-reduce per MLP block instead of two.

Sample Answer

This question is checking whether you can map the abstract idea of tensor parallelism to the concrete mechanics of multi-head attention. You assign $h/8$ attention heads to each of the 8 GPUs, where $h$ is the total number of heads, so each GPU independently computes its subset of Q, K, V projections and attention outputs. After the per-head computation, each GPU holds a partial result of the output projection (split row-wise), so you perform a single all-reduce to sum these partials, costing $2(p-1)/p \cdot b \cdot s \cdot d$ bytes in the ring all-reduce formulation, where $b$ is batch size, $s$ is sequence length, $d$ is hidden dimension, and $p=8$. There is also an all-reduce (or the equivalent conjugate: an all-gather before and reduce-scatter after) in the subsequent MLP block, giving two all-reduces total per transformer layer.

Sample Answer

This question is checking whether you can distinguish the memory and convergence implications of asynchronous versus synchronous pipeline schedules. In GPipe, all micro-batch forward passes complete before any backward pass starts, so every stage must store activations for all $m$ micro-batches simultaneously, giving $O(m)$ peak activation memory per stage. PipeDream's 1F1B interleaves forward and backward passes so that each stage holds activations for at most $p$ micro-batches (where $p$ is the pipeline depth), reducing peak memory to $O(p)$. The tradeoff in the original PipeDream (not PipeDream-2BW or PipeDream Flush) is gradient staleness: weight updates can be applied using gradients computed on older versions of the model parameters, which can hurt convergence or require weight stashing to maintain correctness.

Sample Answer

The standard move is to split layers evenly across stages. But here, stage imbalance matters because the slowest stage determines the throughput of the entire pipeline, and any imbalance directly inflates the bubble. You should profile each layer's compute and memory cost, then assign fewer transformer layers to the first and last stages to compensate for the extra embedding, output projection, and loss computation overhead. For example, you might assign 4 layers to stage 0, 7 layers each to stages 1 through 6, and 4 layers to stage 7. Some frameworks like Megatron-LM support specifying a custom layer-to-stage mapping, and you can iteratively tune this partitioning by measuring per-stage iteration time and rebalancing until the maximum stage latency is minimized.

Sample Answer

The standard move is to adopt ZeRO Stage 3 whenever per-GPU memory is still too tight, since it shards parameters on top of the optimizer states and gradients already sharded in Stage 2. But here, communication volume matters because Stage 3 requires an all-gather of parameters in the forward pass and again in the backward pass, roughly $1.5\times$ the communication of Stage 2. You should measure whether your interconnect bandwidth (NVLink vs PCIe) can absorb this without becoming the bottleneck. If your 8 GPUs are on a single node with NVLink, Stage 3 is usually fine, but across nodes on slower interconnects you may lose more in throughput than you gain in memory, and techniques like activation checkpointing or reducing batch size under Stage 2 could be a better path.

Sample Answer

Get this wrong in production and your model silently stops learning while still consuming full compute budget. The right call is to re-enable dynamic loss scaling. In FP16, the smallest representable normal value is roughly $6 \times 10^{-5}$, so many gradient values that are small but meaningful in FP32 underflow to zero in FP16. Loss scaling multiplies the loss by a large factor (e.g., $2^{16}$) before the backward pass to shift gradients into FP16's representable range, then divides them back down before the optimizer step. Dynamic loss scaling adjusts this factor automatically: it increases the scale when no overflow is detected and halves it when infs or nans appear, keeping gradients in the sweet spot without manual tuning.

Sample Answer

Get this wrong in production and you either lose hours of compute on failure recovery or tank your training throughput with too-frequent blocking checkpoints. The right call is to use asynchronous checkpointing where you snapshot model state to CPU memory or a staging buffer and flush to persistent storage in the background, so GPU training continues unblocked. For choosing interval $N$, you want to minimize expected wasted compute: if mean time between failures is $T_{\text{MTBF}}$ and checkpoint cost is $C$, the optimal interval is approximately $N \approx \sqrt{2 \cdot C \cdot T_{\text{MTBF}}}$ steps. In practice, teams at this scale also keep the last $k$ checkpoints and validate checkpoint integrity with checksums to avoid silent corruption.

Sample Answer

Full restart sounds reasonable but breaks under high failure rates on large clusters, where you could spend more time restarting than training. Elastic training (shrinking the worker pool) doesn't work well when you rely on a fixed global batch size tied to a carefully tuned learning rate schedule, because removing workers changes the effective batch size per step. That leaves a hybrid: use elastic training for short-lived failures where you can redistribute data shards and adjust the per-worker micro-batch size to maintain the global batch size, and fall back to full restart only when the failure is persistent or affects model-parallel ranks that cannot be trivially reassigned. You should also consider that frameworks like TorchElastic handle data-parallel elasticity well, but tensor/pipeline parallel configurations typically require a full restart since rank assignments are tightly coupled to model partitions.

Sample Answer

Most candidates default to a simple bin-packing scheduler that fills nodes greedily, but that fails here because it ignores network topology and can scatter a single job's GPUs across different spine switches, dramatically increasing all-reduce latency. You want topology-aware placement that co-locates a job's ranks within the same rack or pod to keep collective communication on low-latency, high-bandwidth leaf switch links. For pipeline-parallel jobs, place consecutive pipeline stages on adjacent nodes to minimize inter-stage latency, while keeping tensor-parallel groups within NVSwitch domains. A good scheduler also reserves headroom for job elasticity and preemption, and uses gang scheduling so all ranks of a job start simultaneously rather than trickling in and wasting allocated resources.