ML / LLM System Design — Cheat Sheet

Answer: Standard self-attention has time complexity $O(n^2 d)$ ($n$ = sequence length, $d$ = dimension) because it must compute an $n \times n$ attention matrix. Methods to reduce it include:

• FlashAttention: Does not change the mathematical result; reduces wall-clock time by tiling and recomputation to minimize HBM accesses
• Sparse Attention: Longformer/BigBird uses local windows + global tokens, reducing complexity to $O(n \cdot w)$
• Linear Attention: Approximates softmax with a kernel function, reducing complexity to $O(n d^2)$, but usually with some accuracy loss

Follow-up: Why is FlashAttention not considered an "approximate" attention? What low-level optimizations does it perform?

FlashAttention loads Q, K, V in tiles into SRAM, computes the softmax online normalization in SRAM (by maintaining a running max and running sum), then writes the result back to HBM. It is mathematically equivalent to standard attention — it merely reduces the number of HBM read/write operations.

Answer:

• BatchNorm: Computes mean and variance across the batch dimension for each feature. Requires maintaining running mean/var during training and uses fixed statistics at inference. Sensitive to batch size; not suitable for variable-length sequences.
• LayerNorm: Computes mean and variance across the feature dimension for each sample (each token is normalized independently); does not depend on batch statistics. The standard choice in Transformers.

$$\text{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$

Follow-up: What advantage does RMSNorm have over LayerNorm?

RMSNorm removes the mean-centering step and only performs variance normalization: $\text{RMSNorm}(x) = \gamma \odot x / \sqrt{\text{mean}(x^2) + \epsilon}$. Slightly less computation, similar practical performance; adopted by the LLaMA series.

Answer: Gradient clipping constrains the norm of the gradient to a threshold: $$\text{if } \|g\| > c: \quad g \leftarrow g \cdot \frac{c}{\|g\|}$$ In LLM training, a small number of anomalous samples can produce extremely large gradients (gradient spikes), causing abrupt parameter changes or NaN values. Gradient clipping (typically $c = 1.0$) is the standard safeguard against training collapse.

Follow-up: How can you tell whether the gradient clipping threshold is set appropriately?

Monitor the frequency of clipping events in the training log. Occasional triggering (< 5% of steps) is normal; frequent triggering suggests the learning rate may be too large; never triggering during otherwise stable training suggests the threshold may be too loose.

Answer:

• Warmup: Linearly increase the learning rate at the beginning of training (typically for the first 1%–3% of steps), because with random initialization the gradient direction is unstable and a large LR may cause divergence.
• Cosine decay: After warmup, the LR follows a cosine curve from peak to near zero: $\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})(1 + \cos(\pi t / T))$

Follow-up: How does a WSD (Warmup-Stable-Decay) schedule differ from a cosine schedule?

WSD maintains a constant LR after warmup (the stable phase), then decays rapidly at the end. Its advantage is that mid-training checkpoints have good quality, making it suitable for scenarios that need to evaluate downstream tasks from intermediate checkpoints.

Answer: The core of FlashAttention is an IO-aware algorithm design:

1. Split Q, K, V into small blocks, each small enough to fit in GPU SRAM (on-chip memory)
2. Compute softmax and matrix multiplication within SRAM
3. Use online softmax (maintaining row-wise running max and sum statistics) to avoid needing global information for the softmax
4. Avoid writing the $n \times n$ attention matrix to HBM (GPU memory), thereby reducing HBM reads and writes

Source of speedup: standard attention must write/read the attention matrix from HBM, making HBM bandwidth the bottleneck. FlashAttention concentrates computation in SRAM, reducing HBM access from $O(n^2)$ to $O(n^2 d^2 / M)$ ($M$ = SRAM size).

Follow-up: How large is the benefit of FlashAttention for training vs inference respectively?

In training, the main saving is in HBM accesses during backpropagation (the attention matrix is not stored in the forward pass; it is recomputed as needed during backward). In inference, the benefit is primarily in the prefill stage (long prompts); the benefit in the decode stage (single token per step) is smaller.

Answer:

• LoRA / QLoRA: Insert a low-rank bypass ($BA$) alongside a weight matrix; only the bypass parameters are trained. QLoRA further quantizes the base weights to 4-bit.
• Prefix Tuning: Prepend trainable "virtual token" vectors to the keys and values of each attention layer.
• Adapter: Insert a small MLP bottleneck (down-projection → nonlinearity → up-projection) between Transformer sublayers; only adapter parameters are trained.
• Prompt Tuning: Prepend a small number of trainable soft prompt vectors to the input embeddings (only at the input layer).

Follow-up: What is the trade-off between parameter efficiency and expressiveness for these methods?

Fewer parameters save more memory but cap the expressiveness. LoRA, acting directly on weight matrices, typically outperforms adapters and prefix tuning at similar parameter counts. In extreme scenarios (e.g., only tens of examples), fewer parameters can actually prevent overfitting.

Answer:

• Static batching: Collects a batch of requests and waits until all requests finish generation before releasing the batch. If one request is much shorter than others, GPU resources are wasted once it completes (padding and idle waiting).
• Continuous batching (iteration-level scheduling): After each generation step (one token), checks whether any request has finished; completed requests are immediately replaced by new ones. GPU utilization is significantly improved.

Follow-up: Are PagedAttention and continuous batching used together?

Yes. Continuous batching answers "when to schedule requests"; PagedAttention answers "how to allocate KV cache memory" — it splits the KV cache into fixed-size pages, allocates on demand, and avoids memory fragmentation caused by variable request lengths.

Answer:

• ZeRO-1: Shards optimizer states (AdamW's FP32 master copy + $m$ + $v$, 12 bytes/param) across GPUs. Each GPU holds only $1/P$ of the optimizer state; after the update, parameters are gathered via AllGather.
• ZeRO-2: On top of ZeRO-1, also shards gradients. Each GPU keeps only $1/P$ of the gradients (the rest are discarded after Reduce-Scatter).
• ZeRO-3: Also shards parameters. During forward and backward passes, parameters are All-Gathered on demand and released after use.

Communication: ZeRO-1/2 match standard DP ($2|\theta|$ per step). ZeRO-3 needs an AllGather of params in both forward and backward, plus a Reduce-Scatter of grads in backward — $3|\theta|$ total (~$1.5\times$ plain DP's $2|\theta|$).

Follow-up: When is ZeRO-2 better than ZeRO-3?

When model parameters fit on a single GPU but optimizer states do not, ZeRO-2 has lower communication overhead. A typical scenario is fine-tuning a medium-scale model (e.g., 7B–13B) with gradient checkpointing + ZeRO-2.

Answer: Each RLHF-PPO step:

1. Sample a batch of prompts; generate responses with the current policy $\pi_\theta$
2. Score each (prompt, response) pair with the reward model
3. Compute the KL penalty using the reference policy $\pi_{\text{ref}}$
4. Compute advantages (typically with GAE)
5. Update the policy with the PPO clip objective (multiple mini-batch update rounds)

Why KL penalization is needed: Without a KL constraint, the policy quickly drifts into out-of-distribution (OOD) blind spots of the reward model — generating responses that score high under the RM but that humans actually dislike (reward hacking). The KL penalty keeps the policy close to $\pi_{\text{ref}}$ (i.e., the SFT model).

Follow-up: Can you give a concrete example of reward hacking?

For example, if the reward model favors long answers (because good answers in the training data tended to be longer), the policy may learn to produce very long, repetitive responses to any question to obtain a high score — even though a human evaluator would find them verbose and unhelpful.

Answer: Common approaches:

• Replay / mixed training: Mix some general instruction data or pretraining data into the SFT data
• LoRA / PEFT: Only a small number of parameters are updated; pretrained knowledge is preserved in the frozen base weights
• Regularization: Methods like EWC (Elastic Weight Consolidation) penalize large deviations of important parameters
• Low learning rate: When doing full fine-tuning, use an LR 1–2 orders of magnitude lower than pre-training

Follow-up: How do you quantify the degree of catastrophic forgetting?

Evaluate on both a general benchmark (e.g., MMLU, HellaSwag) and the target task benchmark before and after fine-tuning. If performance on the general benchmark drops by more than a few percentage points, significant forgetting has occurred.

Answer:

• GPTQ (Optimal Brain Quantization series): Quantizes layer by layer, using second-order information (inverse Hessian) to minimize the reconstruction error of the layer output. Quantizes column by column; after quantizing each column, compensates the remaining columns.
• AWQ (Activation-Aware Weight Quantization): The core observation is that a small number of "salient channels" (channels with large activations) are critical to output quality. AWQ protects the weights of those channels (e.g., using per-channel scaling to effectively increase precision), rather than quantizing all weights uniformly.

Follow-up: When quantizing to INT4, why is SmoothQuant on activations so important?

Activations often contain outliers (abnormally large values), which stretch the quantization range and reduce effective precision. SmoothQuant migrates activation outliers to weights via a mathematically equivalent per-channel scaling, making the activation distribution more uniform so that both weights and activations can be quantized to lower bit-widths.

Answer: In Megatron-LM's design:

• TP splits linear layers (weight matrices of attention and MLP)
• SP (Sequence Parallelism) splits activations of non-linear operations (LayerNorm, Dropout) — along the sequence dimension
• Junction: a TP layer ends with AllReduce (or ReduceScatter); an SP layer also requires communication. Megatron-LM fuses these two communication operations, so no additional communication is introduced.

Benefit: after the TP AllReduce, each GPU holds the full sequence in its activations (redundant). SP eliminates this redundancy; each GPU holds only $1/P$ of the sequence activations, significantly reducing activation memory.

Follow-up: Does SP help with gradient checkpointing?

Yes. SP reduces the activation volume stored on each GPU. Without gradient checkpointing, activation memory drops from $O(L \cdot n \cdot d)$ to $O(L \cdot n/P \cdot d)$. Even with gradient checkpointing, the temporary memory used during recomputation is reduced proportionally.

Answer:

• Training: Uses the Bradley-Terry preference model. Given prompt $x$ and a response pair $(y_w, y_l)$ ($y_w$ labeled as better), the reward model loss is $-\log \sigma(r(x, y_w) - r(x, y_l))$. The model is typically initialized from the SFT model, with the language model head replaced by a scalar output head.
• Evaluation metrics:
  • Preference prediction accuracy: Accuracy at predicting which of a held-out preference pair is better
  • Reward distribution separability: Whether the reward distributions of chosen and rejected responses are sufficiently separated
  • Reward hack robustness: Whether the reward still ranks OOD responses (generated by the policy) reasonably

Follow-up: Why does the reward model need to be updated periodically?

Because the policy changes continuously during RL training, the distribution of generated responses gradually shifts away from the RM's training distribution (the SFT model's output distribution). On out-of-distribution data, the RM may produce inaccurate scores, leading to reward hacking.

Answer:

• Problem: Traditional KV cache pre-allocates a contiguous memory block for each request (at maximum sequence length). But actual generation lengths vary, causing significant memory waste (internal fragmentation) and preventing sharing across requests (external fragmentation).
• PagedAttention mechanism: Borrows the virtual memory paging idea from operating systems:
  1. KV cache is divided into fixed-size blocks (e.g., each block stores KV for 16 tokens)
  2. A block table records the mapping from logical blocks to physical blocks for each request
  3. New blocks are dynamically allocated when new tokens are generated; blocks are freed when a request finishes
  4. Supports copy-on-write: for beam search candidates sharing the same prefix, KV blocks can be shared

Follow-up: Does PagedAttention negatively affect latency?

The indirect addressing via the block table introduces a small overhead (compared to direct access with contiguous memory), but in practice this overhead is very small (typically < 5%), since attention computation itself is compute-bound or memory-bound and the addressing overhead is not the bottleneck.

Answer:

• Task abstraction: Each task defines a dataset, prompt template (few-shot format), metric, and output type (generation / loglikelihood)
• Evaluation modes:
  • Likelihood-based (e.g., MMLU): compute log-prob for each option, select the maximum
  • Generation-based (e.g., GSM8K): generate output then evaluate with rules / code execution
  • LLM-as-judge (e.g., MT-Bench): score with a stronger model
• Reproducibility: Fix seeds; record prompt templates and few-shot examples; temperature=0 (or fixed)
• Efficiency: Likelihood tasks suit large batches; generation tasks should be sorted by length to reduce padding
• Contamination detection: Check for n-gram overlap between training data and test sets

Follow-up: Why distinguish between "knowledge" and "reasoning" evaluation?

Because a model may perform well on knowledge-heavy tasks (e.g., factual questions in MMLU) but poorly on reasoning-heavy tasks (e.g., math, code), or vice versa. Evaluating them separately helps pinpoint where the model's capabilities are weak.

Answer: Factors to consider:

• Task complexity: Simple classification/extraction tasks usually need only r=4–16; complex reasoning/generation tasks may require r=32–64
• Data size: Use a small rank with little data to prevent overfitting; increase rank when data is plentiful to raise capacity
• Target modules: Applying LoRA only to q_proj and v_proj (fewest parameters) vs. all linear layers (q/k/v/o + MLP gate/up/down) trades parameter count against effectiveness — all linear layers typically yield better results
• Common practice: Start from r=16, α=2r; compare r=8/16/32/64 on the validation set

Follow-up: Can LoRA be combined with QLoRA? Is the accuracy loss large when using 4-bit quantized base weights + LoRA?

Yes — QLoRA is exactly this approach. In practice, 4-bit NF4 quantized base weights + LoRA fine-tuning matches FP16 full fine-tuning within an acceptable margin on most tasks (typically within 1–3 percentage points) while saving enormous memory.

Answer:

Naive co-located PPO runs rollout and training serially on the same set of GPUs:

• Rollout phase: the actor performs autoregressive inference (memory-bound; throughput limited by HBM bandwidth); the trainer waits idle.
• Train phase: PPO backpropagation is compute-intensive; the rollout worker waits idle.

The two phases alternate, and overall GPU utilization is the weighted average of each phase's utilization — well below peak training utilization.

How the disaggregated architecture solves it:

1. Independent rollout workers (vLLM/SGLang engines) continuously generate responses, producing a rollout buffer.
2. Independent train workers (ZeRO-3/FSDP) pull data from the buffer and continuously execute PPO/GRPO updates.
3. The two sets of workers run concurrently; weights are synchronized at a fixed frequency (typically each iteration).

This way, rollout and training are each optimized for their own workload (inference engine vs. training framework) without blocking each other.

Follow-up 1: How much weight-sync bandwidth is needed between rollout workers and train workers in a disaggregated architecture?

For a 7B BF16 model, one complete weight sync transfers ~14 GB. If syncing once per minute, that is ~14 GB ÷ 60 s ≈ 0.23 GB/s, well below the bandwidth ceiling of NVLink/RDMA (sync overhead is negligible). With LoRA-in-RL, only LoRA parameters need syncing (~100 MB scale), greatly reducing sync overhead.

Follow-up 2: What effect does staleness from async rollout have on PPO? How can it be mitigated?

Staleness causes rollout to generate data with old parameters $\pi_{\theta_{\text{old}}}$, introducing an off-policy bias. PPO's importance ratio clip ($\epsilon \approx 0.1\text{–}0.2$) tolerates mild staleness, but when staleness is large, gradient estimate variance grows and training becomes unstable. Mitigation: control the weight sync frequency (no more than a few mini-batch updates), or use more aggressive importance sampling correction.

Answer:

Consider two consecutive linear transforms $Y = GELU(XA)B$ (an MLP block), $A \in \mathbb{R}^{h \times 4h}$, $B \in \mathbb{R}^{4h \times h}$:

1. Column-Parallel $A$: Split $A$ column-wise as $[A_1, A_2]$; each GPU computes $GELU(X A_i)$ independently, no communication needed. GeLU is element-wise and naturally separable.
2. Row-Parallel $B$: Split $B$ row-wise as $\begin{bmatrix} B_1 \\ B_2 \end{bmatrix}$; each GPU computes $Y_i = GELU(XA_i) B_i$.
3. Final AllReduce: $Y = Y_1 + Y_2$ (one AllReduce).

Key insight: Column-Parallel output is exactly the input of Row-Parallel; the intermediate nonlinearity (GeLU) is element-wise and requires no communication. Therefore the entire MLP block needs only one AllReduce (forward), and one in the backward pass as well.

Without this design, each layer would require a separate AllReduce, doubling the communication volume.

Follow-up: Can the same trick be applied to the QKV projection and output projection in the attention block?

Yes. QKV projection uses Column-Parallel (outputs are distributed to each head, which naturally splits column-wise); the output projection uses Row-Parallel, followed by AllReduce. The entire attention block also needs only one AllReduce.

Answer:

Let target model distribution be $p(x)$ and draft model distribution be $q(x)$.

Accept-reject sampling:

1. Sample token $x$ from $q(x)$
2. If $p(x) \geq q(x)$: accept (probability 1)
3. If $p(x) < q(x)$: accept with probability $p(x)/q(x)$

Total probability of accepting token $x$:

• Sampled from $q$ and accepted: $q(x) \cdot \min(1, p(x)/q(x)) = \min(p(x), q(x))$
• Sampled from $q$, rejected, then resampled to $x$: more complex but derivable

Final effective probability:

$$P(\text{output}=x) = \min(p(x), q(x)) + \underbrace{\Big(1 - \sum_i \min(p(i), q(i))\Big)}_{P(\text{reject})} \cdot \frac{\max(0, p(x) - q(x))}{1 - \sum_i \min(p(i), q(i))}$$

The second term = reject probability $\times$ normalized residual; the reject probability cancels the denominator to give $\max(0, p(x)-q(x))$, so $P(\text{output}=x) = \min(p,q) + \max(0,\,p-q) = p(x)$ — exactly the target distribution.

Core intuition: when $p(x) > q(x)$, the draft model "under-sampled" $x$ and the deficit is compensated from the residual probability mass after rejection; when $p(x) < q(x)$, rejection removes the excess probability.

Follow-up: Where is the efficiency bottleneck of speculative decoding?

The bottleneck is the draft model's acceptance rate. If the distributions of the draft model and target model diverge significantly, the acceptance rate is low, most draft tokens are rejected, and the speedup is poor. Improvements include: Medusa-style multi-head prediction, or selecting a draft model whose distribution is closer to the target model.

Answer:

Step 1: The Bradley-Terry model assumes the optimal policy $\pi^*$ satisfies:

$$p(y_w \succ y_l | x) = \sigma(r^*(x, y_w) - r^*(x, y_l))$$

Step 2: Under a KL constraint, the closed-form solution for the optimal policy is:

$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y|x) \exp\!\left(\frac{r(x,y)}{\beta}\right)$$

where $Z(x)$ is the partition function.

Step 3: Solve for the reward:

$$r(x, y) = \beta \log \frac{\pi^*(y|x)}{\pi_{\text{ref}}(y|x)} + \beta \log Z(x)$$

Step 4: Substitute $r$ into the Bradley-Terry model; $Z(x)$ cancels in the difference:

$$p(y_w \succ y_l | x) = \sigma\!\left(\beta \log \frac{\pi^*(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - \beta \log \frac{\pi^*(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\right)$$

Step 5: Replace $\pi^*$ with the trainable $\pi_\theta$ and take the negative log-likelihood to obtain the DPO loss.

Follow-up: DPO's derivation assumes preference data comes from the optimal policy; what practical problems does this assumption cause?

In practice, preference data usually comes from the SFT model (not the optimal policy), which means the reward implicitly learned by DPO may be inaccurate. This is also why online DPO (iterative DPO, where each round generates data with the latest policy) typically outperforms offline DPO.

Answer:

• Problem: When mainstream models score near the ceiling on a benchmark (e.g., >90% on MMLU), discriminability drops. Possible causes include:
  • Training data contamination (test set data leaked into the training set)
  • Insufficient task difficulty (primarily knowledge retrieval, not deep reasoning)
  • Format optimization (models tuned to the benchmark's prompt format)
• Approaches:
  • Use harder benchmarks (e.g., MMLU-Pro, GPQA, MATH)
  • Use dynamically generated evaluation questions
  • Rely on human evaluation (e.g., Chatbot Arena Elo rankings)
  • Detect and report data contamination

Follow-up: What are the design philosophies of HELM and lm-evaluation-harness?

HELM (Stanford) emphasizes "comprehensiveness" — covering multiple dimensions (accuracy, calibration, robustness, fairness, efficiency) with detailed documentation and standardized evaluation procedures for each scenario, but extending to new tasks is relatively heavy. lm-evaluation-harness (EleutherAI) emphasizes "flexibility and community contribution" — tasks are defined concisely (config-driven); the community can quickly add new tasks; 400+ tasks provide broad coverage, though standardization is relatively lower.

Answer:

Motivation: Prefill (processing the prompt) and decode (generating tokens one at a time) have completely different computational characteristics:

Disaggregated Serving Design:

• Prefill nodes: high-compute configuration; process prompts in large batches → generate KV cache
• Decode nodes: high-bandwidth configuration; receive KV cache → generate tokens one at a time
• KV cache is transferred between nodes over a high-speed network (RDMA/NCCL)

Benefits: The two stages can be scaled independently, preventing the memory-bound nature of the decode stage from dragging down the compute utilization of the prefill stage.

Follow-up: How large is the bandwidth requirement for KV cache transfer?

For a 70B model with sequence length 4K and FP16 KV cache, the KV cache per request is on the order of a few hundred MB. If decode nodes need to ingest KV caches from tens of requests per second, tens of GB/s of network bandwidth is required — feasible on modern data-center RDMA networks.

Answer:

• Principle: During the forward pass, intermediate activations are not saved; only some "checkpoints" are kept (typically one per layer boundary). During backpropagation, activations are recomputed from the nearest checkpoint.
• Memory: Reduced from $O(L \cdot a)$ ($a$ = activation size per layer) to $O(\sqrt{L} \cdot a)$ or $O(L')$ ($L'$ = number of checkpoints)
• Compute: Approximately 33% extra forward computation (each checkpoint segment is recomputed forward once)

Practical choice:

• Do not use if memory is sufficient (saves time)
• Enable when memory is insufficient but a 33% training slowdown is acceptable
• Can be selectively enabled (e.g., only for certain large layers)

Follow-up: For selective gradient checkpointing, how do you choose which layers to checkpoint?

Typically choose layers with the largest activations (e.g., the attention matrix is an $O(n^2)$ memory consumer). Layers with small activations (e.g., LayerNorm, embedding) are not checkpointed, achieving a better balance between memory savings and computation overhead.

Answer:

PPO's clipped surrogate objective:

$$L^{CLIP} = \mathbb{E}\left[\min\left(r_t(\theta) \hat{A}_t, \; \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t\right)\right]$$

where $r_t(\theta) = \pi_\theta(a_t|s_t) / \pi_{\theta_{\text{old}}}(a_t|s_t)$; $\epsilon$ is typically 0.1–0.2.

Purpose: When $r_t$ deviates too far from 1, the clip limits the change in the objective function, preventing excessively large single-step updates.

Special considerations in RLHF:

• In standard RL (games, etc.), the state-action space is large and $r_t$ does not deviate much
• In RLHF, the language model's generation space is exponential and the policy may change rapidly
• Therefore $\epsilon$ may need to be reduced, or the number of PPO update epochs increased to better exploit each batch of samples

Follow-up: How is the value function loss balanced against the policy loss in PPO?

Typically a weighted sum: $L = L^{CLIP} + c_1 L^{VF} - c_2 H(\pi)$, where $L^{VF}$ is the MSE loss of the value function and $H(\pi)$ is an entropy bonus to prevent premature collapse. In RLHF, tuning $c_1$ and $c_2$ is critical for training stability.

Answer:

• N-gram overlap detection: Compute the intersection of n-grams (e.g., 8-gram, 13-gram) from the test set against the training data. If the overlap rate exceeds a threshold, flag as potentially contaminated.
• Membership inference: Check whether the model's perplexity on test set samples is anomalously low compared to held-out data; low perplexity may indicate the sample appeared in training.
• Canonical order test: Shuffle the answer choices; if accuracy drops significantly, the model may have memorized the answer at a specific position (suggesting contamination rather than genuine understanding).
• Canary test: Insert unique "canary" sentences into the test set; after training, check whether the model can reproduce them perfectly.

Follow-up: Why might n-gram overlap detection produce false positives?

Because some common knowledge (e.g., "the sun rises in the east") will appear in both training and test sets; n-gram overlap does not mean genuine "memorization." One needs to distinguish "factual public knowledge" from "verbatim copying of specific test samples."

Answer:

Step 1: Clarify

• 70B actor (~140 GB BF16 parameters) + critic (similar or slightly smaller) + ref model + RM
• Naive co-located memory for all 4 models: parameters + optimizer states on the order of 1 TB (not feasible; separation required)
• Goal: run on 8–64 × 80 GB A100/H100 GPUs with throughput that meets a reasonable training schedule

Step 2: 4-model memory decomposition (order-of-magnitude estimates)

• With full parameter training, the complete training state for actor + critic (parameters + gradients + optimizer) is ~1.5–2 TB scale; ZeRO-3 sharding across train workers requires tens of 80 GB GPUs (exact count depends on whether FP32 master copy, activation, and framework overhead are included).
• With LoRA-in-RL (rank=16–32), trainable parameters drop to $\lesssim 1\%$ of total, optimizer states fall from ~560 GB to a few GB — greatly reducing train worker memory requirements.

Step 3: Topology design

Rollout cluster (inference-optimized)    Train cluster (training-optimized)
┌──────────────────────────┐         ┌─────────────────────────┐
│ vLLM / SGLang            │         │ ZeRO-3 / FSDP           │
│  - actor (FP16 weights)  │◄─weights│  - actor (trainable)    │
│  - ref model (frozen)    │  sync   │  - critic (trainable)   │
│  - RM (frozen)           │         │                         │
│                          │──data──►│  rollout buffer         │
│  continuous rollout,     │         │  PPO / GRPO updates     │
│  output                  │         │                         │
│  (prompt, resp, reward,  │         │                         │
│   log_prob, value)       │         │                         │
└──────────────────────────┘         └─────────────────────────┘

• Rollout and train run concurrently (async) or alternately (sync); weights synced once per iteration.
• Ref model and RM require only inference; placing them on the rollout side saves train-side memory.

Step 4: LoRA-in-RL vs full parameter updates

⚠️ The memory figures above are order-of-magnitude estimates (derived from parameter count × bytes/param formulas). Actual values differ substantially due to activations, KV cache, and framework overhead. In interviews, explicitly state "estimate."

Follow-up: How do you decide the resource ratio of rollout to train workers in a disaggregated architecture?

It depends on the ratio of rollout throughput to train throughput. If rollout is the bottleneck (long responses, large batches), add more rollout workers. If train is the bottleneck (large critic, many PPO mini-batches), add more train workers. In practice, first profile the GPU-hours per iteration for each side, allocate proportionally, then adjust based on observed queue utilization.

Answer (high-level overview):

1. Clarify:

• 10M DAU → estimated QPS of ~100–1000 (assuming 1–3 conversation turns per user per day)
• Latency SLA: P95 < 2s (time to first token), P99 < 5s
• Domain adaptation needed (customer service phrasing, product knowledge)

2. Data:

• Historical customer service conversation logs → clean and anonymize → build SFT data
• Periodically sample online bad cases (low ratings, escalated to human agents) → human annotation → feed back into training
• RAG: build a vector knowledge base from product documentation and FAQs

3. Model:

• Base model: 7B–13B scale (balance quality and inference cost)
• SFT (LoRA) fine-tuned on customer service data
• RAG retrieval augmentation: user query → retrieve relevant documents → append to prompt context

4. Serving:

• Quantization: INT8 or INT4 (GPTQ/AWQ) → reduce per-GPU inference cost
• vLLM / TensorRT-LLM deployment, continuous batching + PagedAttention
• Multiple replicas + load balancing, auto-scaling with traffic

5. Monitoring:

• Online metrics: escalation rate, user satisfaction score, average conversation turns
• Quality drift: regularly run eval on a standard test set and monitor score changes
• Safety: apply sensitive word and harmful content filtering to outputs

Follow-up: In this system, what problems do RAG and fine-tuning each solve? Can they replace each other?

Fine-tuning handles "style and format" — making the model respond in a customer-service tone and follow the correct workflow. RAG handles "knowledge and facts" — providing up-to-date product information and company policy. They are complementary, not interchangeable: fine-tuning alone causes hallucinations about product details; RAG alone makes the model sound like a generic assistant rather than a professional customer service agent. The ideal solution combines both.

Standard attention must write the complete $N \times N$ attention matrix to HBM (High Bandwidth Memory), making IO the bottleneck. FlashAttention uses GPU SRAM (fast but small) via tiling:

1. Split $Q, K, V$ into blocks of size $B_r \times d$ and $B_c \times d$; load only one block into SRAM at a time
2. For each Q block, iterate over all K/V blocks and compute local attention within SRAM
3. Use online softmax by maintaining a running max $m$ and running sum $\ell$: after processing the $j$-th KV block, update the previously accumulated output $O_j$ with a correction factor $e^{m_{j-1} - m_j}$, avoiding the need for global normalization

$$O_j = \text{diag}(\ell_j)^{-1}\Big(e^{m_{j-1}-m_j}\,\ell_{j-1}\,O_{j-1} + \tilde{P}_j V_j\Big)$$

IO complexity drops from $O(N^2 d)$ HBM accesses to $O(N^2 d^2 / M)$ ($M$ = SRAM size); memory drops from $O(N^2)$ to $O(N)$ (the full attention matrix is never materialized).

Follow-up: FlashAttention's backward pass requires recomputing the attention matrix (recomputation) — what are the similarities and differences with gradient checkpointing? In very-long-sequence settings, what additional parallelization optimizations does FlashAttention v2 introduce?

Simple position interpolation (PI) uniformly scales position $m$ to $m \cdot L_{\text{train}} / L_{\text{target}}$. The problem is that RoPE frequencies $\theta_j = 10000^{-2j/d}$ span multiple orders of magnitude:

• Low dimensions (small $j$) → high frequency, encoding fine-grained positional relationships at short distances
• High dimensions (large $j$) → low frequency, encoding coarse positional relationships at long distances

After uniform scaling, the rotation angle in high-frequency dimensions changes too densely; the model cannot distinguish adjacent tokens (high-frequency information is "compressed together"), analogous to applying a low-pass filter to an image and losing edge details.

NTK-aware interpolation rescales the base frequency from $b$ to $b' = b \cdot \alpha^{d/(d-2)}$ ($\alpha = L_{\text{target}}/L_{\text{train}}$):

• Low-dimensional high-frequency components are nearly unchanged → preserving local resolution
• High-dimensional low-frequency components are stretched → encoding longer distances

This is analogous to the NTK theory's observation about the difference in learning difficulty between high-frequency and low-frequency features: high-frequency features require higher resolution; low-frequency features can be safely extrapolated.

Follow-up: YaRN further applies temperature scaling to the attention score on top of NTK-aware. What is the motivation? Why is modifying position encoding alone insufficient to fully recover long-context task performance?

Expert collapse in MoE: a small number of experts are selected frequently while the rest are nearly idle, wasting the model's effective capacity.

Auxiliary load balancing loss:

$$\mathcal{L}_{\text{aux}} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i$$

• $N$ = number of experts, $f_i$ = fraction of tokens routed to expert $i$ (discrete statistics), $P_i$ = average probability the router assigns to expert $i$ (continuous, differentiable)
• The $f_i \cdot P_i$ term encourages both to be uniformly distributed: when an expert is both frequently selected and has high router confidence, the penalty is largest
• $\alpha$ is set to a small value to prevent it from dominating the main training loss

Capacity factor (CF): Limits the maximum number of tokens each expert can process in one batch = $\text{CF} \times T/N$. CF too small → tokens are dropped (overflow) → information loss; CF too large → computational waste (padding). CF needs to be dynamically adjusted based on the degree of load imbalance.

Follow-up: DeepSeek-MoE proposes fine-grained expert segmentation (splitting large experts into multiple smaller ones) and a shared expert mechanism. How does this design fundamentally mitigate the tension between load balancing (requiring uniformity) and model capability (requiring specialization)?

ZeRO-3 requires All-Gather of the complete parameters for each layer before the forward pass can proceed. Naive implementation: All-Gather → wait → compute → free; communication and computation are serial, and the GPU spends a long time waiting.

Overlap strategy (dependency graph analysis using backward as an example):

Forward: compute(L) ← All-Gather(L)          compute(L+1) ← All-Gather(L+1)
          ↓ can overlap: while compute(L) runs, asynchronously prefetch All-Gather(L+1)

• Forward: While computing layer $l$, asynchronously launch the All-Gather for layer $l+1$ parameters (prefetch). Requirement: compute time for layer $l$ ≥ communication time for layer $l+1$.
• Backward: Similarly, while computing layer $l$ gradients, prefetch layer $l-1$ parameters; Reduce-Scatter of layer $l$ gradients can also be overlapped with the next layer's computation.

Cost: Simultaneously holding more parameter copies increases (current layer + prefetch layer), adding memory pressure. Total communication is ~$3 \times |\theta|$ per step (higher than DP's $2 \times |\theta|$); when cross-node bandwidth is limited this can become a bottleneck.

Follow-up: At what model scale and hardware conditions does ZeRO-3's communication overhead become unacceptable, making TP (intra-node NVLink) + ZeRO-2 the better choice? Analyze from the perspective of the ratio of communication volume to computation volume.

The $\log \frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)}$ term in the DPO loss is essentially an importance-weighted reward estimate.

Source of off-policy bias:

• As the divergence between $\pi_\theta$ and $\pi_{\text{ref}}$ grows, the variance of importance weights increases and gradient estimates become unstable
• The training data covers a fixed $y$-space anchored to $\pi_{\text{ref}}$'s support. $\pi_\theta$ may have learned to generate responses not seen in training data, but those responses cannot be evaluated by the DPO loss → optimization signal has blind spots
• Analogous to distribution shift in off-policy RL: the further the policy departs from the data-collection policy, the less reliable the estimates

How iterative DPO mitigates this:

1. Sample new responses with the current $\pi_\theta$
2. Annotate preferences with a reward model or human annotators
3. Use the new $\pi_\theta$ as the new $\pi_{\text{ref}}$; retrain DPO
4. Repeat → training data gradually becomes on-policy

Online DPO goes further: within the training loop, it samples $\pi_\theta$'s outputs in real time, scores them with the RM, and immediately updates.

Follow-up: In online DPO, if the reward model itself has a systematic bias (e.g., preferring verbose answers), how would online iteration amplify that problem? What are the mechanistic similarities and differences with reward hacking in PPO?

This is a manifestation of Goodhart's Law: when a proxy metric is optimized to the extreme, it decouples from the true objective.

Theoretical intuition:

• Let the true reward be $r^*(x,y)$, proxy RM $r_\phi(x,y)$, and their difference $\delta(x,y) = r_\phi - r^*$
• When $\pi_\theta$ optimizes in the direction of $\nabla_\theta \mathbb{E}[r_\phi]$, it not only improves $r^*$ but also exploits $\delta$ — entering regions where $r_\phi$ is overestimated
• As $\text{KL}(\pi_\theta \| \pi_{\text{ref}})$ increases, the policy departs further from the training distribution, and the generalization error of $r_\phi$ (i.e., $|\delta|$) grows monotonically
• Qualitative observation: the proxy reward keeps rising; true quality first rises then falls; the crossing point of the two curves is the "over-optimization inflection point"

Factors influencing the divergence rate:

• Larger RM capacity and more diverse preference data → the inflection point appears later
• Larger policy exploration space (longer, more diverse generation) → easier to find reward-hacking paths

Mitigation strategies: KL penalty, RM ensemble (taking the min or variance penalty across multiple RMs), periodic RM updates.

Follow-up: In practice, how does a reward model ensemble exploit agreement and disagreement among multiple RMs? What are the pros and cons of taking the min, the mean, or using disagreement as an uncertainty signal? How does computational cost affect feasibility?

MLA no longer stores the complete $K, V$; instead it stores a low-dimensional latent vector $c_t^{KV}$, which is decompressed at inference time:

$$c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{d_c}, \quad d_c \ll n_h \cdot d_h$$

The KV cache stores only $c_t^{KV}$ (dimension $d_c$); at attention computation time it projects back:

$$k_t = W^{UK} c_t^{KV}, \quad v_t = W^{UV} c_t^{KV}$$

KV cache size drops from $2 \times L \times n_h \times d_h \times s$ to $L \times d_c \times s$ ($d_c$ can be much smaller than $2 n_h d_h$).

Fundamental difference from GQA:

MLA's advantage: the number of Q heads is no longer directly tied to cache size, enabling large cache compression while retaining many Q heads. Trade-off: inference requires extra projection computation, and the low-rank constraint may limit pattern diversity across heads.

Follow-up: Does MLA's low-rank compression cause different attention heads' patterns to converge (loss of head diversity)? Can the high rank of the projection matrix $W^{UK}$ fully mitigate this risk? In practice, what signals can detect degradation of head diversity?