90 PyTorch problems from real ML/AI interviews at Google, Meta, Anthropic, and more.
I struggled to grind for ML/AI interviews so I went back to the basics and created a list after careful research. These are real problems from first person reports from real engineer interviews.
Important
Don't use GPT. The whole point is to struggle through these yourself. If you paste these into ChatGPT you're wasting your time. The goal is to deeply understand PyTorch, not to get an answer. I used GPT to help write some of the initial code, but I tested and solved every problem myself. That's where the learning happens.
Three question sets, 90 problems total:
| Set | Focus | Questions | Difficulty |
|---|---|---|---|
| v1 | Core PyTorch | 35 | Basic to Hard |
| v2 | LLMs from Scratch | 25 | Easy to Hard |
| v3 | Advanced ML Systems (NEW!) | 30 | Easy to Expert |
v3 questions are tagged with the companies that actually ask them. If you're interviewing at Anthropic, Google, Meta, or any top AI company, start there.
# Install PyTorch
# https://pytorch.org/get-started/locally/
# Pick a problem, fill in the TODOs, compare with the solution
jupyter notebook torch/basic/lin-regression/lin-regression.ipynbEach problem has a question file and a _SOLN solution file. Fill in the ... and #TODO blocks, then check your work.
| # | Problem | Links |
|---|---|---|
| 1 | Implement linear regression | Question / Solution |
| 2 | Write a custom Dataset and Dataloader | Question / Solution |
| 3 | Write a custom activation function | Question / Solution |
| 4 | Implement Custom Loss Function (Huber Loss) | Question / Solution |
| 5 | Implement a Deep Neural Network | Question / Solution |
| 6 | Visualize Training Progress with TensorBoard | Question / Solution |
| 7 | Save and Load Your PyTorch Model | Question / Solution |
| 10 | Implement Softmax function from scratch |
| # | Problem | Links |
|---|---|---|
| 1 | Implement a CNN on CIFAR-10 | Question / Solution |
| 2 | Implement an RNN from Scratch | Question / Solution |
| 3 | Use torchvision.transforms for data augmentation | Question / Solution |
| 4 | Add a benchmark to your PyTorch code | Question / Solution |
| 5 | Train an autoencoder for anomaly detection | Question / Solution |
| 6 | Quantize your language model | Question / Solution |
| 7 | Implement Mixed Precision Training | Question / Solution |
| # | Problem | Links |
|---|---|---|
| 1 | Implement parameter initialization for a CNN | Question / Solution |
| 2 | Implement a CNN from Scratch | Question / Solution |
| 3 | Implement an LSTM from Scratch | Question / Solution |
| 4 | Implement AlexNet from scratch | |
| 5 | Build a Dense Retrieval System | |
| 6 | Implement KNN from scratch in PyTorch | |
| 7 | Train a 3D CNN for segmenting CT images | Question / Solution |
| # | Problem | Links |
|---|---|---|
| 1 | Write a custom Autograd function (SILU) | Question / Solution |
| 2 | Write a Neural Style Transfer | |
| 3 | Build a Graph Neural Network (GNN) from scratch | |
| 4 | Build a Graph Convolutional Network (GCN) | |
| 5 | Write a Transformer | Question / Solution |
| 6 | Write a GAN | Question / Solution |
| 7 | Write Sequence-to-Sequence with Attention | Question / Solution |
| 8 | Enable distributed training (DDP) | |
| 9 | Work with Sparse Tensors | |
| 10 | Add GradCam/SHAP to explain the model | Question / Solution |
| 11 | Linear Probe on CLIP Features | |
| 12 | Cross Modal Embedding Visualization (t-SNE/UMAP) | |
| 13 | Implement a Vision Transformer | |
| 14 | Implement a Variational Autoencoder |
Build a Large Language Model from scratch, one question at a time.
| # | Problem | Links |
|---|---|---|
| 1 | Implement KL Divergence Loss | |
| 2 | Implement RMS Norm | |
| 3 | Implement Byte Pair Encoding from Scratch | Question / Solution |
| 4 | Create a RAG Search of Embeddings | |
| 5 | Implement Speculative Decoding | |
| 6 | Implement Attention from Scratch | Question / Solution |
| 7 | Implement Multi-Head Attention | Question / Solution |
| 8 | Implement Grouped Query Attention | Question / Solution |
| 9 | Implement KV Cache in Multi-Head Attention | |
| 10 | Implement Sinusoidal Embeddings | Question / Solution |
| 11 | Implement ROPE Embeddings | Question / Solution |
| 12 | Implement SmolLM from Scratch | Question / Solution |
| 13 | Implement Quantization (GPTQ) | |
| 14 | Implement Beam Search for decoding | |
| 15 | Implement Top K Sampling | |
| 16 | Implement Top p Sampling | |
| 17 | Implement Temperature Sampling | |
| 18 | Implement LoRA / QLoRA | |
| 19 | Mix two models into a Mixture of Experts | |
| 20 | Apply SFT on SmolLM | |
| 21 | Apply RLHF on SmolLM | |
| 22 | Implement DPO based RLHF | |
| 23 | Add continuous batching to your LLM | |
| 24 | Chunk Data for Dense Passage Retrieval | |
| 25 | Implement 5D Parallelism for Large Scale Training |
30 questions covering what top AI companies actually ask in 2024-2025 interviews. Tagged with real companies.
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 1 | Implement Softmax (numerically stable) | Easy | Apple, Meta, Google, Amazon | Q / S |
| 2 | Implement K-Means Clustering | Easy | Uber, LinkedIn, Google, Amazon | Q / S |
| 3 | Implement KNN in PyTorch | Easy | Uber, LinkedIn, Meta | Q / S |
| 4 | Implement Logistic Regression | Easy | Google, Meta, Amazon | Q / S |
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 5 | Contrastive Loss (InfoNCE) + CLIP | Medium | OpenAI, Anthropic, DeepMind, Midjourney | Q / S |
| 6 | 2D Positional Embeddings | Medium | Anthropic, DeepMind, Midjourney, Runway | Q / S |
| 7 | Top-p (Nucleus) Sampling | Medium | Anthropic, OpenAI, DeepMind, Perplexity | Q / S |
| 8 | Top-k Sampling | Medium | Anthropic, OpenAI, DeepMind, Cohere | Q / S |
| 9 | Beam Search for LLM Decoding | Medium | Google, DeepMind, Meta, Apple | Q / S |
| 10 | Temperature Sampling | Easy | OpenAI, Anthropic, Cohere, Perplexity | Q / S |
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 11 | Implement LoRA on a Linear Layer | Medium | Meta, Google, Anthropic, OpenAI, Databricks | Q / S |
| 14 | Implement DPO Loss from Scratch | Hard | Anthropic, OpenAI, DeepMind, Meta | Q / S |
| 15 | Implement PPO for RLHF | Hard | Anthropic, OpenAI, DeepMind, Meta | Q / S |
| 16 | Implement Gradient Checkpointing | Hard | Meta, Google, NVIDIA, Tesla | Q / S |
| 27 | Implement GRPO (DeepSeek-R1) | Expert | DeepMind, Anthropic, OpenAI | Q / S |
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 12 | Implement KV Cache | Medium | Anthropic, OpenAI, Meta, Perplexity | Q / S |
| 18 | Implement Speculative Decoding | Hard | Google, DeepMind, Anthropic, Apple | Q / S |
| 19 | Implement Continuous Batching | Hard | Perplexity, Together AI, Anyscale, Meta | Q / S |
| 28 | Build a Complete LLM Inference Engine | Expert | Perplexity, Together AI, Anyscale, Fireworks AI | Q / S |
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 13 | Sliding Window Attention | Medium | Mistral, Anthropic, Google, DeepMind | Q / S |
| 17 | Mixture of Experts Layer | Hard | Google, DeepMind, Mistral, Databricks, xAI | Q / S |
| 20 | DDPM (Denoising Diffusion) | Hard | Midjourney, Runway, Stability AI, Adobe, Google | Q / S |
| 21 | DDIM Sampling + Classifier-Free Guidance | Hard | Midjourney, Runway, Stability AI, Adobe | Q / S |
| 22 | Selective State Space Model (Mamba) | Hard | DeepMind, Google, Anthropic | Q / S |
| 23 | Vision Transformer + MAE Pretraining | Hard | Meta, Google, Apple, Tesla, Waymo | Q / S |
| 29 | Knowledge Distillation | Medium | Google, Apple, Meta, Qualcomm, Tesla | Q / S |
| # | Problem | Difficulty | Companies | Links |
|---|---|---|---|---|
| 24 | Fused Softmax Kernel in Triton | Expert | NVIDIA, Meta, Google, xAI, Tesla | Q / S |
| 25 | FlashAttention-2 in Triton | Expert | NVIDIA, Meta, Together AI, xAI | Q / S |
| 26 | FSDP (Fully Sharded Data Parallel) | Expert | Meta, Google, NVIDIA, Anthropic, xAI | Q / S |
| 30 | Ring Attention for Long Contexts | Expert | Anthropic, Google, Meta, xAI | Q / S |
"If I'm interviewing at X, which v3 questions should I do?"
| Company | Priority Questions |
|---|---|
| Anthropic | 5, 6, 7, 8, 10, 12, 13, 14, 15, 18, 22, 26, 27, 30 |
| OpenAI | 5, 7, 8, 10, 11, 12, 14, 15, 27 |
| DeepMind | 5, 6, 7, 8, 9, 13, 14, 15, 17, 18, 22, 27 |
| Meta | 1, 2, 3, 4, 9, 11, 12, 14, 15, 16, 19, 23, 24, 25, 26, 30 |
| 1, 2, 4, 9, 11, 13, 16, 17, 18, 20, 22, 23, 24, 26, 29, 30 | |
| Apple | 1, 5, 9, 18, 23, 29 |
| NVIDIA | 16, 24, 25, 26 |
| Midjourney / Runway / Stability AI | 5, 6, 20, 21 |
| Perplexity / Together AI / Anyscale | 7, 10, 12, 19, 25, 28 |
| Tesla / Waymo | 16, 23, 24, 29 |
| xAI | 17, 24, 25, 26, 30 |
| Mistral / Cohere | 7, 8, 10, 13, 17 |
Found a bug? Have a question from your own interview? PRs are welcome. Follow the notebook structure (question file + _SOLN file) and tag the authors.
If you found this helpful, follow me on Twitter. I post about ML interviews, PyTorch tips, and what I'm building next. Or just send me feedback, I read everything.
