Level 4 — ApplicationScaling & Efficient Architectures

Scaling & Efficient Architectures

⏱ 60 minutes60 marksprintable — key stays hidden on paper

Level 4 — Application (novel problems, no hints) Time limit: 60 minutes Total marks: 60

Question 1 — Compute-Optimal Budget Allocation (12 marks)

A research lab follows the Chinchilla compute-optimal recipe, where the optimal parameter count $N_{opt}$ and training tokens $D_{opt}$ each scale as a power of the compute budget $C$ (in FLOPs):

$N_{opt} \propto C^{0.5}, \qquad D_{opt} \propto C^{0.5}, \qquad C \approx 6\,N\,D$

The lab trained a compute-optimal model with $N = 7 \times 10^9$ parameters on $D = 1.4 \times 10^{11}$ tokens.

(a) Verify that this model is (approximately) compute-optimal by checking the Chinchilla tokens-per-parameter heuristic, and compute the training compute budget $C$ . (4)

(b) The lab now receives a 100× larger compute budget. Under the compute-optimal scaling, give the new $N_{opt}$ and $D_{opt}$ . (4)

(c) A deployment engineer argues: "We will serve this model to billions of users, so inference cost dominates. We should train a smaller model on more tokens than Chinchilla-optimal." Explain, with reference to the compute–parameter tradeoff, why this is rational, and state one quantity that becomes worse during training as a result. (4)

Question 2 — Designing an MoE Layer (14 marks)

You are designing a Mixture-of-Experts feed-forward layer to replace a dense FFN. The dense FFN has $d_{model}=4096$ and hidden size $d_{ff}=16384$ . Your MoE layer uses $E = 64$ experts (each an FFN of the same hidden size), with a top-2 router.

(a) Compute the ratio of total MoE parameters to activated parameters per token (ignore router and attention params). (3)

(b) The router produces logits $g = [2.0, 1.0, 1.0, 0.0, \dots]$ (rest very negative). After softmax over just these top gate values, the top-2 experts are selected. Compute the two gating weights that scale the selected experts' outputs (softmax over the top-2 logits only). (4)

(c) During training you observe that 3 of the 64 experts receive 70% of all tokens. Name this failure mode, explain why it is self-reinforcing, and describe an auxiliary load-balancing loss term (give its functional form in words or symbols). (4)

(d) Explain one reason MoE improves parameter-efficiency of scaling but does not proportionally reduce the memory required to host the model. (3)

Question 3 — KV-Cache & Long-Context Memory (12 marks)

A decoder-only Transformer has $L = 48$ layers, $h = 32$ attention heads, head dimension $d_k = 128$ , and stores KV-cache in FP16 (2 bytes).

(a) Derive a formula for KV-cache memory (bytes) as a function of sequence length $S$ and batch size $B$ , then compute it for $B=1$ , $S = 32{,}000$ . (5)

(b) The team switches to Grouped-Query Attention (GQA) with 4 KV groups (instead of 32). Compute the new KV-cache size and the reduction factor. (3)

(c) A colleague proposes Multi-Query Attention (MQA) instead. State the KV-cache size and one quality/robustness tradeoff versus GQA. (2)

(d) Beyond compressing the KV-cache, name one architectural alternative that avoids the linear-in- $S$ KV memory growth entirely, and state the mechanism that enables this. (2)

Question 4 — Speculative Decoding Throughput (12 marks)

You accelerate a large target model using speculative decoding with a small draft model. The draft proposes $\gamma = 4$ tokens per step; the target verifies them in a single forward pass. Empirically the per-token acceptance probability is $\alpha = 0.8$ (accepts are i.i.d.).

(a) The expected number of tokens accepted per verification step (including the "bonus" token when all are accepted) for a draft of length $\gamma$ with acceptance probability $\alpha$ is:

$\mathbb{E}[\text{tokens}] = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha}$

Compute this expected number for $\alpha = 0.8$ , $\gamma = 4$ . (4)

(b) The draft model costs $c = 0.2$ target-forward-passes per drafted token, and one verification = 1 target forward pass. Compute the speedup factor versus standard autoregressive decoding, defined as (expected accepted tokens per step) / (cost per step in target-forward-pass units). Assume drafting the 4 tokens costs $4c$ and verification costs $1$ . (5)

(c) Speculative decoding is exact — it produces the same output distribution as the target alone. Explain in one or two sentences the mechanism (rejection sampling / correction) that guarantees this, despite the draft model being different. (3)

Question 5 — Parallelism Strategy Design (10 marks)

You must train a 175B-parameter model (FP16, 2 bytes/param) that does not fit in a single 80 GB GPU.

(a) Estimate the memory just for parameters + Adam optimizer states + gradients in FP16/FP32 mixed precision (assume: params FP16 = 2B bytes, gradients FP16 = 2B, and Adam keeps FP32 master weights + 2 moments = 12B bytes per param, where B = number of params). Give the total for 175B params in GB. (4)

(b) Explain how ZeRO Stage 3 / FSDP reduces this per-GPU footprint across $P$ GPUs, and what communication cost this introduces. (3)

(c) Contrast tensor parallelism and pipeline parallelism in one sentence each, and state which one introduces "pipeline bubbles." (3)

Answer keyMark scheme & solutions

Question 1 (12 marks)

(a) Chinchilla heuristic ≈ 20 tokens per parameter. Check: $D/N = 1.4\times10^{11} / 7\times10^9 = 20$ ✓ (1 mark for check, 1 for stating heuristic). Compute: $C = 6ND = 6 \times (7\times10^9)(1.4\times10^{11}) = 6 \times 9.8\times10^{20} = 5.88\times10^{21}$ FLOPs. (2 marks)

(b) New budget $C' = 100C$ . Since $N_{opt}, D_{opt} \propto C^{0.5}$ , scale factor $= \sqrt{100} = 10$ .

$N_{opt}' = 10 \times 7\times10^9 = 7\times10^{10}$ (70B). (2)
$D_{opt}' = 10 \times 1.4\times10^{11} = 1.4\times10^{12}$ (1.4T tokens). (2) (Sanity: $6 N' D' = 6(7\times10^{10})(1.4\times10^{12}) = 5.88\times10^{23} = 100C$ ✓.)

(c) For inference-dominated deployment, total lifetime cost = training + (per-query cost × #queries). Per-query cost scales with $N$ (activated params). A smaller $N$ trained on more $D$ (past Chinchilla-optimal) yields a model with equal/near-equal quality but permanently cheaper inference — worth the extra training compute when queries are enormous. (3 marks) Quantity that gets worse: training compute efficiency / loss-per-FLOP during training — you spend more FLOPs than compute-optimal for the same loss (over-training). (1 mark)

Question 2 (14 marks)

(a) Each expert FFN params ≈ $2 \times d_{model} \times d_{ff} = 2 \times 4096 \times 16384$ (two matrices). Ratio of total to activated = experts total / experts used = $64 / 2 = \mathbf{32\times}$ . (3 marks) (per-token only 2 experts active → 32× more total params than activated.)

(b) Softmax over top-2 logits $\{2.0, 1.0\}$ : $e^{2.0}=7.389,\; e^{1.0}=2.718$ ; sum $=10.107$ . $w_1 = 7.389/10.107 = 0.731$ , $w_2 = 2.718/10.107 = 0.269$ . (4 marks) (2 for exp values, 2 for normalized weights.)

(c) Failure mode: routing collapse / expert imbalance (a few experts dominate). (1) Self-reinforcing because the router is trained jointly: experts receiving more tokens get more gradient updates → become better → router routes even more tokens to them (rich-get-richer). (2) Auxiliary load-balance loss (Switch/GShard style): $\mathcal{L}_{aux} = \alpha \cdot E \sum_{i=1}^{E} f_i \, P_i$ where $f_i$ = fraction of tokens dispatched to expert $i$ , $P_i$ = mean router probability for expert $i$ ; minimized when load is uniform ( $f_i = P_i = 1/E$ ). (1)

(d) MoE reduces FLOPs/activated params per token but all experts' weights must still reside in memory (any token might route to any expert), so total memory footprint scales with total params, not activated params. (3 marks)

Question 3 (12 marks)

(a) KV-cache stores K and V (factor 2), across all layers, all heads: $\text{Bytes} = 2 \times L \times h \times d_k \times S \times B \times (\text{bytes/elem})$ $= 2 \times 48 \times 32 \times 128 \times S \times B \times 2$ . Per token per batch = $2\times48\times32\times128\times2 = 1{,}572{,}864$ bytes ≈ 1.5 MB. For $S=32000, B=1$ : $1{,}572{,}864 \times 32000 = 5.033\times10^{10}$ bytes ≈ 46.9 GiB (or ~50.3 GB decimal). (5 marks: 2 formula, 3 compute.)

(b) GQA with 4 KV groups: K/V heads = 4 instead of 32 → factor $32/4 = 8$ reduction. New size $= 46.9 / 8 ≈ 5.86$ GiB. Reduction factor 8×. (3 marks)

(c) MQA = 1 KV head → cache = $46.9/32 ≈ 1.47$ GiB (32× reduction). Tradeoff: MQA can degrade quality/training stability more than GQA because all query heads share one K/V; GQA is the middle ground. (2 marks)

(d) State-space models (Mamba/S4) — maintain a fixed-size recurrent state instead of a growing KV-cache, so memory is $O(1)$ in sequence length (linear-time, constant-state recurrence). (2 marks)

Question 4 (12 marks)

(a) $\mathbb{E}[\text{tokens}] = \dfrac{1-\alpha^{\gamma+1}}{1-\alpha} = \dfrac{1 - 0.8^{5}}{1-0.8} = \dfrac{1 - 0.32768}{0.2} = \dfrac{0.67232}{0.2} = \mathbf{3.3616}$ tokens/step. (4 marks)

(b) Cost per step in target-forward-pass units = drafting cost + verification = $4c + 1 = 4(0.2)+1 = 1.8$ . Baseline autoregressive: 1 token per 1 target forward pass. Speedup $= \dfrac{\text{tokens/step}}{\text{cost/step}} = \dfrac{3.3616}{1.8} = \mathbf{1.868\times}$ . (5 marks)

(c) For each drafted token, accept it with probability $\min(1, p_{target}/p_{draft})$ ; on rejection, resample from the corrected residual distribution $\propto \max(0, p_{target}-p_{draft})$ . This modified rejection-sampling scheme provably yields samples from exactly the target distribution regardless of the draft. (3 marks)

Question 5 (10 marks)

(a) Per-param bytes: FP16 params 2 + FP16 grads 2 + Adam (FP32 master 4 + 2 moments 8 = 12) = 16 bytes/param. For $175\times10^9$ params: $16 \times 175\times10^9 = 2.8\times10^{12}$ bytes = 2.8 TB (≈ 2.55 TiB). (4 marks) (This is why it needs many GPUs.)

(b) ZeRO-3 / FSDP shards params, gradients, and optimizer states across $P$ GPUs, so per-GPU footprint ≈ (total)/ $P$ . Cost: parameters must be all-gathered just-in-time for each layer's forward/backward (and re-sharded after), adding all-gather + reduce-scatter communication each step. (3 marks)

(c) Tensor parallelism: splits individual weight matrices/ops within a layer across GPUs (each GPU computes a slice, needs all-reduce). Pipeline parallelism: assigns different layers to different GPUs, passing activations stage-to-stage. Pipeline parallelism introduces pipeline bubbles (idle time). (3 marks)

[
  {"claim":"Q1a compute C = 5.88e21 FLOPs and tokens/param=20","code":"N=7e9; D=1.4e11; ratio=D/N; C=6*N*D; result = (ratio==20) and abs(C-5.88e21)/5.88e21 < 1e-9"},
  {"claim":"Q1b scaling factor sqrt(100)=10 gives N=7e10, D=1.4e12","code":"f=sqrt(Integer(100)); result = (f==10) and (10*7e9==7e10) and (10*1.4e11==1.4e12)"},
  {"claim":"Q2b top-2 softmax weights are approx 0.731 and 0.269","code":"import math; a=math.exp(2.0); b=math.exp(1.0); s=a+b; w1=a/s; w2=b/s; result = abs(w1-0.731)<0.002 and abs(w2-0.269)<0.002"},
  {"claim":"Q3a KV cache for S=32000,B=1 approx 5.033e10 bytes","code":"bytes_per=2*48*32*128*2; total=bytes_per*32000; result = abs(total-5.033e10)/5.033e10 < 0.01"},
  {"claim":"Q4a expected tokens = 3.3616","code":"a=Rational(8,10); g=4; val=(1-a**(g+1))/(1-a); result = abs(float(val)-3.3616) < 1e-4"},
  {"claim":"Q4b speedup = 3.3616/1.8 approx 1.868","code":"tokens=(1-0.8**5)/(1-0.8); cost=4*0.2+1; sp=tokens/cost; result = abs(sp-1.868) < 0.005"},
  {"claim":"Q5a memory 175B params at 16 bytes = 2.8e12 bytes","code":"mem=16*175e9; result = mem==2.8e12"}
]