Debugging and Testing ML Systems Questions
Finding, diagnosing, and fixing problems in ML code, data, and models, and building tests that catch these problems before they reach users. Covers common ML pitfalls (data leakage, shape mismatches, silent training bugs, mis-specified loss or metrics), root-cause analysis of model regressions and production incidents (accuracy drops, calibration drift, intermittent or hard-to-reproduce failures), distributed-training-specific failures (multi-GPU divergence, intermittent OOM, precision-related instability), and the diagnostic tooling that supports it (reproducibility artifacts, structured logging, instrumentation). Also covers testing ML systems directly: unit tests for data and feature pipelines, validation checks for datasets and features, test oracles and acceptance criteria for probabilistic or non-deterministic model outputs, and integration and regression tests that catch model or pipeline regressions before deployment. Emphasizes the engineering rigor that keeps ML systems correct and maintainable.
A large model-training run fails with a GPU out-of-memory error when you increase the batch size. List the practical mitigation strategies available to you, and describe the trade-off each one makes (extra compute time, implementation complexity, or a change in effective batch-size semantics).
Sample Answer
Direct answer. When a larger batch size pushes training past available GPU memory, the available fixes trade off differently between compute time, implementation complexity, and whether they change the actual math of what "batch size" means, so the right choice depends on which of those you can afford to spend.
The mitigation menu.
- Gradient accumulation: run several smaller micro-batches, summing their gradients before a single optimizer step, simulating a larger effective batch size without ever holding the full batch in memory at once. Trade-off: nearly free in implementation complexity and mathematically equivalent to true large-batch training for most loss functions, but costs extra wall-clock time since you're doing more forward/backward passes for the same effective batch. The exception is the batch-size-semantics one, and it is worth knowing by name: any layer that computes statistics over the batch, batch normalization above all, sees only the micro-batch, not the full effective batch, so a network containing batch normalization is NOT equivalent under accumulation and needs synchronized batch normalization across micro-batches or a normalization layer that does not depend on the batch at all.
- Mixed precision: store activations (the intermediate outputs each layer produces during the forward pass, which are kept in memory because the backward pass needs them, and which usually dominate memory during training) and weights in fp16 or bf16 (bf16 being a 16-bit format with fp32's exponent range but fewer mantissa bits, so it trades precision for far less overflow risk) instead of fp32, roughly halving memory footprint. Trade-off: often a "free" win in both memory and speed on modern hardware, but introduces its own numerical-stability considerations (dynamic loss scaling, meaning the loss is multiplied by a large factor before backward and the gradients divided by it afterwards, so that small gradient values do not vanish to zero in fp16, with the factor adjusted automatically when it causes an overflow; plus watching for NaNs from fp16 underflow and overflow), so it is not entirely free of engineering cost.
- Model or pipeline parallelism: split the model itself (not just the data) across multiple devices, so no single device needs to hold the full model or full activations. Trade-off: the most memory-scalable option but the most complex to implement correctly, and introduces communication overhead between devices.
- Activation checkpointing: don't store intermediate activations from the forward pass, recompute them during the backward pass instead. Trade-off: significantly reduces memory (often the single biggest lever for very deep networks) at the direct cost of extra compute, since the forward pass effectively runs twice for the checkpointed segments.
- Reducing sequence length: for sequence models, truncating or bucketing to a shorter maximum length directly reduces the memory footprint of attention and activations. Trade-off: free in implementation cost, but changes what the model actually sees, potentially losing information for the fraction of examples whose true length exceeds the new cap.
- Reducing the batch size itself, with the learning rate re-tuned to match: the option that is easy to forget precisely because it looks like giving up. Trade-off: free in both compute and implementation complexity, and the one option on this list whose cost lands squarely on effective batch-size semantics. A smaller batch means noisier gradient estimates and a different effective noise scale, so the learning rate and schedule usually need re-tuning (often downward), and the run's convergence path, and sometimes its final quality, genuinely changes rather than merely being reorganized.
- Parameter sharding (splitting optimizer state and/or parameters themselves across devices, as in ZeRO-style approaches, ZeRO being a family of techniques that partitions optimizer state, gradients and eventually parameters across devices instead of replicating them on every one): substantially reduces per-device memory for the optimizer state specifically (which is often larger than the model weights themselves for adaptive optimizers, that is, optimizers such as Adam that keep one or two running statistics per parameter and therefore carry several times the model's own size in state). Trade-off: adds communication overhead to gather sharded parameters when needed, and is more complex to set up than gradient accumulation or mixed precision.
How to choose. The two to know cold are mixed precision and gradient accumulation: both are close to free in engineering cost and, batch-normalization layers aside, neither changes what the model learns. Activation checkpointing is the third to reach for, since it is a pure compute-for-memory trade with no semantic change at all. Only then reach for model or pipeline parallelism, sharding, reducing sequence length, or simply shrinking the batch, since each of those trades away something the first three do not: implementation complexity and communication overhead for parallelism and sharding, information in the input for a shorter sequence, and effective batch-size semantics for a smaller batch with a re-tuned learning rate.
A new model version passed all pre-deployment tests but shows notable degradation once in production. Walk through a forensic post-deployment incident-response plan: what evidence you would gather (logs, feature snapshots, model versions), how you would compare pre- and post-deployment input distributions and schemas, how you would check for a data-pipeline or feature-transformation change that the pre-deployment tests didn't cover, and what you would change about your test suite so a similar regression is caught next time.
Sample Answer
Direct answer
A model that passed every pre-deployment test and still regressed in production tells you something specific: the tests were measuring the wrong thing, or measuring the right thing on the wrong data. My first move is forensic, not corrective: freeze and gather evidence before anything about the environment changes further, compare the world the tests ran against to the world production actually sees, and only then propose a fix. The closing, and in my view most important, part of the response is diagnosing exactly which gap in the test suite let this through, because fixing this one regression without closing that gap just means the next one repeats it.
Structured elaboration
Evidence to gather, before it rotates out of retention. Pull the exact model artifact and its version identifier that was deployed (not "the version that was supposed to be deployed"; deployment mismatches are themselves a real cause), the serving configuration active at the time, and a sample of real inference logs from the incident window (inputs as hashes or fingerprints if raw payloads cannot be retained, plus outputs, confidence scores, and latencies). Separately capture a feature-store snapshot for the incident window, since features can be point-in-time and the values a request saw when it happened may not match what you would compute by re-querying now. Note the exact deployment timestamp and cross-reference it against any other change that landed around the same time (a feature-store schema change, an upstream data-source migration, a dependency bump), since "new model version" and "new bug" being simultaneous does not prove the model version caused it.
Compare pre- and post-deployment input distributions and schemas. For every feature the model consumes, compare the distribution in the incident window's production traffic against both the training distribution and the pre-deployment test distribution, using a two-sample test such as Kolmogorov-Smirnov (KS) for continuous features or a chi-square test for categorical ones. Separately, and more cheaply, diff the schema itself: field names, types, nullability, and cardinality, since a schema drift (a field silently becoming nullable, a categorical feature gaining new values) is a harder, more binary signal than a distributional shift and is usually faster to confirm or rule out.
Check for a data-pipeline or feature-transformation change the pre-deployment tests did not cover. This is the step that specifically explains "passed all tests and still failed": walk the feature-transformation code path end to end and ask, for each step, whether the pre-deployment test suite exercised it against data shaped like what actually reached production. Common gaps: the test suite used a static, curated golden dataset (a fixed, checked-in set of example inputs with their expected outputs, captured once and used as the pre-deploy reference ever since) that no longer represents the live feature distribution; the test suite validated the model in isolation but not the full serving path, so a change in an upstream feature-extraction service between test and deploy went unchecked; or the test suite covered the model's own code but a shared preprocessing library was bumped independently and neither the model tests nor a compatibility test caught the interaction. Confirm by re-running the pre-deployment test suite's own inputs through the CURRENT (post-deployment) feature pipeline and diffing the outputs against what the test suite recorded at deploy time; a difference here directly proves the pipeline moved out from under a test that never re-checked itself against fresh production-shaped data.
What to change about the test suite so this class of regression is caught next time. Once the specific gap is identified, the fix is almost always one of three shapes: replace a static golden dataset with a rolling, periodically-refreshed sample of real (or synthetic but representative) recent production data, so the test data cannot silently drift out of sync with what production actually sees; extend the pre-deployment suite from model-only testing to full serving-path integration testing, so a change anywhere in the request's actual path is exercised, not just the model's own forward pass; and add a schema and distribution CONTRACT check as a hard pre-deploy gate (not just a test that can be skipped under time pressure), so any feature schema or distribution change beyond a defined tolerance blocks the deploy automatically rather than depending on a human noticing. Define that tolerance in numbers or it will not survive its first deadline: block the deploy if any feature's Kolmogorov-Smirnov statistic against the last 7 days of production exceeds 0.1, if its null rate moves by more than 2 percentage points, or if its median moves by more than 20 percent. A 0.1 gate is deliberately loose relative to sampling noise: two clean samples of the sizes below land around 0.01 to 0.02, so 0.1 blocks real movement without firing on ordinary variation.
Worked example
Concretely: precision at the deployed threshold fell from 71 percent to 63 percent over the first two days after the release, an 8-point drop. The post-deployment schema diff shows nothing (all fields present, all types unchanged), but the KS test on one input feature, "days since last activity," shows a clear shift between the pre-deployment test distribution and the incident-window production distribution. The numbers: the golden set holds 5,000 rows with a median of 14 days since last activity; the incident window holds 20,000 production rows with a median of 31 days; the two-sample KS statistic between them is 0.294 with a p-value below 10−300. For scale, two clean samples drawn from the same distribution at those sizes give a KS statistic of about 0.013, with 0.020 at the 95th percentile, so 0.294 is not a borderline call, it is roughly fifteen times the noise floor. Re-running the pre-deployment test suite's own golden inputs through the current feature pipeline reproduces the ORIGINAL test-time values exactly (the pipeline code did not change), which rules out a pipeline bug. The remaining explanation, confirmed by checking the golden dataset's last-refreshed timestamp, is that the golden dataset itself was captured eight months earlier and the live user population's activity pattern had genuinely moved since then; the pre-deployment tests kept passing because they were comparing the new model against stale data, not because the new model actually generalized to current production inputs. Why this version and not the earlier ones that tested against the same stale golden set: the new version was tuned considerably harder against that distribution (more capacity and a longer hyperparameter search, selected on golden-set performance), so it fit the eight-month-old activity pattern more tightly and had correspondingly more to lose when production stopped matching it. The earlier, less tightly-fit versions were leaving performance on the table against the golden set, which is exactly what made them less sensitive to the golden set being wrong. The test-suite fix that follows directly from this evidence is the rolling-refresh golden dataset described above, specifically for this case, not a full serving-path integration test, since the pipeline code itself was proven unchanged.
Trade-offs and pitfalls
The most common wrong turn is rolling back and closing the incident once metrics recover, without completing the forensic comparison. A rollback restores service; it does not tell you whether the next model version will hit the exact same gap, and if the golden dataset is stale, every future release trained and tested the same way is at risk until that is fixed.
A second pitfall is over-correcting the test suite into something too expensive to run on every deploy. A full production-scale, fully-refreshed integration suite on every commit can become slow enough that teams start skipping it under deadline pressure, recreating the exact failure mode you were trying to close. Split the strategy by cost: fast, deterministic schema and distribution CONTRACT checks as a hard, cheap gate on every deploy, and a periodic (not necessarily per-deploy) refresh of the golden dataset and full-path integration run, so the properties that change slowly are not re-verified at a cost that does not match how often they actually move.
A third pitfall is confirming only ONE candidate cause and stopping. In this example, the schema diff, pipeline replay, and distribution test were run as a set specifically because any one check alone (say, the KS test showing drift) would not have distinguished "the golden dataset is stale" from "the pipeline itself changed and happens to now match old test values by coincidence." Running the full battery, and confirming the pipeline replay reproduced the ORIGINAL recorded values, is what turned a plausible hypothesis into a confirmed root cause.
Write a detailed bug report for an intermittent ML training failure that surfaces as a non-deterministic CUDA kernel launch happening only on specific GPU types. Include what timestamps and environment matrix (drivers, CUDA, cuDNN, framework versions) you would capture, minimal steps to reproduce, sample seeds, which logs to attach, an impact assessment, and your recommended next experiment or workaround.
Sample Answer
Direct answer
A bug report for an intermittent, GPU-type-specific training failure has to do the work a live debugging conversation would otherwise do informally, because the person who eventually picks it up will not have access to the machine, the moment, or your memory of what you tried. That means an exact environment matrix (down to driver, CUDA, the NVIDIA GPU computing toolkit, and cuDNN, NVIDIA's deep-learning kernel library, versions, since this exact bug class is version-sensitive by nature), a minimal reproduction that isolates the GPU type as the variable, every seed pinned so any attempted repro is comparing apples to apples, the specific logs to attach and why each one matters, an honest impact assessment, and a concrete next experiment, not just a description of the symptom.
Structured elaboration
Why this bug class specifically demands more than a symptom description. "Training sometimes crashes" is not actionable; a bug report is only as useful as the next person's ability to reproduce, or rule out, the same failure without your involvement. For a bug that is intermittent (does not happen every run) and hardware-specific (only certain GPU types), the report has to actively narrow the search space: which axis is the variable, hardware, driver version, a specific kernel, and provide enough detail that the report itself becomes evidence, not just a complaint.
Timestamps, and why they belong in the report as their own field, not folded into "environment." Record the exact wall-clock time of every failing run, alongside the exact wall-clock time of every successful run in the same investigation, in a consistent timezone, UTC (Coordinated Universal Time, the timezone-independent standard most infrastructure logs already use). This is what lets a later reader correlate a failure against something that changed at a specific moment but is not itself part of the environment matrix: a driver auto-update that rolled out to part of the fleet, a concurrent job that briefly contended for the same GPU, or a thermal-throttling event visible in a separate infrastructure dashboard. A bug report that only states "this happens sometimes" gives a later investigator nothing to line up against those external timelines; a report with exact timestamps lets them check in minutes whether the failures cluster around a specific external event instead of being uniformly distributed, which is itself a strong clue about the root cause.
What the environment matrix needs to capture, and why each field matters. GPU driver version and CUDA (Compute Unified Device Architecture, NVIDIA's parallel-computing platform) toolkit version together determine which low-level kernel implementations get selected for a given operation; a driver/CUDA combination that works on one GPU architecture can expose a genuine race condition or an unsupported code path on another. cuDNN (CUDA Deep Neural Network library, NVIDIA's library of GPU-accelerated primitives for deep learning) version matters separately because it ships its own kernel selection heuristics, independent of the CUDA toolkit version. Framework version (the training library itself) matters because kernel dispatch logic lives partly in the framework, not only in CUDA/cuDNN. Recording all four together, rather than just one, is what lets someone later correlate this exact failure against a known upstream bug report in any of the four projects.
Minimal reproduction steps. State the smallest model, batch size, and number of training steps that still reproduces the failure, and the specific GPU type(s) it reproduces on versus the GPU type(s) it was also tried on and did not reproduce on, since the negative result (GPU type X: 0 failures in N attempts) narrows the search space just as much as the positive one does.
Seeds. Pin and report every seed used in the reproduction attempts, since without them a colleague's "I couldn't reproduce it" is uninformative, they might simply have drawn a different random path through the same non-deterministic kernel. Also report whether the failure reproduces on the SAME seed across multiple attempts (which would suggest the non-determinism lives somewhere other than the obvious PyTorch/NumPy random state, likely the CUDA kernel scheduling itself) or only appears on some seeds and not others.
Which logs to attach. The full stdout/stderr from a failing run (the actual CUDA error string and code, not a paraphrase of it), the nvidia-smi output captured at or near the time of failure (GPU memory usage and utilization can distinguish an out-of-memory-adjacent race from a pure kernel-launch bug), and, if available, output from CUDA_LAUNCH_BLOCKING=1 (an environment variable that forces synchronous kernel launches, which converts an asynchronous, hard-to-localize CUDA error into one that points at the actual failing kernel call, at the cost of much slower execution, which is exactly why it is a debugging-only flag rather than something left on by default).
Impact assessment. State concretely which training runs and which GPU types are affected, and the failure rate observed so far in this specific investigation (framed honestly as "N failures in M attempts on this GPU type, this seed set," not extrapolated into a general production failure rate unless that rate has actually been measured at production scale).
Recommended next experiment or workaround. Propose the next diagnostic step that would most efficiently narrow the remaining hypothesis space (for example, running with CUDA_LAUNCH_BLOCKING=1 to localize the failing kernel, or bisecting across driver versions on the affected GPU type specifically), and a concrete interim workaround if one exists (routing training for the affected GPU type to a known-good driver version, or avoiding the specific operation suspected of triggering it), so the team is not blocked on training while the root cause is still under investigation.
Worked example
A filled bug report for exactly the scenario in the question:
Title: Intermittent CUDA kernel launch failure during training, reproduces only on GPU type A100, not on V100 or H100 (A100, V100, and H100 are successive generations of NVIDIA's data-center training GPUs)
Environment matrix
Field Affected runs Unaffected runs GPU type NVIDIA A100 (80GB) NVIDIA V100, NVIDIA H100 GPU driver 535.104.05 535.104.05 (same driver, different GPU type) CUDA toolkit 12.2 12.2 cuDNN 8.9.4 8.9.4 Framework PyTorch 2.3.0 PyTorch 2.3.0 Observed: 6 failures in 40 training attempts on A100 (15% failure rate in this investigation's sample; not yet measured at full production scale). 0 failures in 25 attempts on V100 and 0 failures in 15 attempts on H100, same code, same driver/CUDA/cuDNN/framework versions.
Timestamps of each A100 failure (UTC): 2026-07-14 03:12, 2026-07-14 09:47, 2026-07-15 03:20, 2026-07-16 03:15, 2026-07-18 03:09, 2026-07-19 22:58.
Tally of that list (a clustering claim that has not actually been counted is not evidence, so the count belongs in the report next to the raw timestamps):
classification timestamps count inside the 03:00-03:20 UTC window 07-14 03:12, 07-15 03:20, 07-16 03:15, 07-18 03:09 4 of 6, on 4 distinct days outside it 07-14 09:47, 07-19 22:58 2 of 6 Four of six landing inside a 20-minute window is the detail worth flagging first, ahead of the raw failure count. Twenty minutes is 20/1440 = 1.39% of a day, so if failures were spread uniformly through the clock you would expect 6 x 0.0139 = 0.08 of them to land in that window, not 4; the probability of seeing 4 or more of 6 there by chance is about 5.5e-07, roughly 1 in 1.8 million. That is far too strong to write off as coincidence, so it suggests correlation with a recurring scheduled job on the shared cluster (a nightly maintenance window is the working hypothesis, not yet confirmed) rather than a purely code-triggered race condition. The two failures outside the window are exactly why it stays a working hypothesis: a maintenance-window cause accounts for 4 of the 6 observations and something else, a second trigger or plain chance, still has to account for 07-14 09:47 and 07-19 22:58.
Minimal reproduction: a 2-layer transformer block, batch size 32, sequence length 128, fails intermittently within the first 50 training steps on A100. Does not require the full model or dataset; the minimal repro script isolates a single forward/backward pass repeated in a loop.
Seeds: reproduced with seeds 0, 17, and 42 on A100 (not seed-specific: same seed 0 both failed and succeeded across different attempts, which points toward non-determinism in kernel scheduling itself rather than in the framework-level random state).
Logs attached: (1) full stdout/stderr from a failing run, including the raw CUDA error string and code, not a summary; (2)
nvidia-smioutput captured within 5 seconds of the failure; (3) a run withCUDA_LAUNCH_BLOCKING=1set, which converts the normally-asynchronous CUDA error into one that names the specific failing kernel call, at the cost of roughly 3 to 5x slower execution for that run, which is why it is not run by default.Impact: blocks reliable overnight training runs specifically on the team's A100 nodes; V100 and H100 nodes are unaffected and available as an immediate workaround.
Recommended next experiment: (1) immediate workaround: route affected training jobs to V100 or H100 nodes while this is investigated; (2) next diagnostic step: check the cluster's job scheduler and infrastructure logs for anything recurring in the 03:00-03:20 UTC window to confirm or rule out the maintenance-window hypothesis the 4-of-6 clustering points at, and separately check whether the two out-of-window failures (07-14 09:47, 07-19 22:58) share any property with each other that the four in-window ones do not, since a hypothesis that only ever explains 4 of 6 is not yet a complete account; (3) in parallel, re-run the minimal repro with
CUDA_LAUNCH_BLOCKING=1on A100 across 20 attempts to identify the specific failing kernel by name, then check that kernel against NVIDIA's published A100 driver 535.x errata before assuming this is a novel bug rather than a known one.
Trade-offs and pitfalls
- Reporting failure rate without stating the sample size invites the reader to misjudge severity in either direction. "15% failure rate" sounds alarming or mild depending on whether it is 6 of 40 or 150 of 1000; always report both numbers together, and be explicit about whether the rate has been measured at production scale or only within this specific investigation's sample, per the reproducibility discipline of never presenting an estimate as a broader measurement than it actually is.
CUDA_LAUNCH_BLOCKING=1is a debugging tool, not a permanent fix. It serializes kernel launches, which can slow training by several times; attaching a run captured with it on is valuable evidence, but the recommendation should never be to leave it on in production training as a workaround for the underlying bug.- A same-seed-both-outcomes observation is a stronger clue than it looks, and is easy to bury in a wall of text. It specifically points toward non-determinism at the CUDA kernel-scheduling level rather than the framework's own random-number state, which changes where the next investigator should look first; surface it explicitly near the top of the report, not just in a logs appendix.
- A stated pattern that has not been counted is the same mistake as a failure rate quoted without its sample size. "The failures cluster overnight" is an assertion; "4 of 6, on 4 distinct days, inside a 20-minute window where 0.08 would be expected by chance" is evidence the next investigator can check and either build on or refute. Put the tally in the report next to the raw list, and let the prose summarise the tally rather than the other way round.
- A minimal reproduction that is not actually minimal wastes the next investigator's time. If the smallest failing case still requires the full model and dataset, keep reducing it (smaller batch size, fewer layers, synthetic data instead of the real dataset) before filing, since every extra unnecessary component in the repro is one more variable a colleague has to also rule out.
You scale synchronous distributed training from a small number of GPUs to many, or you enable mixed precision (AMP) and gradient accumulation, and training that was stable before now diverges or intermittently produces NaNs in the backward pass. Enumerate the plausible causes and propose a prioritized diagnostic plan: which single change you would test first, and what you would expect to observe if each cause were the real one.
Sample Answer
Direct answer. Scaling up GPU count or enabling mixed precision changes several things simultaneously (effective batch size, gradient synchronization behavior, numeric precision), so a training run that was stable before and diverges after needs a plan that isolates ONE variable at a time rather than reverting the whole change and guessing.
Three terms this depends on. A replica is one full copy of the model living on one GPU. DDP (DistributedDataParallel) is the synchronous data-parallel wrapper that runs one process per GPU and keeps the replicas in step. All-reduce is the collective operation that combines a tensor across every replica and hands every replica back the same combined result, which is how the per-replica gradients become one shared gradient each step.
Plausible causes, and what to test first. The single highest-prior cause when scaling GPU count is learning-rate scaling: if the effective batch size grows with GPU count but the learning rate doesn't, or is scaled incorrectly (linear scaling isn't always the right rule for every optimizer), training can become unstable purely from that mismatch, and it's the cheapest hypothesis to test. Worked with real numbers so "adjusted for the new batch size" is computable: a stable baseline of 8 GPUs at 32 examples per device is an effective batch of 256 at LR 0.1. Scaling to 64 GPUs at the same 32 per device gives an effective batch of 2048, an 8x increase. The linear rule says LR = 0.1 x 2048/256 = 0.8. The square-root rule says LR = 0.1 x sqrt(8) = 0.283. If the job was launched at 0.8 and diverges, the diagnostic rerun is at 0.283, and the sanity rerun (which trades throughput for certainty) is at the original 0.1.
Next, check batch-norm statistics: if BatchNorm computes per-replica statistics independently rather than synchronizing across replicas, a network with BatchNorm can behave differently, and less stably, at high replica counts purely because each replica now sees a smaller effective batch for normalization purposes. Then check all-reduce precision (is gradient synchronization happening in a lower precision than intended) and synchronization/seed-order differences (are all replicas actually initialized identically before training starts).
For mixed precision (AMP) specifically. Gradient-accumulation interactions with AMP are a common culprit. The mechanics, since the ordering is the whole bug: loss scaling multiplies the loss by a large factor S before the backward pass so that small fp16 gradients do not underflow to zero, and the gradients must later be divided by S ("unscaled") before the optimizer step. With gradient accumulation over K micro-batches, the correct order is accumulate-then-unscale: keep every micro-batch's gradients scaled, sum all K of them, and unscale ONCE immediately before the optimizer step. The wrong order, unscale-then-accumulate, divides each micro-batch's gradients by S as it goes and then accumulates already-unscaled values, while the scaler still believes it owes the division, so the effective gradient magnitude ends up off by a factor tied to K and the scaler's overflow check runs against tensors that are no longer on the scale it expects. Confirm which order your framework's AMP integration actually uses before assuming the bug is elsewhere. Separately, some DDP jobs show gradient divergence where SOME workers produce NaN and others don't purely from AMP's dynamic loss scaling adjusting independently per worker if it isn't explicitly synchronized, worth checking whether your framework synchronizes the loss-scale factor across replicas or lets each pick its own.
What you would expect to observe if each cause were the real one. This is the half that turns the list into a diagnosis:
- Learning-rate / batch-size mismatch: divergence arrives within the first few hundred steps and arrives EARLIER the more GPUs you add, and the step at which it blows up moves monotonically as you sweep LR. Rerunning the same 64-GPU job at 0.283 (or at 0.1) with everything else untouched restores stability. Gradient norms grow smoothly over several steps before the blow-up rather than jumping in one step.
- BatchNorm computing per-replica statistics: dump each replica's BatchNorm running mean and running variance buffers and compare them across ranks. Synchronized normalization keeps them near-identical; the unsynchronized failure shows them visibly diverging, with the spread widening as per-device batch size falls. Crucially, the gradients themselves stay finite the whole time, so this looks like instability without any numeric red flag. It gets worse as you hold total batch fixed and add GPUs (shrinking per-device batch), and it disappears when you swap in synchronized BatchNorm or raise the per-device batch back up.
- All-reduce running in a lower precision than intended: checksum or norm each rank's gradient tensor immediately BEFORE the collective and immediately AFTER it. Before, the ranks legitimately differ, since each saw a different data shard. After, every rank must hold a bit-identical tensor. The tell is that post-collective gradients agree only to the first three or four significant digits rather than exactly, and that the disagreement grows with rank count, because more partial sums mean more accumulated rounding. Compare against a reference reduction performed in fp32 to size the error.
- Seed or initialization-order mismatch: hash every replica's parameters at step 0, before any data is seen. They must be identical. If rank 0's hash differs from rank 1's at step 0, the replicas never started from the same weights, and the shared averaged gradient is being applied to divergent parameter sets. The signature is that divergence appears essentially at step 1, not after a hundred healthy steps.
- AMP unscale-then-accumulate ordering: instrument the global gradient norm immediately before the optimizer step, and compare against the same total batch run with K=1 (no accumulation). The two should agree within numerical noise. The ordering bug shows up as a clean multiplicative discrepancy tied to K, a factor of K or 1/K rather than a vague drift, which is what makes it identifiable rather than merely suspicious.
- AMP loss-scale desync across workers: log the current loss-scale value per rank per step. Healthy runs show every rank holding the same value and stepping it in lockstep. The failure shows ranks holding different scales at the same step (one at 1024 while another is at 32768), SOME ranks reporting NaN while others report clean gradients, and the identity of the NaN-reporting ranks changing from run to run.
Why the effective batch size is the thread running through all of this. Scaling GPU count in synchronous data-parallel training multiplies the effective batch size by the number of new replicas, which is exactly what makes the learning-rate hypothesis the highest-prior one and why the arithmetic above matters. It also means that if the instability appears only past a specific scale rather than at every scale, that threshold is itself evidence: a genuine LR/batch mismatch only becomes destabilizing once the step size and the curvature disagree badly enough, whereas a hardware or environment fault would be present regardless of how many replicas are running.
Prioritized order to test. Order by what each check actually costs to run, then by prior probability, because two of these are nearly free and one of them is a precondition for trusting any of the others.
- Free, so run these before spending a single GPU-hour on a rerun. Hash every replica's parameters at step 0: one hash per rank, zero training steps, and if the hashes disagree you are done, because none of the cross-rank comparisons below mean anything until the replicas provably start from the same weights. Under AMP, also log each rank's loss-scale value per step, which is one scalar per rank and settles the desync hypothesis outright. Neither of these is the most likely cause; they go first because they are cheaper than everything else and because they are the two that can invalidate the other measurements.
- Learning-rate scaling, the highest-prior cause and the first hypothesis worth real GPU-hours: rerun the 64-GPU job at 0.283, and at the original 0.1 if that is still unstable. One short run, one variable, directly falsifiable.
- AMP accumulate-then-unscale ordering, whenever gradient accumulation is on: compare the pre-optimizer global gradient norm against the same total batch run at K=1. One extra short run and one scalar, and the tell is a clean factor of K or 1/K rather than a drift.
- BatchNorm synchronization: dump the running-mean and running-variance buffers on every rank and diff them. More instrumentation than anything above, still far less than the last item.
- All-reduce precision, last: checksumming every rank's gradient tensor before and after the collective is the heaviest instrumentation in this list, and it is also the least common root cause for a symptom that appeared specifically at a scaling boundary.
Create a prioritized checklist of automated tests and validations you would include in ML continuous-integration for a production model, ordered from the earliest pipeline stage to the latest. For each item, explain briefly why it matters and which common production failure it prevents.
Sample Answer
Direct answer. A prioritized ML CI checklist should be ordered by where a bug is cheapest to catch versus how expensive it is to let it slip through, which generally means data checks first, training-time sanity checks second, pre-deploy checks third, and post-deploy smoke tests last as a final safety net.
The checklist, with the failure each item prevents.
- Feature/data validation (schema checks, null/range checks on incoming training data): prevents training on silently corrupted or malformed data, the cheapest bug to catch and the most expensive to discover after the fact, since a model trained on bad data has to be retrained, not just patched.
- Training-time sanity tests (a tiny-batch overfit test, a check that loss actually decreases over the first N steps): prevents shipping a model whose training pipeline is fundamentally broken (a wiring bug, a frozen layer, a disconnected loss), catching it in minutes instead of after a full multi-hour training run completes with a mysteriously bad result.
- Pre-deploy checks (a regression test against a frozen golden dataset, a minimum-accuracy threshold gate): prevents promoting a model that's measurably worse than what's currently in production, the last automated checkpoint before a real user is affected.
- Post-deploy smoke tests (a simple end-to-end request against the live endpoint immediately after rollout): prevents a broken DEPLOYMENT (not a broken model, a broken serving path, wrong artifact loaded, a crashed container) from silently serving errors or garbage before anyone notices.
Why this order matters, not just this list. A team that only has post-deploy smoke tests catches problems at the most expensive possible point, after a bad model or bad deploy has already reached users. A team that only has data validation catches data problems but has no safety net if the training code itself has a bug, or if a legitimately-trained model is simply worse than the one it's replacing. The checklist is valuable specifically because each layer catches a DIFFERENT class of failure that the layers before and after it cannot, not because more checks are categorically better; a team with limited time should build layer 1 and 2 first, since they're both cheap and catch the most common real bugs before any compute is wasted on a bad training run.
Unlock Full Question Bank
Get access to all Debugging and Testing ML Systems interview questions and detailed answers.
Sign in to ContinueJoin thousands of developers preparing for their dream job.