Model Training Infrastructure and Distributed Training Questions
Scaling model training across hardware and time. Covers GPU/accelerator considerations, data and model parallelism, distributed and large-scale training, experiment tracking and training infrastructure, and the training-versus-inference compute tradeoff. Focuses on the systems and resource decisions that make large-model training feasible.
Compare PyTorch DistributedDataParallel (DDP), Horovod, and DeepSpeed for distributed training. For each, explain typical use cases, integration complexity, memory and communication trade-offs, and support for model parallelism or optimization techniques like ZeRO.
Sample Answer
Direct answer
PyTorch DDP, Horovod, and DeepSpeed all support distributed data-parallel training but differ in scope and integration complexity: DDP is PyTorch-native and the default choice for most PyTorch distributed training, Horovod is a framework-agnostic library (works with PyTorch, TensorFlow, and others) originally built to bring easy, efficient distributed training to frameworks that lacked strong native support, and DeepSpeed goes beyond basic data parallelism to provide advanced memory-optimization features (ZeRO) and additional parallelism strategies (pipeline, and integration with tensor parallelism) as a more comprehensive large-model training toolkit.
Structured elaboration
- PyTorch DDP: native to PyTorch, tightly integrated with the rest of the ecosystem (autograd hooks,
torch.compile, nativeDTensor-based tensor parallelism in recent versions), the natural default for standard PyTorch data-parallel training; less feature-rich out of the box for advanced memory optimization (ZeRO-3-equivalent full sharding requires PyTorch's separate FSDP, not DDP itself) or non-data-parallel parallelism strategies. - Horovod: framework-agnostic (a real differentiator: works across PyTorch, TensorFlow, and other frameworks with a consistent API), historically valuable for teams needing efficient distributed training in frameworks (like older TensorFlow versions) that lacked strong native distributed support; integration complexity is moderate (wrapping the optimizer and adding a few Horovod-specific calls), and its relevance has diminished somewhat as native framework support (like PyTorch DDP/FSDP) has matured, though it remains relevant for genuinely multi-framework organizations wanting one consistent distributed-training approach. On communication and capability: Horovod's original differentiator was implementing an efficient ring-allreduce for gradient synchronization at a time when TensorFlow and PyTorch lacked one natively (PyTorch's own DDP now uses a comparable bucketed NCCL all-reduce, closing much of this gap); Horovod provides no model-parallelism or ZeRO-style memory-sharding support at all, it is purely a data-parallel scaling library, so for a model too large to fit via data parallelism alone, Horovod offers nothing DeepSpeed's memory-sharding features provide.
- DeepSpeed: goes well beyond basic data parallelism, providing ZeRO's memory-sharding stages, pipeline parallelism, and various other large-model-training optimizations (as discussed extensively elsewhere in this topic) as an integrated toolkit; integration complexity is higher than plain DDP (a separate config file/engine wrapping the training loop) but delivers capabilities (fitting far larger models, more aggressive memory optimization) that DDP alone doesn't provide.
- Choosing between them: DDP (or DDP plus FSDP for memory-constrained large models) is the natural default for a PyTorch-only team without extreme memory constraints; Horovod is worth considering specifically for multi-framework organizations; DeepSpeed is worth its added integration complexity specifically when training models large enough to need its advanced memory-sharding and parallelism features beyond what DDP/FSDP alone provide.
Worked example
A team training a 1B-parameter model that comfortably fits with standard data parallelism uses plain PyTorch DDP, the simplest, most tightly-integrated option for their actual needs; a different team training a 70B-parameter model requiring aggressive memory sharding, pipeline parallelism, and CPU/NVMe offload adopts DeepSpeed specifically for those advanced capabilities, accepting its added configuration complexity as a worthwhile trade for training a model that wouldn't otherwise fit on their available hardware.
Trade-offs & pitfalls
A common mistake is reaching for DeepSpeed's full complexity (or Horovod, for a single-framework team with no multi-framework need) when plain DDP would suffice for the actual model scale and memory requirements; matching the tool's complexity to the actual need (rather than defaulting to the most feature-rich option) keeps the training pipeline simpler and easier to debug.
Explain the differences between NVLink (and NVSwitch) and PCIe for GPU-to-GPU communication. Discuss topology, raw bandwidth, latency, and why NVLink/NVSwitch can improve multi-GPU scalability for synchronous training. Give an example of a situation where PCIe becomes the bottleneck and how NVLink changes your parallelism choices.
Sample Answer
Direct answer
NVLink (and its switch fabric, NVSwitch) is a high-bandwidth, low-latency interconnect purpose-built for direct GPU-to-GPU communication within a node, offering substantially higher bandwidth than PCIe, which is a general-purpose bus shared by many devices (GPUs, storage, network cards) and not specifically optimized for GPU-to-GPU traffic.
Structured elaboration
- Topology: NVLink connects GPUs directly to each other (or, with NVSwitch, through a dedicated switch fabric enabling any-to-any GPU connectivity within a node at full bandwidth); PCIe connects each GPU to the host system's PCIe root complex, meaning GPU-to-GPU traffic over PCIe alone must traverse through that shared, general-purpose bus, competing with other PCIe traffic and paying the overhead of going through the host's PCIe topology rather than a direct GPU-to-GPU link.
- Raw bandwidth: NVLink generations offer substantially higher per-link bandwidth than PCIe generations available at the same time (often several times higher aggregate bandwidth per GPU), which directly matters for any collective communication pattern (all-reduce, tensor-parallel activation exchange) that needs to move large tensors between GPUs frequently.
- Latency: NVLink's direct point-to-point (or switch-fabric-mediated) connection also offers lower latency than routing through the host's PCIe complex, which matters especially for smaller, more frequent messages where fixed per-hop latency (not just bandwidth) dominates communication time.
- Why this matters for multi-GPU workloads specifically: workloads with high inter-GPU communication intensity relative to their compute (tensor parallelism, which needs an all-reduce or similar collective within every layer, is the extreme case) benefit disproportionately from NVLink/NVSwitch's advantage over PCIe, which is exactly why tensor parallelism is almost always kept strictly within a single NVLink-connected node rather than spanning nodes connected only by PCIe-to-network paths.
Worked example
An 8-GPU node with full NVSwitch connectivity lets any pair of the 8 GPUs communicate at full NVLink bandwidth simultaneously (a genuine any-to-any high-bandwidth fabric); the same 8 GPUs communicating only over PCIe (in a system without NVLink) would have their inter-GPU traffic bottlenecked by the shared PCIe root complex's aggregate bandwidth, which is both lower in raw capacity and shared across all 8 GPUs' traffic simultaneously rather than each pair getting a dedicated high-bandwidth path.
Trade-offs & pitfalls
Not every multi-GPU system has NVLink/NVSwitch (it's typically found on higher-end data-center GPU systems, not consumer or lower-tier configurations), so PCIe-only multi-GPU communication is still common in less specialized hardware; for workloads with genuinely low inter-GPU communication intensity relative to compute (data parallelism with infrequent, well-overlapped gradient synchronization, for instance), the PCIe-versus-NVLink distinction matters much less than it does for communication-intensive workloads like tensor parallelism.
You're planning on-prem GPU procurement for large-scale training. Draft a plan covering GPU model selection, rack and rack-power sizing (kW per rack), cooling requirements, UPS, networking (RDMA, 100GbE), procurement timelines, vendor support, capacity planning for target utilization, and how these physical choices affect training throughput, single-job latency, and total cost of ownership.
Sample Answer
Direct answer
Planning on-prem GPU procurement for large-scale training requires a plan spanning GPU model selection (matched to the actual workload's compute/memory needs), physical infrastructure sizing (rack space and power, since dense GPU nodes draw far more power per rack than typical enterprise servers), and cooling requirements sufficient for that power density, since underestimating any one of these three typically becomes the binding constraint regardless of how well the other two are planned.
Structured elaboration
- GPU model selection: match the specific GPU generation/model to the actual workload requirements (memory capacity per GPU for the target model sizes, precision support needed like native bf16/fp8, and interconnect capability like NVLink generation), rather than defaulting to "the newest/most powerful available," since procurement lead time and cost scale with how cutting-edge the chosen hardware is, and a slightly older generation may be substantially cheaper and faster to procure while still meeting the actual workload's needs.
- Rack and rack-power sizing: dense GPU servers (e.g. 8-GPU nodes with high-end data-center GPUs) can draw substantially more power per rack unit than typical enterprise compute hardware, often requiring specific high-density rack power provisioning (many kW per rack) well beyond what a typical data center's existing rack power allocation provides; this needs to be sized explicitly against the specific GPU model and node density chosen, not assumed to fit within existing, non-GPU-optimized rack power budgets.
- Cooling requirements: power drawn becomes heat that must be removed; at the power densities modern GPU racks require, some data centers need liquid cooling or other high-density cooling solutions beyond standard air cooling, which is both a cost and a facility-capability question that needs to be resolved (does the target facility support this, or does a facility upgrade/different facility choice become necessary) before hardware even arrives.
- Networking infrastructure: beyond the GPUs themselves, the inter-node networking fabric (InfiniBand or high-speed Ethernet, matched to the target training scale's communication needs) needs its own procurement and installation planning, often on a similarly long lead time to the GPUs themselves, and needs to be sized for the target cluster's actual communication pattern (all-reduce-heavy synchronous training needs particularly high bandwidth and low latency between nodes). Concretely this typically means RDMA-capable fabric (InfiniBand, or RoCE over 100GbE or faster Ethernet) rather than standard TCP/IP networking, since RDMA's ability to transfer data between GPUs across nodes without routing through the host CPU is what makes multi-node all-reduce fast enough not to become the dominant bottleneck at scale.
- UPS (uninterruptible power supply): sized to bridge the gap until backup generators engage, or to allow a controlled checkpoint-and-shutdown of in-flight training jobs, since an ungraceful power loss mid-training risks losing the last unsaved checkpoint's worth of progress and can corrupt a checkpoint file that was mid-write at the moment power dropped; UPS capacity needs to be sized against the same power-density numbers used for rack and cooling planning, not treated as a separate, smaller-scale concern.
- Vendor support: negotiate support/SLA contracts with a bounded hardware-replacement turnaround time, since in a large synchronous training job a single failed GPU or failed NIC can stall the entire job until it's replaced or worked around, and on-site spare-parts agreements for the most failure-prone components meaningfully reduce that downtime risk compared to a standard mail-in RMA process.
- Capacity planning for target utilization: provision against a realistic target utilization (commonly in the 70-85% range) rather than assuming the cluster runs at 100% of nameplate capacity continuously, since job-scheduling gaps, planned maintenance windows, and failed-node downtime all reduce achievable utilization below the theoretical maximum; sizing the cluster as if it will hit 100% utilization leads to a plan that can't actually deliver the training throughput it was sized for.
- Lead time and phased procurement: GPU hardware (especially cutting-edge models) can have long procurement lead times (many months, sometimes longer during periods of high industry demand), which needs to be planned well ahead of the target training start date, potentially with a phased procurement approach (securing an initial tranche of capacity while the remainder is still in the supply pipeline) rather than assuming all hardware arrives simultaneously on a fixed date.
Worked example
Planning a 64-GPU on-prem cluster (8 nodes of 8 GPUs each): GPU model selection settles on a specific generation balancing memory-per-GPU against cost and availability lead time; rack power sizing calculates each 8-GPU node's power draw (summing GPU TDP, host CPU/memory/storage draw, and typical power-supply overhead) and provisions rack power circuits with headroom above that calculated draw; cooling capacity is validated against the facility's existing cooling infrastructure, with a liquid-cooling retrofit budgeted if the calculated heat output exceeds what existing air cooling can handle; and procurement is initiated with enough lead time (informed by current vendor-quoted lead times, not optimistic assumptions) before the planned training start date, with network fabric procurement running on a parallel timeline.
Trade-offs & pitfalls
The most common on-prem GPU procurement planning mistake is treating GPU hardware acquisition as the only real constraint and underestimating power/cooling infrastructure needs, which can become the actual bottleneck (a facility physically unable to support the power density of the newly-arrived GPU hardware) well after the (harder-to-quickly-fix) hardware procurement itself is complete; power and cooling capacity should be validated and, if necessary, upgraded on a timeline that keeps pace with hardware procurement, not treated as an afterthought to be solved once hardware arrives. Each of these choices also has a direct, traceable effect on throughput, latency, and total cost of ownership: insufficient rack power or cooling headroom forces GPUs to thermal-throttle, directly reducing achieved training throughput below the hardware's rated capability regardless of how well everything else was planned; insufficient network bandwidth caps achievable scaling efficiency as GPU count grows (per the earlier scaling-efficiency discussion), which shows up as worse single-job latency at scale; and TCO must account for ongoing power, cooling, and support-contract operating costs, not just GPU sticker price, since a cheaper GPU with worse power efficiency or a weaker support contract can end up costing more over its operating lifetime than a pricier, better-supported alternative.
You need to implement a custom fused transformer-attention kernel to leverage Tensor Cores for better throughput. Explain design choices and implementation steps either using CUDA WMMA APIs or Triton: data layout, tile sizes, alignment constraints, memory staging (shared memory), avoiding bank conflicts, and validation for numerical correctness and performance regression testing.
Sample Answer
Direct answer
Implementing a custom fused transformer-attention kernel to leverage Tensor Cores means combining several steps of the attention computation (the QK^T matrix multiply, the softmax, and the subsequent multiply by V) into a single GPU kernel that keeps intermediate results in fast on-chip memory rather than round-tripping through HBM between each step, while structuring the matrix multiplications specifically to match Tensor Cores' preferred tile sizes and precision.
Structured elaboration
- Why fusion matters here specifically: a naive, unfused attention implementation computes QK^T (writing the full attention score matrix to HBM), then reads it back for softmax (writing the result back to HBM again), then reads it back again for the final multiply by V; for long sequences, this attention-score matrix is large (quadratic in sequence length) and this round-tripping is a substantial, avoidable HBM bandwidth cost. Fusing these steps into one kernel (as FlashAttention-style implementations do) keeps intermediate results in fast shared memory/registers, computing the whole attention operation for a tile of the sequence without ever materializing the full attention-score matrix in HBM.
- Tiling for Tensor Cores: the QK^T and attention-weights-times-V matrix multiplications need to be structured as tiled operations matching Tensor Cores' preferred small-matrix-tile granularity (Tensor Cores operate on fixed small tile sizes, e.g. 16x16 or similar, depending on generation and precision), requiring the kernel to explicitly manage how the sequence and head dimensions are partitioned into Tensor-Core-sized tiles, rather than treating the operation as one large, unstructured matrix multiply.
- Precision choices: using fp16/bf16 for the matrix multiplications (to engage Tensor Cores) while keeping the softmax's numerically sensitive operations (the exponentiation and normalization) in fp32 internally, converting back to the lower precision only for the final output, balancing Tensor Core throughput against the numerical stability softmax specifically needs.
- Online softmax (a key algorithmic trick): since the full attention-score matrix is never materialized, softmax needs to be computed incrementally (an "online" or "streaming" softmax) as tiles of the score matrix are produced, maintaining a running maximum and running sum that get corrected as new, potentially-larger values are seen, rather than the standard softmax algorithm's assumption that the full row of values is available at once.
- Alignment constraints: tile dimensions and shared-memory buffer strides need to be sized and aligned to the hardware's requirements (e.g. global-memory accesses aligned to 128-byte boundaries for efficient coalescing, and matrix dimensions padded to the Tensor Core's tile-size multiple, as discussed in the Tensor Core tile-shapes topic) so that both the global-memory loads staging data into shared memory and the Tensor Core operations themselves run on their fastest supported path rather than an unaligned, slower fallback.
- Avoiding shared-memory bank conflicts: shared memory is physically divided into banks, and when multiple threads in the same warp read or write different addresses that happen to map to the same bank simultaneously, those accesses serialize instead of completing in parallel, a "bank conflict" that silently degrades the kernel's effective shared-memory bandwidth; staging Q/K/V tiles into shared memory needs a layout that avoids this, typically via padding (adding an extra unused column so a tile's row stride no longer aligns with the bank count) or an explicit swizzled/permuted addressing scheme, a technique used deliberately in FlashAttention-style kernels to keep every thread in a warp hitting a distinct bank.
Worked example
Design choices in code (framework-level, since a full custom CUDA kernel is out of scope for this format): implement attention using PyTorch's torch.nn.functional.scaled_dot_product_attention with a Flash-Attention-style backend selected (which internally performs exactly this fusion and Tensor-Core-tiled computation), or, if building a fully custom kernel, use a library like Triton to write the fused kernel at a higher level of abstraction than raw CUDA while still controlling tile sizes and the online-softmax algorithm explicitly; either path avoids materializing the full O(sequence_length^2) attention-score matrix in HBM, which is the core design win regardless of implementation level.
Trade-offs & pitfalls
Writing a genuinely optimal custom fused attention kernel from scratch (rather than using an existing, heavily-optimized implementation like FlashAttention) is a substantial engineering undertaking, requiring careful tuning of tile sizes for the specific hardware generation and getting the online-softmax numerics exactly right; for most teams, the practical answer is using an existing, well-validated fused-attention implementation rather than reimplementing this from scratch, reserving custom kernel work for cases with genuinely novel attention variants not covered by existing libraries. Validating a custom kernel also needs its own explicit process: check numerical correctness by comparing the kernel's output against a reference eager-mode (unfused) attention implementation using a tolerance-based comparison (e.g. torch.allclose with an atol/rtol appropriate to the reduced precision in use, since fp16/bf16 accumulation will not match fp32 bit-for-bit) across a range of sequence lengths, batch sizes, and edge cases (very short sequences, causal/masked attention, sequences not evenly divisible by the tile size); and set up a performance regression test that benchmarks the kernel's measured throughput/latency against a stored baseline on every change, failing the check if performance regresses beyond a defined threshold, since a kernel change can stay numerically correct while silently becoming slower.
Implement a simple Python function that, given GPU peak FLOPS (TFLOPS), memory bandwidth (GB/s), per-sample memory traffic (bytes), and batch size, estimates whether a training step is compute-bound or memory-bound using a heuristic. Define the compute and memory time estimates and use a threshold (e.g., compute_time > 2 * memory_time => compute-bound). Provide example outputs for different configurations.
Sample Answer
Direct answer
Classifying a training step as compute-bound or memory-bandwidth-bound from hardware specs (peak FLOPS, memory bandwidth) and the workload's per-sample memory traffic and batch size means comparing the time each resource would take if it were the sole bottleneck: whichever takes longer is the actual bottleneck (the step's real time is at least that long, assuming imperfect overlap between compute and memory access).
Structured elaboration
- Compute time estimate: given the workload's FLOPs per sample (not provided directly here, but implied by whatever operation is being analyzed) and the GPU's peak FLOPS, compute_time = (FLOPs_per_sample * batch_size) / peak_FLOPS.
- Memory time estimate: given per-sample memory traffic (bytes read/written per sample) and the GPU's memory bandwidth, memory_time = (bytes_per_sample * batch_size) / memory_bandwidth.
- Classification: if memory_time > compute_time, the operation is memory-bandwidth-bound (the GPU's compute units would sit idle waiting for data even if compute were instantaneous); if compute_time > memory_time, it's compute-bound (the memory system could keep up with idealized data delivery, and compute is the limiting factor).
- Arithmetic intensity as the underlying concept: the ratio of FLOPs to bytes moved (FLOPs per byte) is what actually determines this classification, independent of batch size (since both FLOPs and bytes scale linearly with batch size for a simple, non-batched-interaction operation, batch size cancels out of the ratio, though it still affects each individual estimate's absolute magnitude).
Worked example
def classify_compute_or_memory_bound(peak_tflops: float, memory_bandwidth_gbs: float,
flops_per_sample: float, bytes_per_sample: float,
batch_size: int) -> dict:
peak_flops = peak_tflops * 1e12
bandwidth = memory_bandwidth_gbs * 1e9
compute_time = (flops_per_sample * batch_size) / peak_flops
memory_time = (bytes_per_sample * batch_size) / bandwidth
classification = "memory-bound" if memory_time > compute_time else "compute-bound"
return {
"compute_time_s": compute_time, "memory_time_s": memory_time,
"classification": classification,
"arithmetic_intensity": flops_per_sample / bytes_per_sample,
}
# Verification: an elementwise operation (low arithmetic intensity) vs a matmul-like operation (high)
elementwise = classify_compute_or_memory_bound(
peak_tflops=312, memory_bandwidth_gbs=2000,
flops_per_sample=1024, bytes_per_sample=1024*4, # 1 FLOP per byte roughly, read+write a value
batch_size=1_000_000,
)
matmul_like = classify_compute_or_memory_bound(
peak_tflops=312, memory_bandwidth_gbs=2000,
flops_per_sample=1024*1024, bytes_per_sample=1024*4, # much higher FLOPs per byte moved
batch_size=1_000_000,
)
assert elementwise["classification"] == "memory-bound"
assert matmul_like["classification"] == "compute-bound"
Running this confirms the elementwise-style operation (low FLOPs relative to bytes moved) classifies as memory-bound, while the matmul-like operation (much higher FLOPs per byte, from reusing loaded data across many multiply-accumulate operations) classifies as compute-bound, matching the general expectation that elementwise operations tend to be memory-bound and dense matrix multiplies tend to be compute-bound on modern hardware.
Trade-offs & pitfalls
This model assumes compute and memory access happen sequentially (not overlapped) for the estimate; real GPU execution often overlaps compute and memory access to some degree (especially with good kernel design and enough independent work to hide memory latency behind ongoing compute), so real measured time is often somewhat better than this simple additive/max-of-the-two model predicts, meaning this classification is a useful diagnostic heuristic, not a precise performance predictor.
Unlock Full Question Bank
Get access to all Model Training Infrastructure and Distributed Training interview questions and detailed answers.
Sign in to ContinueJoin thousands of developers preparing for their dream job.