Mistral AI interview questions

Project once, split into heads, attend in each head, merge.

def mhsa(x, w_qkv, w_o, h):            # x: (B, T, d)
    B, T, d = x.shape
    q, k, v = (x @ w_qkv).split(d, -1)
    heads = lambda t: t.view(B, T, h, -1).transpose(1, 2)
    q, k, v = heads(q), heads(k), heads(v)  # (B, h, T, d/h)
    s = q @ k.transpose(-2, -1) / (d // h) ** 0.5
    future = torch.ones(T, T, dtype=torch.bool,
                        device=x.device).triu(1)
    s = s.masked_fill(future, float("-inf"))
    y = s.softmax(-1) @ v
    return y.transpose(1, 2).reshape(B, T, d) @ w_o

Interviewers watch for the mask before softmax, scaling by the head size, not d, and the transpose before the final reshape.

Pre-norm should read x + attn(norm(x)), then x + mlp(norm(x)): normalise each sublayer's input and add its output to the un-normalised stream.

Common bugs: adding to the normed tensor, so the residual path stops being an identity; normalising the sublayer's output instead of its input; one norm shared where two are meant; and no final norm before the output layer. A quick test: zero both output projections, and the block must return x exactly.

LayerNorm subtracts the mean, divides by the standard deviation, then applies a learned scale and shift. RMSNorm skips the mean: it divides by the root mean square and keeps only the scale. That's less computation with comparable quality in the paper's experiments, so models like Llama use it, in the same pre-norm positions.

Greedy takes the argmax. Both samplers sort the probabilities, keep a prefix, renormalise it and sample.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

def pick(logits, rng, mode="greedy", k=5, p=0.9,
         temp=1.0):
    if mode == "greedy":
        return int(np.argmax(logits))
    probs = softmax(logits / temp)
    order = np.argsort(probs)[::-1]   # high to low
    if mode == "top_k":
        keep = order[:k]
    else:                             # top_p
        cum = np.cumsum(probs[order])
        keep = order[:np.searchsorted(cum, p) + 1]
    q = probs[keep] / probs[keep].sum()
    return int(rng.choice(keep, p=q))

Interviewers check the max subtracted in softmax, that top-p keeps the token that crosses p, and the renormalisation before sampling.

It stores every processed token's keys and values, per layer and head, so each decode step computes only the new token's.

It's exact because of the causal mask: earlier tokens never depend on later ones, and greedy output was identical with and without the cache. After a 960-token prompt on our test machine, a step took 9.6 ms with the cache against 228.8 ms without, about 24 times faster.

In multi-head attention every head has its own queries, keys and values. Grouped-query attention lets a group of query heads share one key-value head.

The cache holds keys and values per key-value head, so fewer of them means a proportionally smaller KV cache: Llama 2 70B's 64 query heads share 8, an 8× cut, with little quality loss.

First prove it: a profiler shows each step waiting on the next batch, and GPU utilisation dips in a sawtooth. A quick test is to repeat one cached batch; if throughput jumps, the input pipeline is the bottleneck.

Then move work out of the hot path. Tokenise, filter and dedupe offline, writing packed, pre-tokenised sequences into large shards instead of many small files. At train time, stream shards with parallel workers, prefetching and pinned memory, shuffle by shard order plus a buffer, and make the loader resumable from a saved position.

Sort failures by stage before changing anything: log the retrieved passages for every bad answer. Right passage missing: fix retrieval with better chunking, hybrid search, query rewriting, metadata filters, or more candidates plus a reranker. Passage present, answer wrong: fix generation by putting it first, trimming distractors, and asking the model to quote and to abstain. Source wrong or stale: fix the data. Then add each failure to the eval set, so the fix stays fixed.

Ingest: parse documents to text, keep metadata like source, date and permissions, and split by structure into chunks of a few hundred tokens. Embed them in batches and store vectors with the metadata.

At query time, embed the question with the same model, retrieve the top k with filters, and give the chat model chunks with ids to cite. Make retrieval a tool, so the agent can search again with a rewritten query.

Both learn from preference pairs, a chosen and a rejected answer. RLHF trains a reward model on them, then optimises the policy against it with PPO, held near a reference model: powerful, but several models and a touchy RL loop.

DPO drops the reward model and the RL: one loss on the pairs raises the chosen answer's probability relative to the reference. Simpler and steadier.

An MoE layer swaps one feed-forward block for many experts and a small router that sends each token to the top few, often 2.

Parameters grow, but each token uses a slice: Mixtral 8×7B has about 47B parameters and uses about 13B per token. Expert parallelism places experts on different GPUs and ships tokens to them in an all-to-all exchange, so the costs become communication and load balance.

Tensor parallelism splits each weight matrix across GPUs, so every layer needs a collective step: fast only over very fast links, usually inside one machine. Pipeline parallelism gives each GPU a run of consecutive layers and passes activations along: little communication, but GPUs sit idle in the pipeline bubble unless the batch is split into micro-batches. Large runs combine both with data parallelism.

ZeRO removes data parallelism's redundancy: every GPU normally holds full parameters, gradients and Adam states. ZeRO-1 shards the optimiser states, ZeRO-2 also shards gradients, using reduce-scatter instead of all-reduce, and ZeRO-3 also shards parameters, all-gathering each layer's weights just before its forward and backward, then freeing them. FSDP is PyTorch's ZeRO-3-style implementation, sharded per wrapped module.

The price is communication, so overlap it: prefetch the next layer's all-gather while this one computes, and reduce-scatter each layer's gradients as soon as its backward finishes.

Adam keeps running averages of each gradient and its square, and divides the first by the root of the second, so every weight gets its own step size. Both start at zero, so bias correction divides by one minus beta to the power t.

def adamw(w, g, m, v, t, lr=1e-3, b1=0.9,
          b2=0.999, eps=1e-8, wd=0.01):
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    m_hat = m / (1 - b1 ** t)
    v_hat = v / (1 - b2 ** t)
    w = w - lr * (m_hat / (v_hat ** 0.5 + eps) + wd * w)
    return w, m, v

An L2 term in the loss gets divided by that root too, so high-gradient weights decay less. AdamW decouples it: shrink the weights directly, outside the adaptive step.

Split it into stages with different costs: render pages and detect layout, then recognise each region with the right model, such as text OCR, tables or formulas, then assemble Markdown in reading order. OCR only pages without a reliable text layer.

Make the API asynchronous: upload returns a job id, pages fan out on a queue, and a failed page retries alone. Batch GPU stages across jobs, keep a fast lane for small requests, and cache by file hash. Evaluate each stage, with edit distance to reference Markdown.

The forward process is fixed: it adds Gaussian noise over T steps until data becomes noise, with a closed form for any step. The reverse process is learned: a network removes noise one step at a time.

# one DDPM training step, noise-prediction loss
t = torch.randint(0, T, (B,))
eps = torch.randn_like(x0)
a = alpha_bar[t].view(B, 1, 1, 1)
xt = a.sqrt() * x0 + (1 - a).sqrt() * eps
loss = F.mse_loss(model(xt, t), eps)

Training samples one random step per example. Generation starts from pure noise and runs the steps sequentially, so it's slow; DDIM samplers skip steps.