AI research engineer system design interview questions

As asked

Design a shared GPU cluster for research training jobs, batch fine-tunes, and low-latency inference. How do you schedule work fairly and efficiently?

Sample answer outline

Separate workload classes because inference and long training jobs have different latency and preemption tolerance. Use quotas, priority classes, gang scheduling for distributed training, and preemptible queues for exploratory work. Strong answers discuss GPU fragmentation, topology awareness, checkpointing, fair-share policies, and visibility into queue time by team. Inference should usually live in reserved pools with autoscaling and admission control rather than competing directly with training. The common failure is maximising utilisation while making urgent production workloads wait behind low-priority experiments.

Expect these follow-ups

• How do you handle a training job that needs eight GPUs on the same node group?
• What metrics tell you the scheduler policy is unfair?
• When would you split research and production into separate clusters?

As asked

Walk me through the infrastructure design for pretraining a 70B dense transformer on 1 trillion tokens using 512 A100 GPUs. Address parallelism strategy, fault tolerance, checkpointing, and data loading.

Sample answer outline

Use 3D parallelism: tensor parallelism within a node (8 GPUs per tensor parallel group since intra-node NVLink bandwidth is high), pipeline parallelism across nodes for very large depth, and data parallelism across pipeline replicas using FSDP or ZeRO-3 for optimizer state sharding. Implement checkpointing every 500 to 1000 steps with distributed save to object storage (S3 or equivalent) to tolerate GPU failures; Elastic training can restart from the last checkpoint with a reduced number of GPUs. Use a high-throughput data loader with prefetched shards, multiple workers per GPU, and tokenized data stored in memory-mapped files (MMapIndexedDataset) to saturate GPU compute.

Expect these follow-ups

• How do you handle a single GPU failure at step 50,000 in a pipeline-parallel setup?
• What is the communication bottleneck between nodes and how does your network topology (InfiniBand vs Ethernet) affect parallelism choices?

As asked

Design a retrieval-augmented generation system that can search 100 million documents and return relevant passages to an LLM within 200ms. Cover embedding, indexing, retrieval, reranking, and context assembly.

Sample answer outline

Embed all documents offline with a bi-encoder (like BGE or E5) and store dense vectors in a vector database (Faiss or Qdrant) alongside sparse BM25 inverted index for hybrid retrieval. At query time, embed the query and run approximate nearest neighbor search (HNSW) to retrieve top-100 candidates in under 50ms, then run a cross-encoder reranker on the top-100 to produce final top-5 with precise relevance scores within the remaining budget. Assemble context by formatting passages with metadata and passing to the LLM. The reranker is the main latency risk; cache common queries and quantize the reranker to INT8.

Expect these follow-ups

• How do you handle documents that update frequently and require index freshness within minutes?
• What chunking strategy do you use for long documents and how does chunk size affect retrieval quality?

As asked

Your team produces a new model checkpoint every two days. Design an automated evaluation pipeline that produces comparable, trustworthy benchmark results across checkpoints and flags regressions before they reach production.

Sample answer outline

Build a benchmark runner that pins dataset versions, prompt templates, decoding parameters (temperature=0, max_tokens fixed), and model config to a hash stored alongside every result. Run a suite of benchmarks (MMLU, GSM8K, HumanEval, BBH) deterministically by fixing random seeds and storing per-example predictions. Compare new results against a baseline checkpoint using a statistical test (bootstrap confidence intervals) to flag regressions. Store all results in a database with checkpoint metadata for historical comparison. Use a CI-style trigger that blocks deployment if any benchmark drops more than a threshold.

Expect these follow-ups

• How do you handle a benchmark where the model's output format changes and old scoring scripts break?
• What is the minimum set of benchmarks you would run if compute is limited?

As asked

Design an inference serving system for a 13B LLM that needs to handle 10,000 requests per second with a P99 latency under 2 seconds for 200-token responses. Cover batching, scaling, and hardware choices.

Sample answer outline

Use continuous batching (as in vLLM) rather than static batching: new requests join the batch as soon as a previous request finishes, maximizing GPU utilization. Shard the model across multiple GPUs with tensor parallelism and deploy multiple model replicas behind a load balancer. INT8 quantization reduces memory, allowing more requests per GPU. Use PagedAttention for KV cache management to avoid memory fragmentation and support variable sequence lengths. Auto-scale the number of replicas based on queue depth. Prefill and decode can be split onto separate hardware since prefill is compute-bound and decode is memory-bandwidth-bound.

Expect these follow-ups

• What is the difference between prefill latency and decode latency, and why do they have different hardware requirements?
• How would you handle requests with very different output lengths without letting short requests wait for long ones?

As asked

Design a system where human raters provide preference feedback on model outputs in near-real-time, and a reward model is continuously updated to drive online RL fine-tuning. Address data freshness, reward model staleness, and stability.

Sample answer outline

Separate the system into a serving cluster generating responses, a rating interface collecting (winner, loser) pairs from raters, a reward model training job that periodically retrains on accumulated data, and a policy training job running PPO against the current reward model. Use a versioned reward model registry so the RL job always uses a fixed reward model for a training epoch, preventing instability from mid-training reward shifts. Implement drift detection to flag when the policy has moved far enough from the reference that a fresh reward model is needed. Buffer at least several thousand new preference pairs before updating the reward model to avoid overfitting to the current policy's outputs.

Expect these follow-ups

• What is reward model overoptimization and how do you detect it in a live system?
• How would you handle adversarial raters who consistently label wrong answers as correct?

As asked

You have a general-purpose 7B LLM and 50B tokens of high-quality medical text. Design a continual pretraining pipeline that adapts the model to the medical domain without catastrophically forgetting general capabilities.

Sample answer outline

Mix the medical corpus with a fraction of general domain data (typically 5 to 20%) to prevent catastrophic forgetting, tuning the mix ratio based on a forgetting metric on general benchmarks. Use a lower peak learning rate than the original pretraining (1/3 to 1/10) with a short warmup to avoid destroying pre-learned representations. Evaluate on both domain benchmarks (MedQA, PubMedQA) and general benchmarks (MMLU, HellaSwag) at regular checkpoints. If forgetting is severe, add elastic weight consolidation (EWC) or similar regularization targeting the most important parameters. Select a checkpoint that balances both sets of metrics rather than optimizing solely for domain performance.

Expect these follow-ups

• How do you determine what mixing ratio prevents forgetting without wasting compute on irrelevant data?
• What is the difference between continual pretraining and domain-adaptive fine-tuning in terms of learning rate and data expectations?

As asked

Your research team wants to run 200 small-scale ablation experiments (1B model, 10B tokens each) concurrently on a shared cluster of 1,000 GPUs. Design the job scheduling, resource allocation, experiment tracking, and result aggregation system.

Sample answer outline

Use a cluster scheduler (SLURM or Kubernetes with GPU operators) with a priority queue so urgent experiments can preempt lower-priority ones. Each experiment is a containerized job with a config file defining all hyperparameters, committed to a git hash for reproducibility. Results are streamed to a central metrics store (W&B or MLflow) with structured metadata enabling SQL-style queries across experiments. Build a result aggregation dashboard that plots metric distributions across seeds and highlights which ablations most influence the target metric. Implement automatic retries for jobs that fail due to preemption versus those that fail due to bugs.

Expect these follow-ups

• How do you prevent one team member's large experiment from starving everyone else's small ablations?
• What are the tradeoffs between running fewer seeds per experiment versus more experiments with one seed each?

As asked

Design the data pipeline, architecture choices, and training curriculum for pretraining a vision-language model that accepts interleaved images and text. Assume you have 400M image-caption pairs and 1T text tokens.

Sample answer outline

Use a pre-trained vision encoder (ViT) frozen initially and a pre-trained LLM backbone. Connect them via a lightweight projection layer (MLP or cross-attention) trained to align visual and text tokens. Training in three stages: first align the projection using image-caption pairs with frozen encoder and LLM; second unfreeze the encoder and train with image-text interleaved data; third fine-tune on instruction-following data for downstream tasks. Use separate data pipelines for image and text that merge at the dataloader level, ensuring balanced sampling. Evaluate on VQA, captioning, and OCR benchmarks at each stage.

Expect these follow-ups

• How do you handle images of very different resolutions and aspect ratios in the visual encoder?
• What is the tradeoff between using cross-attention versus token concatenation to fuse visual and text representations?

As asked

You want to improve a model's mathematical reasoning using only synthetic data generated by a stronger teacher model. Design the full pipeline from problem generation to filtering to training, including how you prevent the student from overfitting to the teacher's style.

Sample answer outline

Use the teacher to generate diverse problems by seeding with varied difficulty parameters, then generate multiple chain-of-thought solutions per problem and keep only those whose final answer can be verified (for math, verify numerically). Filter by difficulty using the student's pass rate: problems the student always gets right or always fails are uninformative; keep problems in the 20 to 80% student accuracy range. Diverse decoding (temperature sampling from the teacher) prevents style overfitting. Evaluate on held-out human-written benchmarks (MATH, AIME) rather than teacher-generated tests to avoid distribution collapse.

Expect these follow-ups

• What is the risk of the student imitating the teacher's reasoning errors rather than the correct approach?
• How do you decide when synthetic data is no longer improving the model and more real data is needed?