ML Engineer interview prep

Short answer: Data parallelism — when the model fits on one GPU: a replica on every device, the batch is split, gradients are combined with all-reduce. Tensor parallelism splits individual matmuls across GPUs — needs NVLink-class links, intra-node. Pipeline parallelism splits layers into stages; micro-batches fight the bubble. The largest models combine all three — 3D parallelism.

In depth:

1. Order of choice — start with data parallelism; a layer doesn't fit → tensor parallelism within a node; the whole model doesn't fit → pipeline across nodes.
2. The bubble — stages sit idle at the start and end of each step; mitigated by slicing the batch into micro-batches.

⚠️ Common mistake: conflating tensor and pipeline parallelism: TP splits matrices inside a layer (all GPUs compute the same layer together), PP splits the model by layers into sequential stages.

Short answer: The KV-cache stores the per-layer K and V matrices of every token processed so far. Without it, autoregressive generation would recompute K and V for the whole prefix at every step — O(n²) work per sequence; with the cache, each decode step computes K, V and Q only for the single new token — O(n) total.

In depth:

1. Why it works — K and V of past tokens do not depend on future ones: compute them once, reuse them at every later step.
2. What is cached — only K and V, separately per layer. The current token's Q is computed fresh: it is needed exactly once — for this token to attend over the past — and is never reused.
3. The cost — memory: the cache grows linearly with context length and batch size, and at long contexts it rivals the weights themselves in size.

No cache:   step t → recompute K,V for tokens 1..t  (O(t) per step, O(n²) total)
With cache: step t → K,V for token t only + read
            cached K,V of tokens 1..t-1 from HBM    (O(1) compute, O(n) total)

⚠️ Common mistake: saying "Q, K and V are cached." Q is not cached — a token's query is used exactly once, at the moment it is generated.

Short answer: Separate the deterministic from the stochastic. Transforms, losses and metrics get ordinary unit tests on hand-computable cases; training is covered by smoke tests (overfit a tiny batch, shape and gradient-flow checks) and behavioral tests with tolerances.

In depth:

1. Test the code, not the model — deterministic pieces (features, losses) deserve exact asserts; stochastic ones get tolerances and invariants.
2. Smoke test in CI — a short 2–3 step run: the pipeline assembles, the loss is finite and decreasing.

⚠️ Common mistake: "we look at test accuracy" — that's model evaluation, not code testing. A dataloader bug can cost a couple of metric points and never surface.

Short answer: Subtract the row max before exp — exp(89) already overflows float32. For the loss, don't compute softmax and log separately; combine them as log-sum-exp: CE = logsumexp(z) − z[y].

In depth:

1. Shift by the max — softmax is shift-invariant: exp(z − m)/Σexp(z − m) is the same value, minus the overflow.
2. log-softmax in one piece — log(softmax(z)) = z − m − log Σ exp(z − m); we never take the log of a near-zero.
3. CE as indexing — minus log-softmax at the correct class position, averaged over the batch.

import numpy as np

def log_softmax(z):
    z = z - z.max(axis=1, keepdims=True)  # exp(89) overflows fp32
    return z - np.log(np.exp(z).sum(axis=1, keepdims=True))

def cross_entropy(z, y):
    # CE = logsumexp(z) - z[y]; no separate softmax -> log
    n = len(y)
    return -log_softmax(z)[np.arange(n), y].mean()

⚠️ Common mistake: computing softmax, then taking the log — under saturation the probability rounds to 0 and log yields −inf. This is exactly why PyTorch's F.cross_entropy takes logits, not probabilities.

Short answer: Overfitting means the model learned the noise of the training set instead of the pattern: low error on train, high error on new data. You detect it via the train/validation gap and fix it with regularization, a simpler model, early stopping, and more data.

In depth:

1. Detection — compare train vs validation error; learning curves: train error keeps falling while validation error stalls or rises → the model started memorizing noise.
2. Treatment — reduce capacity or add signal, don't just chase the metric.

⚠️ Common mistake: saying "cross-validation prevents overfitting." CV only detects it — it gives an honest quality estimate. Prevention comes from regularization, simpler models, and data.