AI infrastructure engineer role-specific interview questions

As asked

Walk me through the three sharding strategies in PyTorch FSDP: FULL_SHARD, SHARD_GRAD_OP, and NO_SHARD. When would you pick each one, and what are the memory and communication tradeoffs at scale?

Sample answer outline

FULL_SHARD shards parameters, gradients, and optimizer state across all ranks, giving the lowest peak memory but requiring two all-gather passes per forward and one reduce-scatter per backward. SHARD_GRAD_OP keeps parameters unsharded during the forward pass, trading memory for fewer collective calls. NO_SHARD is DDP. A strong answer covers the tradeoff between peak GPU memory and all-reduce/all-gather volume, and explains when gradient checkpointing interacts with FSDP prefetching.

Expect these follow-ups

• How does FSDP's auto_wrap_policy interact with transformer layer boundaries, and what happens if you wrap at the wrong granularity?
• How do you profile the all-gather overhead in FSDP and what tools do you use?

As asked

You have a training loop where you want to overlap CUDA kernel execution with CPU-side data preprocessing. Walk me through how you'd structure streams and events to do this correctly, and what bugs commonly appear when developers get the synchronization wrong.

Sample answer outline

Create a separate CUDA stream for compute and another for data transfer; use cudaEvent_t to signal when a transfer is complete before the compute stream proceeds. Common bugs: forgetting to record the event on the correct stream, calling cudaStreamSynchronize on the wrong stream blocking the CPU, and assuming the default stream is non-blocking. A strong answer also covers how torch.cuda.Stream and stream context managers map onto these primitives.

Expect these follow-ups

• What is the difference between cudaStreamSynchronize and cudaDeviceSynchronize in terms of performance impact?

As asked

Describe what DeepSpeed ZeRO eliminates at each stage. Given a 70B parameter model, how would you estimate peak memory per GPU at ZeRO-3 with mixed precision, and what communication overhead does ZeRO-3 add over ZeRO-1?

Sample answer outline

ZeRO-1 shards optimizer states, ZeRO-2 adds gradient sharding, ZeRO-3 shards parameters as well. For a 70B model in bf16 with Adam, ZeRO-3 on N GPUs keeps roughly (70B * 16 bytes) / N bytes of optimizer state plus activation memory. ZeRO-3 adds parameter gather collectives on every forward/backward pass, roughly doubling communication volume compared to ZeRO-1. Mention ZeRO-Infinity for NVMe offload and the latency implications.

Expect these follow-ups

• How does ZeRO-3 interact with activation checkpointing, and can you combine them safely?

As asked

How does activation (gradient) checkpointing reduce memory, and what is the theoretical recompute cost as a fraction of total training FLOPs? How do you decide which layers to checkpoint in a transformer, and what goes wrong if you checkpoint too aggressively?

Sample answer outline

Checkpointing stores only a subset of activations and recomputes the rest during the backward pass, reducing activation memory from O(L) to O(sqrt(L)) with optimal placement. The recompute adds roughly one extra forward pass worth of FLOPs. In transformers, checkpointing at the attention or MLP boundary is standard. Checkpointing too aggressively increases recompute overhead disproportionately and can also interfere with FSDP's prefetching if the recompute triggers unexpected collective calls.

Expect these follow-ups

• How does selective activation recomputation in Megatron-LM differ from full per-layer checkpointing?

As asked

What is MFU (model FLOP utilization), how do you calculate it for a GPT-style transformer training run, and what is a realistic target MFU on A100s with bf16? What are the main culprits when MFU is significantly below target?

Sample answer outline

MFU = achieved_FLOPS / peak_theoretical_FLOPS. For a transformer with L layers, H heads, d_model, and sequence length S, the forward FLOPs per token are roughly 6 * N (N = params) for a standard pass. On A100 bf16 peak is 312 TFLOPS; a well-tuned run hits 35 to 55% MFU. Common culprits below target: pipeline bubbles, data-loading stalls, small batch sizes causing underutilization, memory-bound operations in attention, and excessive host-device synchronization.

Expect these follow-ups

• How would you distinguish between a communication bottleneck and a compute bottleneck when MFU is low?

As asked

For a 13B parameter model served with a 4096-token context window, walk me through how you'd estimate the KV cache memory per request, how PagedAttention (vLLM) manages this differently from a static allocation, and what happens under memory pressure.

Sample answer outline

KV cache per request is 2 * L * H * d_head * S bytes (two tensors per layer, one for K one for V). For 13B with 40 layers, 40 heads, d_head=128, 4096 tokens in fp16 that is 2 * 40 * 40 * 128 * 4096 * 2 bytes, roughly 3.1 GB per request. PagedAttention allocates fixed-size blocks lazily, so GPU memory is not reserved upfront and can be shared across requests via copy-on-write for shared prefixes. Under memory pressure, vLLM will preempt lower-priority sequences and swap their KV blocks to CPU, recomputing them when rescheduled.

Expect these follow-ups

• How does prefix caching in vLLM work, and what workloads benefit most from it?

As asked

You are running a 70B model training job with checkpoints needed every 30 minutes. Naive checkpointing blocks training for several minutes. What techniques reduce checkpoint latency, and how do you handle partial or corrupt checkpoints if a node fails mid-save?

Sample answer outline

Async checkpointing copies tensors to pinned CPU memory on a background thread without blocking the GPU, then flushes to storage independently. Sharded checkpoints (each rank writes its own shard) parallelize I/O across nodes. For corruption safety, write to a temp path and atomic-rename only on completion; keep N-1 and N checkpoints for rollback. Mention PyTorch Distributed Checkpoint (DCP), fsync semantics on distributed filesystems like Lustre, and the difference between saving optimizer state (large) vs. model weights only.

Expect these follow-ups

• How does PyTorch DCP's save_planner differ from the older torch.save for distributed checkpoints?

As asked

Explain continuous batching and why it improves GPU utilization compared to static batching for autoregressive LLM inference. What is iteration-level scheduling, and what are its failure modes at high concurrency?

Sample answer outline

Static batching waits for all requests in a batch to finish before adding new ones, leaving GPUs idle when short sequences finish early. Continuous batching (iteration-level scheduling) inserts new requests at each decode step as slots free up, improving GPU utilization significantly. Failure modes at high concurrency include KV cache exhaustion causing preemption storms, head-of-line blocking when one very long sequence blocks slot reuse, and priority inversion if FCFS scheduling is used without admission control.

Expect these follow-ups

• How does chunked prefill interact with continuous batching to reduce time-to-first-token variance?

As asked

Why is loss scaling necessary in fp16 mixed-precision training but not usually in bf16? How does dynamic loss scaling work, what does it do when it detects a gradient overflow, and what hyperparameters control its behavior?

Sample answer outline

fp16 has a smaller dynamic range than bf16; small gradients underflow to zero. Loss scaling multiplies the loss by a large factor before backprop and divides gradients afterward, keeping values in fp16 range. Dynamic scaling starts with a high scale factor, checks for inf/nan gradients each step, halves the scale on overflow and skips that update, and increases the scale gradually when no overflow is detected. bf16 has a larger exponent range that covers typical gradient magnitudes without scaling. Key hyperparameters: initial_scale, scale_factor, scale_window.

Expect these follow-ups

• What happens to the training step counter when a gradient overflow is detected and the update is skipped?

As asked

Three days into a 10-day training run, a node fails and the job crashes. Your last checkpoint is from 12 hours ago. How do you decide whether to restart from the checkpoint, investigate the failed node, or do something else first?

Sample answer outline

First, determine whether the node failure is isolated or systemic (check other nodes in the cluster, look for hardware alerts). If isolated: drain the failed node, replace it with a healthy spare, and restart the job from the 12-hour-old checkpoint. Assess whether the 12-hour loss is acceptable or whether more frequent checkpointing needs to be implemented immediately. Communicate ETA to the team. If systemic: pause all jobs, investigate root cause (network flap, storage issue, DCGM alerts) before restarting to avoid repeated failures.

Expect these follow-ups

• After this incident, what checkpoint frequency would you set and how would you justify it in terms of compute cost?

As asked

Your LLM serving SLA is 2s p95 TTFT. At 3pm on a Tuesday, p95 TTFT spikes to 8s with no change in traffic volume. What is your investigation sequence and what hypotheses do you rule out in what order?

Sample answer outline

Check whether the spike correlates with any deployment, config change, or external event (e.g., a long-context request flooding the queue). Rule out in order: traffic pattern change (longer prompts increase prefill time), a misbehaving single replica dragging down p95, KV cache exhaustion causing preemption/recompute, a noisy neighbor on shared GPU hardware, or a thermal throttle event on one node. Tools: check p50 vs p95 divergence (if p50 is fine, it is a tail issue), look at per-replica metrics in Grafana, check GPU utilization and memory stats on each node.

Expect these follow-ups

• If you found that one replica was serving abnormally long requests that were blocking the queue, how would you handle it without dropping those requests?

As asked

Your LLM inference bill has doubled in a month while traffic grew by only 20 percent. How would you investigate and control cost?

Sample answer outline

Break cost down by model, route, tenant, prompt tokens, completion tokens, cache hit rate, retries, and hardware utilisation. Look for longer prompts, lower batching efficiency, runaway agent loops, increased fallback to larger models, and retry storms after provider or gateway errors. Strong answers propose guardrails such as token budgets, model routing, semantic or prefix caching, prompt compression, and per-tenant cost dashboards. The candidate should avoid blind quality regression by pairing cost changes with evals and customer-impact checks. A common miss is treating unit cost as fixed when application behaviour often drives the bill.

Expect these follow-ups

• How do you detect an agent loop from billing data?
• When would you route a request to a smaller model?
• What cost metric should product managers see weekly?

As asked

Explain how NCCL implements ring-allreduce. What determines the chunk size it uses, how does it pipeline reduce-scatter and all-gather, and what breaks down when a node in the ring has a slower NIC?

Sample answer outline

Ring-allreduce splits tensors into chunks, passes them around a ring doing partial reductions, then redistributes. NCCL chooses chunk size based on buffer size, algorithm threshold (tree vs. ring), and topology. Pipelining overlaps the reduce-scatter phase with data transmission. A slow NIC creates a bottleneck that backs up the entire ring, so the answer should mention NCCL_ALGO, topology detection via nvlink vs. PCIe vs. IB, and how NCCL can fall back to tree algorithms on heterogeneous clusters.

Expect these follow-ups

• When would NCCL prefer tree-based collectives over ring for allreduce?

As asked

Explain how GPUDirect RDMA works and why it matters for GPU cluster training. What does it bypass, what kernel or driver components must be in place, and what are the failure modes you'd check when GPUDirect is enabled but throughput is still poor?

Sample answer outline

GPUDirect RDMA allows the HCA to DMA directly into GPU memory, bypassing the CPU and system memory. It requires the nvidia-peermem kernel module (or nv_peer_mem), a compatible HCA, and PCIe topology where GPU and HCA share the same root complex. Poor throughput with GPUDirect enabled often traces to IOMMU misconfiguration, missing hugepages, firmware mismatches, or NUMA topology causing the DMA to cross a QPI/UPI hop. Strong answers also mention ibv_devinfo and perftest utilities for diagnosis.

Expect these follow-ups

• How would you verify GPUDirect RDMA is actually active and not silently falling back to bounce buffers?

As asked

Your custom CUDA kernel is running at 40% theoretical occupancy. Walk me through how you'd diagnose the root cause using Nsight Compute, and what changes to block size, shared memory allocation, or register pressure would you try first?

Sample answer outline

Low occupancy is usually caused by too many registers per thread, excessive shared memory per block, or a block size that does not divide evenly into warps of 32. Nsight Compute's occupancy section shows the limiting factor. Reducing register pressure via __launch_bounds__, splitting shared memory across multiple kernels, or picking a block size closer to a multiple of the warp size are the first levers. A strong answer discusses the roofline model and distinguishes between compute-bound and memory-bound kernels.

Expect these follow-ups

• What is the difference between theoretical occupancy and achieved occupancy, and which one actually matters for throughput?

As asked

Describe how tensor parallelism, pipeline parallelism, and data parallelism are combined in Megatron-LM's 3D parallelism. How do you decide the TP, PP, and DP degrees for a 175B model on a cluster of 1024 A100s?

Sample answer outline

TP splits individual layer matrices across GPUs within a node (NVLink bandwidth), PP assigns layers to pipeline stages across nodes, and DP replicates the full model across pipeline groups. For 175B on 1024 A100s, a common choice is TP=8 (one node), PP=8 or 16 (pipeline depth tied to micro-batch count for good bubble fraction), and DP filling the rest. The answer should discuss pipeline bubble fraction (1 - 1/PP), micro-batch size for bubble hiding, and why TP degree is usually bounded by NVLink topology.

Expect these follow-ups

• How do you minimize the pipeline bubble, and what is the 1F1B schedule in Megatron?

As asked

Explain the core idea behind FlashAttention's tiling approach. Why does it reduce HBM reads and writes compared to standard attention, and what is the IO complexity saving? How does FlashAttention 2 improve on the original?

Sample answer outline

Standard attention materializes the N x N attention matrix in HBM; FlashAttention tiles Q, K, V into SRAM blocks and computes softmax incrementally using an online normalization trick, never writing the full N x N matrix to HBM. IO complexity drops from O(N^2) to O(N^2 / M) where M is SRAM size. FlashAttention 2 improves parallelism by splitting work across the sequence dimension for better warp utilization and reduces the number of non-matmul operations, achieving roughly 2x throughput improvement.

Expect these follow-ups

• What is the causal mask handling in FlashAttention and how does it affect the tiling schedule?

As asked

Distributed training jobs need all pods to start simultaneously or the job stalls. How does gang scheduling solve this, and how would you implement it in a Kubernetes cluster using Volcano or Kueue? What failure modes do you watch for?

Sample answer outline

Gang scheduling holds all pods in a pending state until the full set can be scheduled atomically, preventing partial allocations that hold GPUs idle. In Volcano this is the PodGroup primitive; in Kueue it is the ClusterQueue workload admission. Failure modes include: deadlock when two jobs each hold partial resources, slow pod startup causing NCCL timeout, and node preemption mid-job. A strong answer also covers minMember thresholds and checkpoint-on-preemption strategies.

Expect these follow-ups

• How do you handle node failure mid-training in a gang-scheduled job, and what is your checkpoint strategy?

As asked

Compare PTQ and QAT for deploying a large language model at INT8 or INT4. When does PTQ with GPTQ or AWQ give you acceptable accuracy, and when do you need QAT? What hardware-level considerations drive the choice between INT8 and INT4?

Sample answer outline

PTQ (GPTQ, AWQ) calibrates quantization parameters on a small dataset without retraining; it works well for large models where individual weight errors average out, and is preferred for deployment speed. QAT is needed when PTQ degrades accuracy beyond tolerance, especially for smaller models or tasks sensitive to outlier activations. INT8 is well-supported by Tensor Cores on Ampere/Hopper; INT4 doubles throughput but requires kernel support (e.g., CUTLASS, bitsandbytes) and can have larger accuracy gaps. Hardware SMs must support the target dtype natively.

Expect these follow-ups

• How does AWQ differ from GPTQ in terms of which weights it prioritizes for accurate quantization?

As asked

You have a server with 8 A100s connected via NVLink and an additional 4 servers connected via InfiniBand. How does the collective communication strategy change for an allreduce that spans all 32 GPUs compared to one that is intra-node only?

Sample answer outline

Intra-node allreduce uses NVLink ring or tree with very high bandwidth (600 GB/s per A100). Cross-node allreduce must use InfiniBand, which at 200 Gb/s HDR is a bandwidth bottleneck. NCCL handles this hierarchically: intra-node reduce-scatter via NVLink, then inter-node allreduce via IB, then intra-node broadcast. Optimal strategy depends on message size; for small tensors, latency dominates and a tree over IB is preferred. A strong answer also discusses NCCL_IB_HCA pinning and ensuring affinity between the GPU and HCA on the same PCIe root complex.

Expect these follow-ups

• How would you use nvidia-smi topo and ibstat to verify the topology assumptions before a large run?

As asked

Walk me through how you'd use Nsight Systems to profile a multi-GPU training run and identify whether the bottleneck is compute, memory bandwidth, or inter-GPU communication. What markers or annotations do you add to the training code to make the trace interpretable?

Sample answer outline

Run nsys profile with CUDA and NVTX tracing enabled; the timeline shows kernel gaps (compute idle), cudaMemcpy operations (data movement), and NCCL calls (communication). Add torch.cuda.nvtx.range_push/pop or PyTorch Profiler annotations around forward, backward, and optimizer steps. Identify the bottleneck by comparing kernel duration to theoretical compute time, and checking if NCCL calls show up on the critical path. A strong answer discusses the difference between overlapping and non-overlapping communication.

Expect these follow-ups

• What would you look for in the Nsight Systems trace to confirm that gradient communication is successfully overlapping with backward compute?

As asked

In Megatron-style tensor parallelism, how are linear layers split into column-parallel and row-parallel variants? What collective operations are needed, and why does this layout minimize communication to one allreduce per transformer sublayer?

Sample answer outline

Column-parallel splits the weight matrix along the output dimension; each rank computes a partial output, requiring no communication between ranks in the forward pass for the column layer, but an allreduce after the row-parallel layer that follows. Row-parallel splits the input dimension; the partial products are summed via allreduce. The MLP block is column-then-row, so one allreduce completes the block. For self-attention, QKV projection is column-parallel and the output projection is row-parallel. Strong answers note that the allreduce can be replaced with reduce-scatter + all-gather (sequence parallelism) to also shard activations.

Expect these follow-ups

• What is sequence parallelism, and how does it complement tensor parallelism to reduce activation memory?

As asked

Describe how PyTorch's caching memory allocator works. What causes GPU memory fragmentation during training, how do you detect it, and what knobs in the allocator or training loop can reduce out-of-memory errors caused by fragmentation rather than insufficient total memory?

Sample answer outline

PyTorch maintains a pool of cached blocks to avoid repeated cudaMalloc calls. Fragmentation occurs when many differently-sized tensors are allocated and freed, leaving small non-contiguous free blocks. torch.cuda.memory_stats() shows fragmentation ratio; torch.cuda.memory_snapshot() gives a detailed view. Mitigations include setting PYTORCH_CUDA_ALLOC_CONF with max_split_size_mb to limit block splitting, calling torch.cuda.empty_cache() strategically, and ensuring batch size and sequence length are stable across steps to reduce allocation churn.

Expect these follow-ups

• When would you use expandable_segments in the PyTorch CUDA allocator, and what problem does it solve?

As asked

Explain how MIG (Multi-Instance GPU) works on A100/H100 and when you'd use it for inference serving. What are the isolation guarantees, what are the performance tradeoffs compared to running a single large model, and what operational complexity does it add?

Sample answer outline

MIG partitions the GPU into up to 7 compute instances with dedicated SM slices, memory partitions, and bandwidth, providing hard isolation between tenants. It is useful when you have many small models that each do not fill a full A100, improving utilization. Tradeoffs: each instance has less total memory, so large models will not fit; the partition size is fixed at the GPC granularity; reconfiguring instances requires taking the GPU offline. Operationally, MIG profiles must be declared at the Kubernetes device plugin level and managed as part of node provisioning.

Expect these follow-ups

• How does MIG interact with NCCL for collective operations that need to span GPU instances?

As asked

Explain how speculative decoding accelerates autoregressive inference. What determines the acceptance rate of draft tokens, how do you select the draft model size, and under what conditions does speculative decoding actually slow things down?

Sample answer outline

The draft model generates K tokens cheaply; the target model verifies all K in one forward pass via parallel scoring. Accepted tokens pass with probability min(1, p_target / p_draft). Acceptance rate depends on how well the draft distribution matches the target. Draft model should be 10 to 50x smaller than the target for a net speedup; too large and verification dominates. It slows things down when acceptance rate is very low (divergent distributions) or when the target model is compute-bound and parallel verification does not amortize the draft overhead.

Expect these follow-ups

• How does self-speculative decoding (using the target model's own early exit layers) differ from using a separate draft model?