Format. An MLE onsite round where you build and train a model, not just discuss one. One report gives 60 minutes with a coding agent allowed; another describes it as an hour of coding followed by a 30-minute review. Both reports place it in the MLE loop next to a system design round and the faulty-comparator algorithm round.
The task. Implement an autoregressive language model and train it on the Shakespeare dataset. It has to run end to end, from the data through training. One report calls it nanoGPT and labels the round ML system design: an hour of coding, then a 30-minute review. (Generating text is the natural way to show the model works; the reports don't name it.)
The review. After the build, the interviewer asks Transformer fundamentals. The reported topics:
What the reports say about the company. One candidate says Mercor weighs communication heavily: follow the interviewer's lead and don't go deep into areas they don't know. The MLE opening has two roles, Marketplace and Frontier Data; a reply says HR described Marketplace as rebranded into Frontier Data. One candidate heard the result by email the same evening.
Ask before you build (none of these are answered in the reports):
A 60-minute build (our walkthrough, not from the reports). A character-level GPT in plain PyTorch. Writing the attention by hand, rather than calling nn.MultiheadAttention, gives you something concrete to point at when the review asks about self-attention, the causal mask and pre-norm. We ran this file: changing tokens after position 40 leaves every output before position 40 unchanged, a top-p near zero reproduces greedy decoding, and the sample at the end reads like broken Shakespeare. The production equivalent of the hand-written attention is F.scaled_dot_product_attention(q, k, v, is_causal=True).
import torch
import torch.nn as nn
import torch.nn.functional as F
# ---- data: character-level -------------------------------------------------
text = open("input.txt").read() # e.g. tiny Shakespeare
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]
block_size, batch_size = 128, 32
device = "cuda" if torch.cuda.is_available() else "cpu"
def get_batch(split):
d = train_data if split == "train" else val_data
ix = torch.randint(len(d) - block_size - 1, (batch_size,))
x = torch.stack([d[i:i + block_size] for i in ix])
y = torch.stack([d[i + 1:i + block_size + 1] for i in ix]) # next-token targets
return x.to(device), y.to(device)
# ---- model -------------------------------------------------------------------
class CausalSelfAttention(nn.Module):
def __init__(self, d_model, n_head, dropout):
super().__init__()
assert d_model % n_head == 0
self.n_head = n_head
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
self.drop = nn.Dropout(dropout)
def forward(self, x):
B, T, C = x.shape
q, k, v = self.qkv(x).split(C, dim=2)
# (B, T, C) -> (B, n_head, T, head_dim)
q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v))
att = (q @ k.transpose(-2, -1)) / (k.size(-1) ** 0.5) # scaled dot product
mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
att = att.masked_fill(~mask, float("-inf")) # no attending to the future
att = self.drop(F.softmax(att, dim=-1))
y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
return self.proj(y)
class Block(nn.Module):
"""Pre-norm: x + f(LayerNorm(x)), so the residual path stays an identity."""
def __init__(self, d_model, n_head, dropout):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = CausalSelfAttention(d_model, n_head, dropout)
self.ln2 = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, 4 * d_model), nn.GELU(),
nn.Linear(4 * d_model, d_model), nn.Dropout(dropout),
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
return x + self.mlp(self.ln2(x))
class GPT(nn.Module):
def __init__(self, vocab_size, d_model=128, n_head=4, n_layer=4, dropout=0.1):
super().__init__()
self.tok_emb = nn.Embedding(vocab_size, d_model)
self.pos_emb = nn.Embedding(block_size, d_model) # learned absolute positions
self.blocks = nn.Sequential(*[Block(d_model, n_head, dropout) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(d_model) # pre-norm needs a final norm
self.head = nn.Linear(d_model, vocab_size, bias=False)
def forward(self, idx, targets=None):
T = idx.size(1)
pos = torch.arange(T, device=idx.device)
x = self.tok_emb(idx) + self.pos_emb(pos)
logits = self.head(self.ln_f(self.blocks(x)))
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
self.eval() # dropout off while sampling
for _ in range(max_new_tokens):
logits, _ = self(idx[:, -block_size:]) # crop to the context window
logits = logits[:, -1, :] / temperature # temperature > 0; greedy = argmax
if top_k is not None:
kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, [-1]]
logits = logits.masked_fill(logits < kth, float("-inf"))
if top_p is not None: # nucleus: smallest set with mass >= p
sorted_logits, order = torch.sort(logits, descending=True)
cum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
drop = cum - F.softmax(sorted_logits, dim=-1) >= top_p # always keeps the top token
sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
logits = torch.full_like(logits, float("-inf")).scatter(1, order, sorted_logits)
next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
idx = torch.cat([idx, next_id], dim=1)
return idx
# ---- train -------------------------------------------------------------------
model = GPT(len(chars)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
@torch.no_grad()
def estimate_loss(iters=20):
model.eval()
out = {}
for split in ("train", "val"):
out[split] = torch.stack([model(*get_batch(split))[1] for _ in range(iters)]).mean().item()
model.train()
return out
for step in range(2000):
if step % 200 == 0:
print(step, estimate_loss())
x, y = get_batch("train")
_, loss = model(x, y)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
print("final", estimate_loss())
start = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(start, 300, temperature=0.8, top_p=0.9)[0].tolist()))
Short answers to the review topics (ours, not from the reports).
x = LN(x + f(x)). Pre-norm normalizes the sublayer's input: x = x + f(LN(x)), which leaves a clean identity path through the residual stream, so deep stacks train more stably and are less sensitive to learning-rate warmup. Pre-norm needs a final LayerNorm before the output layer; post-norm can do slightly better when it trains well, but is harder to train.softmax(QKᵀ / √d_k) V. With roughly unit-variance entries, the variance of q·k grows linearly with d_k, so dividing by √d_k keeps the logits at unit scale and the softmax from saturating. The causal mask sets future positions to −∞ before the softmax. Multiple heads attend to different subspaces in parallel. Cost is O(T²) in sequence length.Interviewer hints: The review after the build covers the Transformer architecture, pre-norm versus post-norm, self-attention, position embeddings, and decoding (greedy, beam search, top-k, top-p).