ML engineer role-specific interview questions

As asked

Walk me through what happens when you call loss.backward() in PyTorch. How does autograd build the computation graph, and what are the memory implications of retaining it?

Sample answer outline

A strong answer explains that PyTorch builds a dynamic computation graph during the forward pass, attaching grad_fn nodes to tensors. backward() traverses the graph in reverse, applying the chain rule at each node. retain_graph=True is needed for multiple backward passes but prevents freeing intermediate activations, which can double GPU memory usage in practice.

Expect these follow-ups

• What happens if you call backward() twice without retain_graph=True?
• How does gradient accumulation work and why is it used in large batch training?

As asked

Describe how torch.cuda.amp works under the hood. What is the GradScaler doing, and what training instabilities have you seen when using mixed precision?

Sample answer outline

AMP runs the forward pass in float16 for speed and memory savings, while keeping a float32 master copy of weights. GradScaler multiplies the loss by a large factor before backward to avoid underflow in float16 gradients, then unscales before the optimizer step and checks for infs or NaNs, skipping the update if found. Common instabilities include loss spikes when the scale factor is too aggressive and divergence in layers that are numerically sensitive like LayerNorm.

Expect these follow-ups

• Why does bfloat16 behave differently from float16 in practice?
• How would you debug a training run where loss suddenly goes to NaN after switching to AMP?

As asked

Your training job shows GPU utilization fluctuating between 40 and 60 percent. Walk me through how you would diagnose and fix this.

Sample answer outline

Low GPU utilization usually means the GPU is starved for data, waiting on CPU-side preprocessing, or blocked on communication. Start with nvidia-smi or PyTorch Profiler to identify the idle periods. Profile the data loader by checking DataLoader num_workers and whether pin_memory is enabled. Check if memory transfers are saturating PCIe bandwidth. If the bottleneck is communication in distributed training, verify NCCL configuration and whether gradient compression helps.

Expect these follow-ups

• How do you use torch.profiler to find the specific kernel that is slow?
• What is the impact of num_workers on multi-GPU training versus single GPU?

As asked

Your model's prediction distribution has shifted over the past week, but labels are not available yet. How do you detect and diagnose the root cause?

Sample answer outline

Without labels, you monitor input feature distributions using statistical tests like Population Stability Index or Kolmogorov-Smirnov against a reference window. If PSI is above 0.2 on a key feature that is high in SHAP importance, that feature is the likely culprit. Check upstream data pipelines for schema changes, null rate spikes, or encoding changes. Model output distribution monitoring with Jensen-Shannon divergence can flag drift before labels arrive.

Expect these follow-ups

• How would you decide when drift is bad enough to trigger a retraining?
• What monitoring tool would you reach for first: Evidently AI, WhyLogs, or a custom solution?

As asked

Walk me through the decision of whether to serve a ranking model in real time versus running it in daily batch. What factors tip you toward each approach?

Sample answer outline

Batch inference is cheaper and simpler when predictions can be precomputed, like email campaign scoring or overnight recommendation refreshes. Real-time serving is required when context changes rapidly or per-request features are needed, like session-based recommendations. Cost, latency SLA, feature freshness requirements, and user session lifetime all factor in. A common hybrid is precomputing candidate sets in batch and re-ranking in real time.

Expect these follow-ups

• How do you handle cold-start users in a batch precomputed ranking system?
• What infrastructure do you need to support sub-100ms p99 latency for a real-time ranking model?

As asked

You need to run a hyperparameter sweep over learning rate, batch size, and weight decay for a transformer fine-tuning job. How do you set this up in Weights and Biases and avoid wasting compute on clearly bad configurations early?

Sample answer outline

W&B Sweeps support random, grid, and Bayesian search strategies. For a three-dimensional space, Bayesian search with early termination using the Hyperband scheduler eliminates low-performing runs before they finish, typically reducing compute by 60 to 80 percent. Log validation loss at each epoch so the scheduler can prune early. Use a logarithmic prior on learning rate and weight decay. Set a minimum number of steps before pruning to avoid killing runs that need warmup.

Expect these follow-ups

• How would you decide between Bayesian search and random search for a sweep with 20 hyperparameters?
• What W&B features would you use to compare two sweep runs that produced similar final metrics but different training dynamics?

As asked

Your fraud detection dataset has 0.1 percent positive labels. What techniques would you apply at the data, model, and evaluation level, and how do you choose between them?

Sample answer outline

At the data level, undersample the majority class or oversample with SMOTE, or use weighted random sampling in the DataLoader. At the model level, use focal loss or set class_weight='balanced' in sklearn models. At the evaluation level, use precision-recall AUC or F1 rather than accuracy or ROC-AUC, and set the decision threshold to the business-optimal point on the precision-recall curve. For fraud specifically, the cost of false negatives versus false positives should drive the threshold choice.

Expect these follow-ups

• Why can ROC-AUC be misleading at 0.1 percent positive rate?
• How do you calibrate predicted probabilities for a model trained with downsampling?

As asked

When would you choose PyTorch FSDP over DDP for a training job, and what are the concrete memory and communication tradeoffs you are making?

Sample answer outline

DDP replicates the full model on every GPU and syncs gradients via all-reduce, so each GPU must fit the whole model. FSDP shards parameters, gradients, and optimizer state across GPUs, allowing models that do not fit on a single GPU, but at the cost of additional all-gather and reduce-scatter communication. FSDP is the right choice when a single GPU cannot hold the model or optimizer state.

Expect these follow-ups

• How do you handle activation checkpointing alongside FSDP?
• What is ZeRO stage 3 and how does it relate to FSDP conceptually?

As asked

You are building a feature store for a real-time recommendation system. Walk me through how you keep the online store consistent with the offline store, and how you handle feature staleness during serving.

Sample answer outline

The offline store holds historical features computed in batch and is used for training, while the online store holds the latest feature values served at low latency. Consistency is maintained via a materialization pipeline that pushes updated features from the offline store to a low-latency store like Redis. Training-serving skew occurs when the online feature computation differs from training, so both pipelines should share transformation logic. Staleness is managed with TTLs and monitoring dashboards tracking feature freshness.

Expect these follow-ups

• How would you detect training-serving skew in production?
• What would you change in this design for features that must be computed in real time from raw events?

As asked

Explain the difference between post-training quantization and quantization-aware training. When does INT8 quantization cause unacceptable accuracy degradation, and how do you recover from it?

Sample answer outline

Post-training quantization calibrates scale factors from a representative dataset without retraining, which is fast but can lose accuracy on models with wide activation distributions. QAT simulates quantization noise during training so the model adapts, typically recovering most of the accuracy loss. Degradation is worst for models with outlier activations, like large transformer models, where per-channel or per-token quantization is needed. SmoothQuant or AWQ can migrate outlier difficulty from activations to weights.

Expect these follow-ups

• How does GPTQ differ from standard PTQ for large language models?
• What hardware constraints drive the choice between INT8 and INT4 quantization?

As asked

Explain how the KV cache works in transformer inference, how it grows with sequence length and batch size, and what tradeoffs you make when deciding how large to allow it.

Sample answer outline

The KV cache stores computed key and value tensors from past tokens so they are not recomputed on each generation step. Its memory footprint scales as 2 x num_layers x num_heads x head_dim x sequence_length x batch_size x bytes_per_element. For large batches or long sequences this dominates GPU memory. Techniques like PagedAttention in vLLM manage KV cache as non-contiguous pages, improving utilization. Flash attention reduces the memory footprint during prefill but does not affect cached inference.

Expect these follow-ups

• How does continuous batching in vLLM differ from static batching for throughput?
• What is multi-query attention and how does it reduce KV cache size?

As asked

Describe how you would design a model registry and promotion workflow so that moving a model from staging to production is safe and auditable, and rollback is fast.

Sample answer outline

A model registry stores versioned artifacts with metadata like training run ID, dataset version, and evaluation metrics. Promotion gates require passing offline eval thresholds and a shadow deployment or A/B test in staging. Rollback is fast when the serving infrastructure reads a model version pointer rather than baking the model into an image, so changing the pointer and restarting serves the previous version in seconds. MLflow, Vertex AI Model Registry, or SageMaker Model Registry all support this pattern.

Expect these follow-ups

• How do you handle model lineage when a model was fine-tuned from a base that itself has multiple versions?
• What would trigger an automatic rollback and how do you avoid flapping?

As asked

Your offline AUC improved by 1.5 points but the A/B test showed no significant lift in the business metric. What are the most likely explanations and how do you debug this?

Sample answer outline

The most common causes are evaluation data leakage or label bias that does not represent live traffic, positional bias in offline labels from a system that serves them differently online, or the metric gap being too small to manifest as a detectable business lift given the A/B test power. Check the offline eval pipeline for temporal leakage, verify label construction matches online behavior, and run a power analysis to confirm the test duration was sufficient to detect the expected lift.

Expect these follow-ups

• How would you construct a held-out eval set that better predicts online performance?
• When would you ship a model that improves AUC but shows neutral A/B results?

As asked

Explain exactly what gradient checkpointing does to memory during training and what you pay in compute to get those savings.

Sample answer outline

Gradient checkpointing trades memory for compute by not storing all intermediate activations during the forward pass. During backward, it recomputes the discarded activations from the nearest checkpoint. This reduces activation memory from O(num_layers) to O(sqrt(num_layers)) with the standard uniform checkpoint strategy, at the cost of approximately 33 percent more compute. It is most useful when activation memory is the binding constraint, not parameters or optimizer state.

Expect these follow-ups

• How do you decide which layers to checkpoint versus which to keep materialized?
• How does selective activation recomputation in Megatron-LM differ from standard gradient checkpointing?

As asked

When running inference with ONNX Runtime, how do you choose and configure an execution provider, and what graph optimizations does ORT apply automatically?

Sample answer outline

ONNX Runtime selects an execution provider based on hardware: CUDA EP for NVIDIA GPUs, TensorRT EP for further optimization, CoreML EP for Apple Silicon, and CPU EP as the fallback. ORT applies graph optimizations in passes: basic optimizations like constant folding and dead node elimination, extended optimizations like operator fusion, and layout optimizations. TensorRT EP provides additional kernel auto-tuning for the specific GPU. Setting opt_level and enabling memory arena can significantly affect throughput.

Expect these follow-ups

• What is the difference between ORT's TensorRT EP and running TensorRT directly?
• How do you profile an ORT session to find which nodes are taking the most time?

As asked

What does idempotency mean in the context of a daily feature engineering DAG, and what concrete design decisions make a pipeline idempotent versus fragile to reruns?

Sample answer outline

An idempotent DAG produces identical output if run multiple times for the same logical date. Concretely, write output to a partition keyed on the logical execution date, overwrite rather than append, and avoid using wall-clock time inside tasks. Use Airflow's ds macro for date context. If a task reads from a mutable source like a database table, snapshot it at the start of the run so reruns see the same input. Avoid deleting and recreating tables when an INSERT OVERWRITE partition is sufficient.

Expect these follow-ups

• How do you handle backfills for a feature pipeline that depends on a model output from a previous day?
• What is the difference between execution_date and data_interval_start in Airflow 2?

As asked

Your recommendation model has a p50 latency of 20ms but a p99 of 400ms. What are the most common causes and how do you systematically reduce p99 without hurting throughput?

Sample answer outline

P99 tail latency spikes often come from garbage collection pauses in JVM-based serving stacks, CUDA kernel launch overhead under bursty load, dynamic batching queue stalls, or CPU-GPU memory transfer jitter. Profile with percentile histograms at each layer: load balancer, serving framework, and kernel. Fixes include pinning memory for async transfers, setting a stricter max_queue_delay, using streaming batch inference with a fixed window, and moving the serving runtime to C++ or Rust to eliminate GC pauses.

Expect these follow-ups

• How do you load test a serving endpoint to reproduce p99 conditions before launch?
• What changes would you make specifically to reduce tail latency in TorchServe?

As asked

Describe how you use DVC to version datasets and link them to model training runs. What breaks when someone skips versioning and how have you seen that play out?

Sample answer outline

DVC tracks dataset files by content hash and stores the actual data in a remote like S3 or GCS, committing only the .dvc pointer files to Git. This means a Git tag on a training commit captures the exact dataset used. When versioning is skipped, teams lose the ability to reproduce a model, debug performance regressions, or audit what data a production model was trained on. A common failure mode is a data pipeline silently overwriting the training parquet files after the model ships.

Expect these follow-ups

• How do you integrate DVC with an experiment tracker like MLflow or Weights and Biases?
• What is the tradeoff between DVC and a purpose-built feature store for dataset versioning?

As asked

Your Spark feature pipeline is slow because a few user IDs generate millions of events while most generate dozens. How do you address the data skew?

Sample answer outline

Identify the skewed keys with a groupBy count and sort descending. Solutions include salting the skewed keys by appending a random suffix and aggregating in two stages, using broadcast joins instead of shuffle joins when one side is small, or applying AQE's skew join optimization in Spark 3 which automatically splits skewed partitions. For feature aggregation specifically, you can precompute heavy-hitter user features in a separate job that reads a filtered partition.

Expect these follow-ups

• How does Spark's Adaptive Query Execution handle skew compared to manual salting?
• What monitoring would you add to detect when a new high-traffic entity joins the dataset and breaks the pipeline?

As asked

Describe exactly how you would implement a shadow deployment for a new ranking model. What data do you collect, how long do you shadow, and what signals trigger a full promotion?

Sample answer outline

In shadow mode, live traffic is duplicated and sent to the new model, but only the existing model's predictions are served. Both models' predictions are logged with the same request context. You compare prediction distributions, latency profiles, and resource consumption. After collecting statistically meaningful volume, typically a few days to a week, you verify the new model's offline metrics on live traffic and check for any edge-case crashes. Promotion is triggered when metrics pass thresholds and no anomalies appear in the shadow period.

Expect these follow-ups

• How do you handle shadow deployment when the new model is significantly more expensive to run?
• What would you do if the shadow model frequently disagrees with the live model but you do not yet have online labels to judge which is better?

As asked

How do you write tests for an ML training pipeline? What do you test with unit tests, integration tests, and end-to-end smoke tests?

Sample answer outline

Unit tests cover individual transforms and feature logic with hand-crafted inputs and known outputs, catching regressions in preprocessing logic. Integration tests run the pipeline on a tiny synthetic dataset to verify the full execution graph, checking shapes, dtypes, and that loss decreases over a handful of steps. End-to-end smoke tests run a short training job on a small data slice and assert that final eval metrics are within a reasonable range. You also want tests that catch training-serving skew by running the same feature transform through both the training and serving code paths.

Expect these follow-ups

• How do you keep smoke tests fast enough to run in CI?
• What would you add to a test suite specifically to catch silent data corruption in feature engineering?

As asked

You need to fine-tune a 7B parameter model across 4 GPUs using Hugging Face Accelerate. Walk me through the accelerate config and what to watch for to avoid OOM errors.

Sample answer outline

Run accelerate config to generate a config file specifying multi-GPU with DDP or DeepSpeed backend, mixed precision, and gradient accumulation steps. For a 7B model, DeepSpeed ZeRO stage 2 or 3 is usually needed to avoid OOM. Monitor per-GPU memory with torch.cuda.memory_reserved() at the first step. Common OOM causes are the optimizer state (8 bytes per param in Adam), activations from large batch sizes, and the tokenizer expanding sequence length beyond expectation.

Expect these follow-ups

• When would you switch from Accelerate's DDP path to using DeepSpeed directly?
• How do you debug an OOM that only happens on step 2 but not step 1?

As asked

Walk me through how LoRA works mathematically, how you choose the rank, and how you decide which weight matrices to apply it to.

Sample answer outline

LoRA freezes the pretrained weights W and adds a low-rank decomposition delta_W = A * B where A is (d, r) and B is (r, k) with r much smaller than d and k. The rank r controls expressivity versus parameter count; typical values are 4 to 64 with 8 or 16 being a common default. Higher rank helps on tasks far from pretraining but risks overfitting on small datasets. Applying LoRA to attention query and value projections is the standard starting point; adding it to the MLP layers or all linear layers helps for more task-specific adaptation.

Expect these follow-ups

• How does QLoRA differ from LoRA and what does the quantization step add?
• How do you merge LoRA adapters back into the base model for deployment?

As asked

Explain how continuous batching works in a system like vLLM, how it differs from static batching, and what specific GPU utilization problem it solves.

Sample answer outline

Static batching groups a fixed set of requests and waits for all to finish before starting the next batch, so a long request blocks short ones from completing and the GPU idles while waiting. Continuous batching inserts new requests into the batch as soon as a slot frees up from a completed sequence. This dramatically improves GPU utilization when request lengths vary widely, which is typical in production. vLLM's PagedAttention manages the KV cache as dynamic pages so arriving and departing sequences do not require contiguous memory.

Expect these follow-ups

• How does continuous batching interact with sampling parameters that differ per request, like temperature?
• What throughput metric would you report to compare a static-batching versus continuous-batching serving system?

As asked

You are serving a pipeline that runs a preprocessing step, a main model, and a postprocessing step. How do you configure this in Triton, and how does dynamic batching interact with ensemble scheduling?

Sample answer outline

Triton ensemble models define a pipeline graph in config.pbtxt where each step is a model with named inputs and outputs wired together. Dynamic batching is configured per model instance with preferred_batch_size and max_queue_delay_microseconds. In an ensemble, requests flow through sequentially so dynamic batching applies to each constituent model independently. The preprocessing and postprocessing steps often use the Python or BLS backend to handle arbitrary logic, while the main model uses the TensorRT or ONNX backend for peak throughput.

Expect these follow-ups

• How do you measure per-step latency in a Triton ensemble?
• What is the Business Logic Scripting backend and when is it better than an ensemble?