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.
What is gradient clipping and why is it used in training large models? Explain at least two clipping strategies (global norm, value clipping) and their impact on training dynamics and stability.
Sample Answer
Direct answer
Gradient clipping caps the magnitude of gradients before the optimizer update to prevent occasional very large gradients (from an unstable loss landscape, a bad mini-batch, or numerical issues) from causing a destructively large parameter update; the two common strategies are global-norm clipping (rescale the entire gradient vector if its overall norm exceeds a threshold) and value clipping (clamp each individual gradient component to a fixed range independently).
Structured elaboration
- Global norm clipping: compute the L2 norm of the concatenated gradient across all parameters; if that norm exceeds a threshold
max_norm, rescale every gradient component by the same factormax_norm / normso the overall gradient direction is preserved but its magnitude is capped. This is the more common choice for deep networks (especially RNNs and transformers) because it preserves the gradient's relative direction across parameters, only shrinking its overall size. - Value clipping: clamp each individual gradient component independently to a fixed range (e.g.
[-c, c]), regardless of the overall norm. This is simpler and cheaper to compute (no need to first compute a global norm across all parameters), but can distort the gradient's direction, since components near the clip boundary get truncated disproportionately relative to smaller components, unlike norm clipping's uniform rescaling. - Impact on training dynamics: clipping prevents a single unstable step from derailing training (a common cause of NaN/divergence in RNNs and transformers, especially early in training or with an aggressive learning rate), at the cost of somewhat biasing the effective update whenever clipping actually triggers (the update is smaller and, for value clipping, differently-directed than the "true" unclipped gradient would have produced).
Worked example
A gradient vector with an unusually large spike (norm 50, versus a typical norm around 2) and a clip threshold max_norm=1.0: global norm clipping rescales every component by 1.0/50 = 0.02, preserving the gradient's direction while capping magnitude to exactly 1.0; value clipping with a threshold c=0.1 instead clamps any individual component exceeding 0.1 to exactly 0.1 regardless of the others, which for a gradient with most components small but a few very large ones changes the update's direction, not just its magnitude, since the large components get disproportionately truncated relative to the small ones.
Trade-offs & pitfalls
Norm clipping's threshold needs to be chosen relative to the model's typical (unclipped) gradient norm during stable training, not an arbitrary small number; a threshold set far below typical gradient norms clips essentially every step, meaningfully slowing learning by systematically shrinking updates that weren't actually problematic.
Explain how batch size interacts with learning rate and gradient accumulation. Describe practical heuristics (e.g., linear learning rate scaling, warmup), how to compute effective batch size, and when gradient accumulation is preferable to increasing physical batch size.
Sample Answer
Direct answer
Batch size, learning rate, and gradient accumulation are tightly linked: increasing the effective batch size (whether by using a larger physical batch or by accumulating gradients over more micro-batches) generally calls for a proportionally larger learning rate to keep the effective step size in a well-behaved range, and a warmup period to avoid instability from applying that larger learning rate to an undertrained model right at the start.
Structured elaboration
- The linear scaling rule: as a practical heuristic, when multiplying the effective batch size by a factor k, multiply the learning rate by the same factor k (starting from a known-good learning rate at some baseline batch size); this heuristic comes from the observation that a k-times-larger batch produces a gradient estimate with roughly k-times-lower variance, so a proportionally larger step can be taken with similar stability to the smaller-batch, smaller-step baseline.
- Warmup: rather than applying the fully-scaled (and therefore larger) learning rate from step one, gradually ramp it up over an initial warmup period (a few hundred to a few thousand steps, depending on scale); this matters more, and for more steps, the larger the target learning rate is, since applying a large step size to a randomly-initialized, not-yet-stabilized model is a common source of early-training instability that warmup specifically mitigates.
- Gradient accumulation's role: since gradient accumulation is mathematically equivalent to a larger physical batch (once losses are correctly averaged across the accumulated micro-batches, as covered elsewhere in this topic), the same linear-scaling-and-warmup logic applies based on the EFFECTIVE batch size (micro-batch size times accumulation steps), not the smaller physical micro-batch size alone; a common mistake is tuning the learning rate as if the smaller micro-batch size were the real batch size, under-scaling the learning rate relative to what the true effective batch size calls for.
- Computing effective batch size precisely: in a distributed, gradient-accumulated setup, effective batch size = (micro-batch size per device) x (number of gradient accumulation steps) x (number of data-parallel devices); dropping any one of these three factors (e.g. forgetting to multiply by the device count in a multi-GPU job) is the specific version of the mistake above that shows up once accumulation and data parallelism are combined.
- When gradient accumulation is preferable to increasing physical batch size: gradient accumulation is the right tool specifically when the physical batch size you want is limited by GPU memory rather than by anything else; it lets you reach a larger EFFECTIVE batch size (for its statistical benefits, like a smoother, lower-variance gradient estimate) without needing more memory, at the cost of proportionally longer wall-clock time per optimizer step, since the accumulation steps run sequentially rather than in parallel. When memory headroom exists to simply increase the physical batch size instead, that is generally preferable, since it reaches the same effective batch size without the added wall-clock cost that accumulation's sequential micro-batches impose.
- Practical tuning approach: rather than trusting the linear scaling rule blindly at very large scale-up factors (where it's known to become less reliable), validate empirically with a learning-rate sweep at the target effective batch size, using the linear-scaling estimate as a reasonable starting point/search-range center rather than a guaranteed-correct final answer.
Worked example
A known-good configuration at batch size 256 with learning rate 1e-3: scaling to an effective batch size of 2048 (a factor of 8, achieved either via a larger physical batch or via accumulating gradients over 8 micro-batches of size 256) suggests a starting learning rate of 8×10−3 per the linear scaling rule, combined with a warmup period (e.g. the first 1000-2000 steps ramping linearly from a small initial value up to 8×10−3) rather than applying 8×10−3 immediately from step one.
Trade-offs & pitfalls
The linear scaling rule is a well-established heuristic, not a law; it tends to work well up to moderate batch-size increases but becomes less reliable at very large scale-ups (very large effective batch sizes), where more sophisticated LR schedules or additional techniques (e.g. layer-wise adaptive rate scaling) are sometimes needed, and empirical validation at the actual target scale remains the reliable final check rather than trusting the heuristic's extrapolation blindly.
Describe how data loading and preprocessing can become a bottleneck in GPU training. Provide three concrete optimizations (for example: prefetching, parallel decompression, persistent workers, sharded datasets on fast storage) and explain when to apply each approach for multi-node training.
Sample Answer
Direct answer
Data loading and preprocessing become a GPU-training bottleneck when the CPU-bound work of reading, decoding, and augmenting each batch takes longer than the GPU takes to consume the previous batch, leaving the (expensive) GPU idle while it waits; the fix is almost always adding parallelism and overlap to the data pipeline, not touching the model or GPU compute at all.
Structured elaboration
- Prefetching: have the data loader prepare the next batch (or several batches ahead) while the GPU is still processing the current one, using a background thread or process so the CPU-bound preparation work overlaps with GPU compute instead of happening serially before each GPU step. For multi-node training specifically, prefetching is applied identically per node (it's a per-node, local pipeline concern) but becomes more important as node count grows, since with more nodes pulling from shared storage simultaneously, any per-request latency variance is more likely to occasionally exceed a shallow prefetch buffer's depth; deeper prefetch queues are the standard response.
- Parallel data-loading workers: use multiple worker processes (not just threads, to avoid Python's GIL limiting true parallelism for CPU-bound preprocessing work) to read and preprocess several samples/batches concurrently, scaling the number of workers to match available CPU cores until the CPU-side throughput comfortably exceeds what the GPU needs. Apply this on every node independently in multi-node training (each node's DataLoader workers are scoped to that node's own CPU cores and its own shard of the data, with no cross-node coordination needed), and increase worker count specifically when profiling shows a given node's GPU is starved even though the cluster's aggregate storage bandwidth is not the bottleneck, i.e. when the constraint is that individual node's CPU decode/augment throughput rather than the shared storage backend.
- Pinned memory and efficient host-to-device transfer: using pinned (page-locked) host memory for the batch before transferring to GPU, combined with asynchronous (non-blocking) transfer, reduces the time spent specifically in the CPU-to-GPU memory copy step, which can itself be a meaningful bottleneck even after the CPU-side preprocessing is fast. This applies per-node in multi-node training exactly as it does single-node (it's a local host-to-device transfer concern), and is worth enabling by default across every node in a multi-node job since its cost is negligible and its benefit compounds identically at any node count.
Worked example
Profiling a training loop shows GPU utilization sitting at 40%, with the profiler's timeline showing large gaps between consecutive GPU compute kernels; increasing the DataLoader's num_workers from 2 to 8 (matching available CPU cores) and enabling pin_memory=True for the host-to-device transfer closes most of these gaps, raising GPU utilization to over 90%, confirming the bottleneck was CPU-side data preparation, not GPU compute capacity.
Trade-offs & pitfalls
Adding too many data-loading worker processes beyond what the machine's CPU core count can genuinely support in parallel doesn't help further and can actually hurt (context-switching overhead, memory pressure from many workers each buffering samples); the right worker count is typically close to the number of physical CPU cores available on the node, not an arbitrarily large number.
Explain synchronous versus asynchronous stochastic gradient descent in a distributed data-parallel setup. Discuss convergence guarantees, staleness, and scenarios where asynchronous updates are attractive despite potential instability.
Sample Answer
Direct answer
Synchronous SGD has every worker compute a gradient against the same, current parameter values and waits for all workers before applying a single combined update, giving convergence behavior equivalent to (or very close to) single-machine SGD at a larger effective batch size; asynchronous SGD lets each worker push its gradient and pull fresh parameters independently, without waiting for others, trading some workers computing gradients against slightly outdated ("stale") parameters for higher hardware utilization.
Structured elaboration
- Synchronous: every worker's gradient this step is computed against identical parameter values (the state after the previous step's update); once all gradients arrive, they're averaged and applied as one update, after which every worker again has identical, up-to-date parameters. Convergence guarantees closely mirror standard SGD's, since the process is mathematically equivalent to computing a gradient over a larger effective batch (the concatenation of every worker's mini-batch).
- Asynchronous: a worker pulls current parameters, computes a gradient, and pushes it back independently of other workers' progress; by the time its push arrives, the server's parameters may have already been updated by other workers' pushes in the meantime, meaning the pushed gradient was computed against parameters that are now "stale" (out of date) relative to the current server state.
- Staleness and its effect: the degree of staleness (how many other updates happened between a worker's pull and its push) tends to grow with more workers and with heterogeneous worker speeds (a slow worker's gradient becomes more stale the longer it takes to compute); staleness biases the effective update direction, since it's technically a gradient of an earlier point on the loss surface being applied to a later point, which can slow or, in extreme cases, destabilize convergence if unbounded.
- When each is chosen: synchronous is the default for most modern large-scale training (predictable convergence behavior, well-supported by AllReduce-based collectives) provided stragglers are managed; asynchronous is chosen specifically when worker heterogeneity or unreliability is severe enough that waiting for the slowest worker every step would be prohibitively wasteful, accepting some convergence-quality cost in exchange for higher aggregate hardware utilization.
Worked example
With 8 workers, one of which is consistently 3x slower than the others (a straggler): synchronous training's every-step wall-clock time is bounded by that slowest worker, wasting the other 7 workers' idle time waiting each step; asynchronous training lets the 7 faster workers keep contributing updates continuously without waiting, at the cost of the slow worker's occasional contributions being noticeably stale (computed against parameters several updates out of date) by the time they arrive.
Trade-offs & pitfalls
Bounded-staleness schemes (allowing async updates but capping how stale any single contribution is allowed to be before it's rejected or down-weighted) are a common middle ground, retaining most of asynchronous training's utilization benefit while limiting the worst-case convergence-bias risk that fully unbounded asynchrony carries.
List common storage formats and ingestion strategies for large ML datasets (examples: TFRecord, Parquet, LMDB, raw images). For each, explain trade-offs in throughput, random-access patterns, compression, and suitability for streaming versus batch ingestion in distributed training.
Sample Answer
Direct answer
Different storage formats for large ML training datasets trade off differently on read throughput, random-access support, schema flexibility, and ecosystem compatibility: TFRecord and similar sharded binary formats optimize for fast sequential streaming reads, Parquet optimizes for columnar analytical access and compression, and LMDB (and similar key-value stores) optimizes for fast random access to individual records.
Structured elaboration
- TFRecord (and similar sharded-binary/WebDataset-style formats): a sequence of serialized records written sequentially into shard files; optimized for fast sequential (streaming) reads, which matches how training typically consumes data (shuffled at the shard or buffer level, not truly randomly accessed record-by-record); poor for random access to a specific individual record, and less flexible for schema evolution than a columnar format.
- Parquet: a columnar storage format, excellent for reading a subset of columns/features efficiently (skip reading columns you don't need) and for strong compression (similar values grouped together compress well); more naturally suited to tabular/structured data and analytical access patterns than to, say, raw image or audio bytes, though it can store binary blobs as a column too; less common as the primary format for large unstructured (image/video/audio) training data specifically because its columnar-read advantage matters less when every "row" needs its full raw content read anyway.
- LMDB (and similar embedded key-value stores): memory-mapped, very fast random access to individual records by key, useful when a workload genuinely needs random (not just shuffled-sequential) access patterns, at the cost of typically worse write throughput and less natural support for very large aggregate dataset sizes spread across a distributed storage backend (LMDB is fundamentally a single-file, single-machine-oriented store, less naturally distributed than sharded object-storage formats).
- Raw images (e.g. individual JPEG/PNG files, whether on a local filesystem or as individual object-storage keys): the simplest possible representation, one file per example, with no format-level wrapping at all; throughput suffers from per-file open/seek/request overhead when there are millions of small files (each read pays a filesystem or object-storage request cost independent of the file's actual byte size), which is precisely the overhead that sharded formats like TFRecord/WebDataset exist to amortize by batching many examples into one larger file; random access to a specific image by ID is straightforward (open that one file) since each example already is its own addressable unit, unlike a packed shard where reaching one record means seeking within a larger file; compression is per-file only (whatever the image codec itself provides, e.g. JPEG's own lossy compression) with no opportunity for cross-example compression gains; well-suited to small-to-medium datasets or to workloads needing frequent random single-example access (interactive labeling, debugging), poorly suited to large-scale distributed training where the sheer number of small-file requests becomes the bottleneck long before raw bytes-per-second does.
- Trade-off summary: for large-scale distributed training where the dominant access pattern is streaming through shuffled shards (the common case), TFRecord-style sharded formats are usually the best default; Parquet earns its place when the workload has genuinely tabular/structured data and benefits from columnar compression/selective-column reads; LMDB earns its place for smaller-to-medium datasets needing fast true-random access, less commonly the right choice at the largest distributed-training scales.
Streaming versus batch ingestion: TFRecord/WebDataset-style formats are built for streaming (sequential read through a shard, consumed as it arrives, no need to have the whole shard resident before starting), which matches the common distributed-training access pattern of continuously feeding a GPU without waiting on a full batch job to materialize first; Parquet is more naturally suited to batch ingestion (a full columnar file is typically read, or partially read by column, as one bounded operation feeding an ETL or analytical job) though it can be read in row-group chunks in a streaming-like fashion; LMDB sits closer to batch/offline ingestion too, since its value comes from having the full key-value store built and available for random lookups, not from being consumed as an in-progress stream; raw images support both trivially (each file is its own unit either way) but pay the small-file-request overhead in either mode.
Worked example
A large image classification dataset (millions of images) is well-served by sharded TFRecord or WebDataset-tar files, streamed sequentially per epoch with shard-level shuffling; a tabular click-prediction dataset with hundreds of numeric/categorical feature columns is well-served by Parquet, letting a feature-selection experiment read only the specific columns needed without paying the I/O cost of the columns being ignored; a dataset needing frequent, genuinely random single-record lookups (e.g. an active-learning loop repeatedly sampling specific individual examples by ID) is a case where LMDB's random-access strength actually matters.
Trade-offs & pitfalls
Choosing a format based on familiarity rather than actual access pattern is a common mistake; a team defaulting to Parquet for large image data (paying columnar-format overhead for a workload that gets none of columnar's benefit, since every "row" needs its full image content regardless) or to TFRecord for a tabular feature-selection workload (losing the ability to cheaply read a subset of columns) both leave real performance on the table.
Unlock Full Question Bank
Get access to all 22 Model Training Infrastructure and Distributed Training interview questions and detailed answers.
Sign in to ContinueJoin thousands of developers preparing for their dream job.