AI research engineer role-specific interview questions

As asked

Standard self-attention has O(n^2) time and memory complexity in sequence length. Walk me through exactly where that quadratic term comes from, and describe at least two concrete algorithmic approaches that reduce it.

Sample answer outline

The quadratic cost comes from computing all n^2 query-key dot products to form the attention matrix, then multiplying by values. FlashAttention avoids materializing the full n x n matrix by tiling the computation in SRAM and fusing the softmax, achieving O(n) memory while keeping exact attention. Sparse attention methods (Longformer sliding window, BigBird random plus global patterns) reduce compute by restricting which token pairs attend to each other, trading exact for approximate attention.

Expect these follow-ups

• How does FlashAttention 2 differ from the original FlashAttention in terms of the tiling strategy?
• What is the causal masking change required in FlashAttention for decoder-only models?

As asked

Explain what the key-value cache is during autoregressive decoding. Given a model with 32 layers, 32 attention heads, head dimension 128, and sequence length 4096, how much GPU memory does the KV cache consume per request in float16?

Sample answer outline

The KV cache stores the key and value projections for all past tokens so they do not need to be recomputed at each new generation step. Memory is 2 (K and V) times layers times heads times head_dim times sequence_length times bytes_per_element: 2 * 32 * 32 * 128 * 4096 * 2 bytes = 2 GB per request in float16. At scale this becomes the primary bottleneck, motivating techniques like grouped-query attention (GQA) which reduces the number of KV heads, and PagedAttention which manages the cache in non-contiguous blocks.

Expect these follow-ups

• How does grouped-query attention differ from multi-query attention in terms of the KV cache size reduction?
• What does PagedAttention in vLLM do differently from a naive KV cache allocation?

As asked

You are training a large transformer and run out of GPU memory. A colleague suggests gradient checkpointing. Explain what it does, what the exact tradeoff is, and where you would place checkpoints in the network to get the best tradeoff.

Sample answer outline

Gradient checkpointing discards intermediate activations during the forward pass instead of storing them, then recomputes them during backpropagation when they are needed for gradients. The cost is one extra forward pass per checkpointed segment, so runtime increases by roughly 30-35% while memory drops significantly. The optimal placement recomputes about every sqrt(n) layers, which gives O(sqrt(n)) memory instead of O(n), balancing memory savings against recompute time. In practice checkpointing is placed at transformer block boundaries because each block's input is a clean recompute boundary.

Expect these follow-ups

• How does selective activation checkpointing differ from checkpointing every block?
• What happens to the memory profile if you checkpoint too aggressively, for example every single layer?

As asked

Describe exactly what AdamW stores per parameter, how the bias correction terms work in the first few steps, and why decoupled weight decay in AdamW is different from L2 regularization in plain Adam.

Sample answer outline

AdamW stores two moment estimates per parameter: m (first moment, exponential moving average of gradients) and v (second moment, EMA of squared gradients). Bias correction divides by (1 - beta1^t) and (1 - beta2^t) to compensate for the zero initialization of m and v, which would otherwise undershoot early steps. In plain Adam, L2 regularization adds lambda * param to the gradient before updating moments, making the effective weight decay scale with the adaptive learning rate. AdamW decouples this by applying weight decay directly to the parameter after the gradient update, so all parameters decay at the same rate regardless of gradient magnitude.

Expect these follow-ups

• Why does large learning rate combined with small beta2 sometimes cause training instability?
• What is the memory cost of AdamW for a 7B parameter model in float32 optimizer states?

As asked

Summarize the Chinchilla scaling law finding. What did Hoffmann et al. show that Kaplan et al. got wrong, and what is the practical implication for deciding how many tokens to train on given a compute budget?

Sample answer outline

Kaplan et al. found that model size matters more than data and recommended scaling parameters faster than tokens. Chinchilla showed both should scale roughly equally: for a given compute budget C (in FLOPs), the optimal model has N parameters and is trained on D tokens where N and D are roughly equal proportions of C divided by 6. The implication is that GPT-3 at 175B was significantly undertrained, and a 70B model trained on 1.4T tokens (Chinchilla) matches or beats it. Practically, this shifted the community toward training smaller models for longer, which also produces better inference economics.

Expect these follow-ups

• Why does the compute-optimal prescription differ from what you would choose if inference cost matters?
• How does the Chinchilla analysis handle data quality, and what happens when you run out of high-quality tokens?

As asked

Explain what LoRA does mathematically, how you choose the rank r in practice, and what kinds of weight updates LoRA cannot represent that full fine-tuning can.

Sample answer outline

LoRA freezes the pre-trained weight matrix W and learns a low-rank update delta_W = A * B, where A is d x r and B is r x d_out, with r much smaller than d. At inference the update is merged back into W with no added latency. Rank is typically chosen between 4 and 64: higher rank allows more expressive updates but costs more memory and training time. LoRA is constrained to updates that lie in a low-rank subspace, so it struggles to represent large distributed changes to all singular directions of the weight matrix, which full fine-tuning can express. For large distribution shifts, higher ranks or full fine-tuning may be necessary.

Expect these follow-ups

• Why is the B matrix initialized to zero and A initialized randomly, and what happens if you initialize both to zero?
• How does QLoRA combine LoRA with quantization, and what are the gradient flow implications?

As asked

Describe the typical learning rate schedule used in large language model pretraining: why you need warmup, what cosine decay does to the schedule, and what the practical consequences are of setting warmup steps too short or the peak learning rate too high.

Sample answer outline

Warmup linearly increases the learning rate from near zero to the peak over the first few hundred to thousand steps, because the optimizer moments are uninitialized and large early steps cause the model weights to take huge destructive updates. Cosine decay then smoothly reduces the learning rate from peak to a small floor (often 10% of peak) over the training horizon. Too short warmup causes early loss spikes because moment estimates are poor. Too high peak learning rate causes instability even after warmup. The learning rate should roughly scale as the inverse square root of the batch size when adjusting across compute configurations.

Expect these follow-ups

• What is the impact on final model quality if you restart training from a checkpoint using a warmup that is too aggressive?
• How does the Warmup-Stable-Decay schedule from MiniCPM differ from cosine decay?

As asked

Compare rotary positional embeddings (RoPE) to learned absolute positional embeddings. What mathematical property makes RoPE attractive for length generalization, and what is actually happening geometrically when you apply the rotation?

Sample answer outline

Absolute positional embeddings add a fixed or learned vector per position index, making extrapolation to unseen lengths unreliable because the model never saw those indices during training. RoPE encodes position by rotating query and key vectors in 2D subspaces using a position-dependent angle, so the dot product between two tokens depends only on their relative distance rather than absolute index. This relative nature means the model can generalize further if the rotation angles are chosen carefully, and YaRN or dynamic NTK scaling can extend the effective context window post-training.

Expect these follow-ups

• Why does vanilla RoPE still struggle beyond the training context length even though it encodes relative positions?
• What does the NTK-aware scaling trick actually modify?

As asked

LLaMA and many modern LLMs switched from LayerNorm to RMSNorm. What does RMSNorm actually compute, what does it drop compared to LayerNorm, and why does that matter for training speed and stability?

Sample answer outline

LayerNorm computes mean and variance over the feature dimension, subtracts the mean (re-centering), divides by standard deviation (rescaling), then applies learned scale and bias parameters. RMSNorm drops the mean subtraction entirely and normalizes by root-mean-square of the activations, keeping only the learned scale. This removes 7-8% of compute from the normalization step and is empirically as stable as LayerNorm because the centering term matters less than the scaling in practice. Pre-norm placement (before attention/FFN instead of after) also plays a large role in modern training stability.

Expect these follow-ups

• Where is normalization placed in a pre-norm transformer versus the original post-norm design, and which is more stable at large scale?
• What happens to training if you accidentally remove the RMSNorm scale parameter?

As asked

Your transformer training job is running smoothly at step 8,000, then the loss spikes from 2.1 to 8.5 in a single step and does not recover. What is your diagnostic process and what are the most likely causes?

Sample answer outline

First check gradient norms logged at that step: an exploding gradient typically precedes the spike and points to a bad batch, a learning rate issue, or a numerical overflow. Check if a particular data shard or corrupted example entered the loader at that step. Inspect the loss scale if using mixed precision since underflow can cause divergence. Rule out hardware failures such as a GPU ECC error by checking CUDA device status. Look at the learning rate schedule and optimizer state for an accidental warmup restart. If gradients are fine, the issue may be in the data pipeline rather than the model.

Expect these follow-ups

• How would you set up pre-emptive gradient norm logging to catch this before it becomes unrecoverable?
• What does grad_norm clipping at 1.0 actually prevent and does it hide the underlying issue?

As asked

Modern LLMs replace the original two-layer ReLU FFN with SwiGLU. What does SwiGLU compute, why is it thought to work better, and how does it change the parameter count compared to a standard FFN of the same hidden dimension?

Sample answer outline

SwiGLU computes (x W1) swish-gate times (x W2), then projects down with W3, requiring three weight matrices instead of two. The gating mechanism allows the model to modulate which parts of the intermediate representation pass through, acting as a learned feature selector. Empirically it outperforms ReLU and GELU activations on language modeling benchmarks. The extra matrix means that to match a reference FFN's compute budget you must reduce the hidden dimension by roughly 2/3, which is why LLaMA's FFN hidden size is 2/3 of 4 times the model dimension rather than the full 4x.

Expect these follow-ups

• What is the difference between SwiGLU and GeGLU, and does it matter empirically?
• Why is the FFN typically 4x the model dimension in the original transformer, and where does that number come from?

As asked

Explain the tiling strategy FlashAttention uses to avoid writing the full attention matrix to HBM. What is the role of the online softmax, and why does this give the same result as naive attention?

Sample answer outline

FlashAttention divides Q, K, V into tiles that fit in SRAM. It processes tiles in an outer loop and maintains running statistics for the online softmax: a running max m and a running normalization sum l. As each K tile is processed, it updates m and l using the log-sum-exp trick to correct previous partial sums, then accumulates the weighted values. At the end, the output for each query block is the correctly normalized sum without ever materializing the full n x n matrix in HBM. Correctness follows because the online softmax is mathematically equivalent to computing softmax over the full sequence.

Expect these follow-ups

• Why is SRAM much faster than HBM, and roughly what is the bandwidth ratio on an A100?
• How does FlashAttention 2 improve on v1, specifically in terms of work partitioning across warps?

As asked

When fine-tuning a language model with PPO against a reward model, what is reward hacking, how does the KL divergence penalty from the reference policy prevent it, and what goes wrong if the KL coefficient is set too high or too low?

Sample answer outline

Reward hacking occurs when the policy discovers high-reward behaviors that do not align with the true objective, for example generating repetitive text that the reward model over-scores. The KL penalty adds beta times KL(policy, reference) to the reward, discouraging the model from drifting too far from the supervised fine-tuned reference. If beta is too small, the policy exploits the reward model's weaknesses and degenerates. If beta is too large, the model barely moves from the reference and the reward improvement is negligible. Adaptive KL targets try to set beta dynamically to maintain a target KL divergence.

Expect these follow-ups

• How does DPO avoid the reward model and the PPO training loop entirely?
• Why is reward model overoptimization harder to detect than reward hacking on a narrow task?

As asked

You quantize a 7B LLM to INT8 using post-training quantization. Which parts of the network are most sensitive to quantization error, what metric do you use to measure quality loss, and which popular technique specifically handles outlier activations in LLMs?

Sample answer outline

Attention projection weights, especially the output projection, and the first and last layers tend to be most sensitive because they handle high-variance distributions. Perplexity on a held-out set is the standard quality metric, with task benchmarks (MMLU, HellaSwag) as secondary checks. LLM.int8() from bitsandbytes addresses the outlier activation problem: it decomposes matrix multiplications into a high-precision float16 path for the small fraction of outlier feature dimensions and INT8 for the rest, avoiding the large quantization errors caused by a few dominant activation values.

Expect these follow-ups

• How does GPTQ differ from LLM.int8() in its approach to quantization?
• What is the accuracy-latency tradeoff between INT8 and INT4 quantization on an A100 GPU?

As asked

Describe how byte-pair encoding builds a vocabulary from scratch, what the merge rules are, and give two concrete failure modes that BPE causes during model training or inference.

Sample answer outline

BPE starts with a character-level vocabulary, then iteratively merges the most frequent adjacent pair of tokens into a new token until the vocabulary reaches the target size. The merge rules are ordered and deterministic. Failure modes include: long numbers or rare words splitting into many tokens, inflating sequence length and degrading numerical reasoning; and inconsistent tokenization of the same string depending on surrounding context (for example, a word at the start of a string versus mid-sentence may tokenize differently due to space-prefix tokens in GPT-style tokenizers).

Expect these follow-ups

• Why did modern tokenizers like LLaMA's SentencePiece BPE operate at the byte level rather than the Unicode character level?
• How does tokenization affect code generation quality compared to natural language generation?

As asked

You are preparing a web-crawl dataset for pretraining a 7B model. Explain why deduplication is critical, what happens when you train on heavily duplicated data, and describe at least two deduplication strategies with their tradeoffs.

Sample answer outline

Training on duplicated data causes the model to memorize repeated passages rather than learning general patterns, increases memorization of training examples (a privacy risk), and wastes compute on redundant examples that do not add to the model's generalization. Exact deduplication uses hashing to remove identical documents and is fast but misses near-duplicates. MinHash LSH computes approximate Jaccard similarity across document n-grams, catching paraphrased or partially copied documents at the cost of higher compute and some false positives. Semantic deduplication using embedding similarity is more powerful but much more expensive at scale.

Expect these follow-ups

• How does the Dolma or RedPajama dataset handle deduplication at the trillion-token scale?
• What is the risk of over-aggressive deduplication for low-resource languages or specialized domains?

As asked

Explain the speculative decoding algorithm. How does the draft model and verifier model work together, and why is the output distribution guaranteed to match the target model's distribution exactly?

Sample answer outline

A small draft model generates k candidate tokens autoregressively, then the large target model processes all k tokens in a single parallel forward pass (since transformers can score all positions simultaneously). For each draft token, if the target model's probability is at least as high as the draft model's probability, the token is accepted. If rejected, a corrected token is sampled from the adjusted distribution to ensure the final output distribution matches the target exactly. Speedups come from amortizing the expensive target model forward pass over multiple accepted tokens, typically giving 2-3x throughput improvement.

Expect these follow-ups

• Why is speculative decoding most effective on non-batch or small-batch inference rather than large batch sizes?
• What makes a draft model a good candidate, and can the same model act as its own draft model?

As asked

You are running ablation studies comparing five architectural variants of a transformer. What do you log to Weights and Biases (or equivalent), how do you structure runs and sweeps, and what mistakes do researchers commonly make that make experiments hard to reproduce?

Sample answer outline

Log hyperparameters, loss curves (train and val), gradient norms, learning rate, hardware utilization, and sample model outputs at key checkpoints. Use tags and groups to link related runs. A sweep config defines the search space and strategy (grid, random, Bayesian). Common mistakes include logging only loss and missing gradient norms which diagnose training health, not pinning the random seed and data order so results cannot be reproduced, logging to stdout only without checkpointing config files, and overwriting the same run name rather than using unique IDs.

Expect these follow-ups

• How do you handle comparing runs that used different random seeds to separate variance from signal?
• When is Bayesian hyperparameter search worth its overhead versus random search?

As asked

Describe how token routing works in a sparse Mixture of Experts transformer. What is the auxiliary load-balancing loss, why do you need it, and what failure mode occurs without it?

Sample answer outline

Each token is routed to the top-k experts (typically k=1 or 2) by a router network that computes scores and selects the highest scoring experts per token. Without load balancing, the router collapses so that one or two experts receive nearly all tokens (expert collapse), wasting the capacity of the others. The auxiliary loss penalizes imbalanced routing by adding a term proportional to the product of average routing probability and average expert load, encouraging uniform distribution of tokens across experts. The loss coefficient must be tuned carefully: too large and routing becomes too uniform, hurting performance; too small and collapse returns.

Expect these follow-ups

• How does expert capacity factor work and what happens to tokens that exceed an expert's capacity?
• What is the difference between switch transformer routing and Mixtral's top-2 routing?

As asked

Your training throughput is lower than expected. Walk me through how you would use PyTorch profiler and NVIDIA Nsight to identify whether the bottleneck is compute-bound, memory-bandwidth-bound, or data pipeline-bound, and what signatures each case shows.

Sample answer outline

Start with torch.profiler to capture a few training steps with CUDA activity and trace stack export to TensorBoard. Look for large gaps between CUDA kernels (indicating CPU-GPU sync or data stall), long memory copy operations (data pipeline bottleneck), or back-to-back GEMM kernels with high SM utilization (compute-bound). In Nsight Compute, check arithmetic intensity versus the roofline model: if the operation is below the memory bandwidth ridge point it is memory-bandwidth-bound; above it is compute-bound. A data-pipeline bottleneck shows up as GPU idle time while waiting for the DataLoader, solvable with more workers or async prefetching.

Expect these follow-ups

• What does the GPU SM efficiency metric tell you, and when is 80% SM efficiency acceptable versus a red flag?
• How does prefetching with CUDA streams help eliminate data pipeline stalls?

As asked

Explain the key algorithmic difference between DPO and PPO for aligning LLMs with human preferences. What implicit model does DPO use, and in what training scenarios does PPO still outperform DPO?

Sample answer outline

PPO trains against an explicit reward model using reinforcement learning, requiring four models simultaneously: the policy, reference policy, reward model, and value function, making it expensive and finicky to tune. DPO reparameterizes the RLHF objective to show that the optimal policy under a KL-constrained reward maximization can be expressed directly in terms of the policy's log probability ratio versus the reference, eliminating the need for explicit reward model or RL training. PPO tends to outperform DPO on tasks requiring online exploration or when the preference data distribution is far from the SFT model's distribution, because DPO is an offline method that can overfit to the preference dataset.

Expect these follow-ups

• What is SimPO and how does it address the reference model dependency in DPO?
• How does the implicit reward in DPO relate to the Bradley-Terry preference model?

As asked

Your new 7B model gets surprisingly high scores on MMLU. Walk me through how benchmark contamination can occur, how you would test whether your model has been contaminated on MMLU specifically, and what design choices in data preparation reduce the risk.

Sample answer outline

Contamination occurs when evaluation examples appear in the pretraining corpus, either verbatim or in paraphrased form. To test, run n-gram overlap analysis between your training data and MMLU questions; if 8-gram or longer matches exist, contamination is likely. You can also perturb MMLU answers (replace correct with wrong) and check if the model's probability shifts, which it would not if it memorized rather than reasoned. Prevention involves deduplicating the training corpus against known benchmarks before training and building a held-out set of new questions for true generalization measurement.

Expect these follow-ups

• What are the limitations of n-gram overlap as a contamination detector for multiple-choice benchmarks?
• How does the GSM8K contamination controversy illustrate the problem at scale?

As asked

Explain how tensor parallelism as implemented in Megatron-LM splits the attention and FFN layers across multiple GPUs. What communication primitives are used, and what is the latency penalty compared to data parallelism?

Sample answer outline

In Megatron-style tensor parallelism, the attention head dimension is split across GPUs so each GPU handles a subset of heads, requiring an all-reduce after the output projection. The FFN intermediate layer is column-split across GPUs, and the output projection is row-split, requiring one all-reduce per FFN block. This means two all-reduces per transformer layer instead of one per model replica update in DDP. The latency penalty is significant for small batch sizes or models that are not large enough to amortize the communication overhead, which is why tensor parallelism is typically reserved for the largest models (70B+) where the model does not fit on one GPU even with FSDP.

Expect these follow-ups

• How does Megatron's sequence parallelism complement tensor parallelism in reducing activation memory?
• Why does tensor parallelism work best within a single node rather than across nodes?

As asked

During pretraining, you track both perplexity on a held-out set and accuracy on a few downstream benchmarks. Describe scenarios where perplexity decreases but downstream accuracy does not improve, and vice versa.

Sample answer outline

Perplexity measures average next-token prediction loss on a general text distribution and can decrease as the model gets better at predicting common words and syntax without improving on reasoning tasks. Downstream accuracy can plateau even as perplexity drops if the specific capability (multi-step math, code, factual recall) requires a capability that emerges more sharply rather than gradually. Conversely, a model trained on domain-specific data (code only) may see perplexity increase on the general held-out set while coding benchmarks improve, because its distribution has shifted.

Expect these follow-ups

• Why are some capabilities (like arithmetic) thought to emerge sharply at certain model scales rather than improving smoothly?
• How would you design a held-out set that better predicts downstream task performance?

As asked

Implement a numerically stable softmax function and explain the log-sum-exp trick. What goes wrong if you compute softmax naively on logits of magnitude 1000, and how does the stable version fix it?

Sample answer outline

Naive softmax computes exp(x_i) / sum(exp(x_j)), which overflows to inf for large positive x or underflows to zero for large negative x. The log-sum-exp trick subtracts the maximum value from all inputs before exponentiation: exp(x_i - max_x) / sum(exp(x_j - max_x)). Subtracting the max ensures the largest exponent is zero, preventing overflow while keeping all other values in a numerically safe range. The result is mathematically identical to naive softmax because the max cancels in numerator and denominator.

Expect these follow-ups

• How does this same principle appear in the online softmax computation in FlashAttention?
• Why is log-softmax implemented separately from log(softmax(x)) in PyTorch, and what is the numerical advantage?