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.
Design an autoscaling policy for an ML training cluster that uses Kubernetes or Slurm: describe triggers, scale step sizes, cooldowns, safety buffers, how to incorporate Spot instance availability and preemption risk, and how to avoid oscillations or under-provisioning for large gang-scheduled jobs.
Sample Answer
Direct answer
An autoscaling policy for a Kubernetes- or Slurm-managed ML training cluster needs a trigger metric that reflects actual queued demand, conservative scale-step sizes and cooldowns to avoid thrashing on GPU nodes, safety buffers for spot/preemption risk, and a specific accommodation for gang-scheduled jobs, which can be starved by an autoscaler that scales up in small increments that never reach the full all-at-once capacity such a job actually needs.
Structured elaboration
- Trigger metric: queue depth (pending jobs or pending GPU-hours) and time-in-queue are better autoscaling triggers than raw utilization.
- Scale step sizes: scale up in larger, coarser steps (e.g. a full node of 8 GPUs at a time).
- Cooldowns: a meaningfully longer cooldown than typical CPU autoscaling, since GPU node boot/join time is commonly several minutes.
- Safety buffers and spot instance availability/preemption risk: on Kubernetes, this typically means running a mixed node-pool (on-demand plus spot/preemptible node pools, using taints/tolerations or a priority-based Cluster Autoscaler expander to prefer spot capacity first) with a fallback to on-demand capacity when spot is unavailable or a preemption notice fires; on Slurm, this means configuring a preemptible partition/QOS with its own priority and a documented requeue behavior (Slurm's
--requeueor an equivalent hook) so a preempted job's allocation is reclaimed and the job automatically resubmits rather than silently failing. Both need the autoscaler to treat spot/preemptible capacity as inherently less reliable, budgeting a safety margin of on-demand capacity so a burst of simultaneous preemptions doesn't stall the whole queue. - Avoiding under-provisioning for large gang-scheduled jobs: a gang-scheduled job (e.g. one requiring all 8 nodes of a 64-GPU synchronous training run simultaneously) cannot start with a partial allocation; if the autoscaler scales up in small increments reactively (adding one node, waiting to see if the job still hasn't started, adding another), it can take many scale-up cycles to reach the full requirement, while the job sits gang-scheduled-but-unstarted the whole time, effectively starving it far longer than a smaller job would ever wait. The fix is for the autoscaler to be gang-scheduling-aware: on Kubernetes, this means integrating with a gang-scheduling-aware autoscaler/scheduler pairing (e.g. Cluster Autoscaler driven by a scheduler plugin like Kueue or Volcano that exposes the FULL pending resource requirement of a gang, not just individual pod requests, so the autoscaler provisions the whole batch of nodes in one scale-up decision rather than one pod's worth at a time); on Slurm, this means the scheduler's own backfill/reservation logic (Slurm natively supports this via job step dependencies and resource reservations) holding capacity as it accumulates specifically for the pending gang job rather than backfilling it away to smaller jobs indefinitely.
- Scale-down caution: scale down more conservatively than scale-up, since removing a node that's mid-job (or about to be needed again soon, including for an accumulating gang-scheduled job) is more disruptive than a slightly delayed scale-down.
Worked example
On a Kubernetes cluster using Kueue for gang-aware scheduling: a policy scaling up by one full 8-GPU node whenever queue depth exceeds 16 GPU-equivalent-hours sustained for 5 minutes, with a 10-minute cooldown; when an 8-node (64-GPU) synchronous job is submitted, Kueue reports the FULL 64-GPU pending requirement to Cluster Autoscaler in one signal (rather than 8 separate 8-GPU pod requests trickling in), so the autoscaler provisions all 8 nodes in one coordinated scale-up rather than adding them one at a time across 8 separate 10-minute-cooldown cycles, which would otherwise delay the job's start by roughly an hour.
Trade-offs & pitfalls
A common mistake is directly porting a CPU-service autoscaling policy onto a GPU training cluster without adjusting for provisioning time and granularity. A second, more subtle mistake specific to this setting is treating every pending job as independently schedulable, which is exactly what starves large gang-scheduled jobs behind a scheduler/autoscaler pairing (vanilla Kubernetes Cluster Autoscaler with default pod-level scheduling, or a Slurm configuration without reservation/backfill tuned for large jobs) that isn't gang-scheduling-aware.
Explain common causes of GPU out-of-memory (OOM) errors and VRAM fragmentation in long-running training jobs. Provide practical mitigations, and describe how you would diagnose fragmentation specifically (as opposed to a genuine capacity shortfall) before applying a fix.
Sample Answer
Direct answer
GPU out-of-memory (OOM) errors and VRAM fragmentation in long-running training jobs happen either because peak memory demand genuinely exceeds the GPU's capacity, or because memory is available in aggregate but fragmented into pieces too small to satisfy a single allocation request; the two causes need different diagnosis and different fixes.
Structured elaboration
- Genuine capacity shortfall: the model's weights, gradients, optimizer state, and activations for the configured batch size genuinely exceed available GPU memory; mitigations are the standard memory-saving toolbox (smaller batch size, mixed precision, gradient checkpointing, sharding) discussed throughout this topic.
- Fragmentation: memory gets allocated and freed repeatedly over the course of training (different-sized activation tensors from step to step, especially with variable-length inputs), and over time the memory allocator's free space becomes fragmented into many small, non-contiguous blocks; even if the SUM of free memory would be enough for a new allocation, no single free block is large enough, causing an OOM despite aggregate free memory being sufficient. This typically manifests as OOM errors appearing only after a job has been running for a while (as fragmentation accumulates), rather than immediately at startup.
- Diagnosing which is happening: check the memory allocator's reported statistics (e.g. PyTorch's
torch.cuda.memory_summary()) for allocated versus reserved memory; a large gap between reserved and allocated, combined with an OOM despite that gap, is a strong fragmentation signal, distinct from a case where allocated memory itself is simply climbing toward the device's total capacity (a genuine capacity issue). - Diagnosing fragmentation further: reproducing the OOM with a fixed input shape (eliminating variable-length-input-driven allocation-pattern variability) and observing whether it still occurs helps confirm whether variable-size allocations are the fragmentation driver.
Fixes to address fragmentation specifically
Setting allocator configuration options that reduce fragmentation (e.g. PyTorch's PYTORCH_CUDA_ALLOC_CONF with an expandable-segments or max-split-size setting), preallocating buffers for known-variable-size tensors to a fixed maximum size rather than letting the allocator repeatedly allocate and free differently-sized tensors, and periodically calling the framework's cache-clearing function at safe points (though this has its own overhead and shouldn't be called every step) to consolidate free memory.
Trade-offs & pitfalls
Treating every OOM as a capacity problem (reflexively reducing batch size) without first checking whether it's actually fragmentation wastes the opportunity to fix the real root cause with an allocator configuration change instead, which preserves the batch size (and therefore training efficiency) that a capacity-focused fix would have unnecessarily sacrificed.
Explain best practices for creating containerized, reproducible training environments. What should be included in the container image (OS libs, CUDA/cuDNN versions, pip/conda packages), and how do you ensure experiments are reproducible across image versions and host kernels?
Sample Answer
Direct answer
A reproducible containerized training environment pins every layer of the software stack INSIDE the image (OS/system libraries, CUDA/cuDNN, framework and dependency versions, code) as a versioned, immutable image, and additionally accounts for the parts of the stack that live OUTSIDE the container on the host, principally the host's NVIDIA GPU driver and kernel version, which containerization alone does not pin.
Structured elaboration
- Base OS and system libraries: pin the base image to a specific tag, including specific versions of system-level libraries.
- CUDA/cuDNN versions: pin exact CUDA and cuDNN versions matching what the framework was validated against.
- Framework and dependency versions: pin exact versions for the framework and every dependency in a lockfile-style manifest (pip's
requirements.txtwith==pins, or a conda lockfile for conda-managed environments). - Code and configuration: bake the training code (or a specific commit hash reference) into the image, or record the commit hash alongside the image tag.
- Host kernel and GPU driver, the part containerization does NOT pin: for GPU workloads, the NVIDIA driver is installed on the HOST machine, not inside the container; the container only ships the CUDA TOOLKIT (compiler, runtime libraries), which must be compatible with, but is architecturally separate from, whatever driver version the host happens to have installed. Two hosts running the byte-identical container image can still produce different behavior (or fail to run at all) if their host driver versions differ enough to fall outside the CUDA toolkit's supported driver-compatibility range, or, more subtly, if a driver update changes low-level kernel scheduling or floating-point behavior in a way that affects numerical results at the margins. The host's Linux KERNEL version matters similarly for anything relying on kernel-level behavior the container doesn't isolate (e.g. specific NCCL/RDMA driver interactions, GPU peer-to-peer access configuration), which is invisible to and unmanaged by the container image itself.
- Mitigations for the host-side gap: document and enforce a MINIMUM supported host driver version (and ideally a validated, narrow range) alongside the image tag, since "any driver new enough to run this CUDA version" is a looser reproducibility bar than "this exact validated driver version"; where infrastructure allows it, standardize the host driver/kernel version across the fleet via the same infrastructure-as-code discipline applied to the container image itself (e.g. a pinned host AMI/base image for the node pool), rather than treating the host environment as out of scope simply because it's outside the container boundary.
- Reproducible builds: use a version-controlled Dockerfile that produces a deterministic image, and tag/publish images immutably.
Worked example
A training image tagged training:v2.3.1-cuda12.1-cudnn8.9-torch2.1.0 runs correctly and reproducibly on host A (driver 535.104) and host B (driver 550.90.07), both within CUDA 12.1's supported driver range, so the container-internal stack is identical, but a subtle numerical difference in one specific fused kernel is later traced to the driver version difference, not anything inside the image; the fix is documenting a validated driver version (or narrow range) as part of the reproducibility contract, not just the image tag, and re-running the validation suite whenever the fleet's host driver is upgraded.
Trade-offs & pitfalls
Pinning every version this strictly trades convenience for reproducibility guarantees. The host driver/kernel gap is a common blind spot precisely because it's easy to assume containerization solves reproducibility completely; teams that only pin the image and never document or standardize the host driver version can spend significant debugging time chasing an apparent "non-determinism" that's actually a host-environment difference the container never controlled.
Outline a CI pipeline for ML training code that runs unit tests, environment reproducibility checks, small-scale integration trainings, and artifact validation before allowing full-scale runs. Describe tools, gating criteria, and how you would prevent flaky non-deterministic behavior from failing CI.
Sample Answer
Direct answer
A CI pipeline for ML training code should run fast unit tests on every commit, environment reproducibility checks, small-scale integration training runs, and artifact validation before any change reaches full-scale training, with the whole pipeline deliberately designed around tolerant, seeded, threshold-based checks rather than exact-match assertions, since naive exact-equality checks on anything involving GPU-kernel or stochastic behavior will be flaky and erode the team's trust in CI.
Structured elaboration
- Unit tests: fast, isolated tests of individual components: a custom loss function's output on known inputs, a data-preprocessing function's correctness, a model architecture's output shape.
- Environment reproducibility checks: verify the dependency lockfile resolves to the exact expected versions.
- Small-scale integration trainings: run the actual training entry point for a handful of steps on a small synthetic or subsampled dataset, checking the loop executes without error, loss is finite and decreasing, and checkpoint save/load round-trips correctly.
- Artifact validation: validate that expected artifacts were produced with expected structure/schema.
- Preventing flaky non-deterministic behavior from failing CI: this needs to be an explicit design principle across all four tiers, not an afterthought. First, every CI training run is seeded (fixed RNG seeds across Python/NumPy/framework, as in the general reproducibility checklist) so at least the SAME CI run is reproducible if re-triggered; deterministic-kernel flags are enabled specifically in CI (accepting the performance cost, since CI runs are tiny) so GPU-kernel-level non-determinism doesn't add noise on top of RNG seeding. Second, CI assertions on numerical outputs use TOLERANCE-based comparisons (e.g.
assert abs(loss - expected) < 1e-4, or a directional check like "loss decreased over N steps" rather than "loss equals exactly X"), since even with seeding and deterministic mode, minor floating-point differences across different CI runner hardware generations are still possible and a bit-exact assertion would fail spuriously on a legitimate hardware change. Third, a genuinely flaky check (one observed to fail intermittently despite no real code change) is treated as a bug in the TEST, not silently retried into passing or ignored: it's investigated (usually an under-seeded random source, exactly the failure mode a good reproducibility checklist is meant to prevent) and fixed or, if it can't be immediately fixed, explicitly quarantined (marked known-flaky, excluded from the merge-blocking gate, and tracked as an open issue) rather than left in the blocking suite where it silently trains the team to re-run CI without investigating red builds.
Worked example
A pull request modifying the data-augmentation pipeline triggers: unit tests confirming correctly-shaped output; an environment check confirming the lockfile resolves cleanly; and a 20-step integration training run on a 100-example synthetic dataset with a fixed seed and deterministic kernels enabled, asserting loss is finite and its value after step 20 is within a small tolerance band of a previously-recorded reference value (not required to match exactly), all completing in under 5 minutes. When this check failed intermittently once despite no code change, the team traced it to the augmentation library's own internally-seeded RNG not being re-seeded from the run's base seed between CI invocations; fixing that seeding gap (not adding a retry-until-green step) resolved the flakiness.
Trade-offs & pitfalls
The integration-training tier needs to be genuinely small and fast to be practical as a CI gate. A common design mistake specific to flakiness is reaching for automatic retries ("just re-run failed CI jobs once") as the default fix; this masks real flakiness sources (like the under-seeded augmentation RNG above) instead of fixing them, and a genuinely broken change can pass on a lucky retry, which is a worse outcome than an honest, investigated red build.
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.
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.