ml-engineering

# Machine Learning Engineering Distilled from **Machine Learning Engineering Open Book** by Stas Bekman - source: https://github.com/stas00/ml-engineering (CC BY-SA 4.0). Know-how from training BLOOM-176B, IDEFICS-80B and production RAG and RL training and inference systems. This skill is a condensed index; each section links back to the full chapter for depth, scripts, and benchmarks. A field-tested, end-to-end guide to **training and serving large models** (LLMs, VLMs, multi-modal, RAG) on real hardware at scale - distilled from actually training BLOOM-176B and IDEFICS-80B and building production inference/RAG systems. It is a practitioner's brain dump: opinionated guidance backed by copy-paste scripts, benchmark tools, and comparison tables, written for the engineers and operators who have to make expensive clusters actually deliver a finished model. It spans the entire stack that decides whether a run succeeds and how much it costs: selecting and *benchmarking* accelerators, storage, and network so the fast compute is never starved; orchestrating jobs with SLURM/Kubernetes; maximizing throughput (MFU) and fitting models in memory via parallelism (DP/TP/PP/ZeRO), activation recomputation, and offload; keeping training numerically stable through loss spikes and NaN/Inf; and surviving the *inevitable* hardware and job failures with frequent checkpointing, spare capacity, and automatic restarts. On the serving side it covers inference latency/throughput/cost trade-offs, KV-cache and memory sizing, and framework selection - plus diagnosing multi-node/multi-GPU hangs and testing the whole thing. Use it as an operator's runbook: figure out which resource is *actually* the bottleneck (compute? memory? network? storage? dataloader?), then jump to the targeted recipe. For pure debugging technique (gdb/strace/py-spy/CUDA), pair this with [The Art of Debugging](https://github.com/stas00/the-art-of-debugging/blob/master/SKILL.md). ## Core principles - **Measure, don't assume.** Vendor/theoretical TFLOPS are marketing; benchmark *your* hardware and software stack before optimizing or buying. Track MFU/throughput, not vibes. - **Find the actual bottleneck.** A training step is gated by the slowest of: accelerator compute, memory bandwidth/capacity, inter/intra-node network, storage IO, or CPU dataloading. Optimizing anything else is wasted effort. - **At scale, failure is the steady state.** With hundreds/thousands of accelerators, hardware *will* fail mid-run. Design for frequent checkpoints, automatic restarts, spare nodes, and kill/save switches from day one. - **Reproduce small and fast.** Debug on a tiny model / few layers / one node before burning cluster time - see [The Art of Debugging](https://github.com/stas00/the-art-of-debugging/blob/master/SKILL.md). - **Watch the logbooks.** Others have already hit your instability; training chronicles document the loss spikes and the fixes. See [LLM/VLM chronicles](https://github.com/stas00/ml-engineering/tree/master/resources#publicly-available-training-llmvlm-logbooks). ## Compute / accelerators Full chapter: [Compute](https://github.com/stas00/ml-engineering/tree/master/compute) · [Accelerators](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md). - **The number that matters is achievable, not theoretical.** Real matmul FLOPS are well below spec. Measure Maximum Achievable Matmul FLOPS with [`mamf-finder.py`](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/benchmarks/mamf-finder.py); compute **MFU** (model FLOPS utilization) to compare setups. See [the most important thing to understand](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md#the-most-important-thing-to-understand) and [what accelerator characteristics we care for](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md#what-accelerator-characteristics-do-we-care-for). - **Memory capacity + bandwidth often bind before FLOPS.** Check the [accelerator memory size and speed table](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md#accelerator-memory-size-and-speed) and the [TFLOPS comparison table](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md#tflops-comparison-table) when choosing hardware. - **Don't forget power and cooling** - throttling silently caps sustained throughput. See [power and cooling](https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/README.md#power-and-cooling). - **CPU and CPU memory matter too** (dataloading, offload, pinned memory). See [CPU](https://github.com/stas00/ml-engineering/tree/master/compute/cpu) and [CPU memory](https://github.com/stas00/ml-engineering/tree/master/compute/cpu-memory). ## Storage (IO) Full chapter: [Storage](https://github.com/stas00/ml-engineering/blob/master/storage/README.md). - **Pick the file system for the job:** distributed/parallel FS for shared checkpoints and datasets; fast **local** NVMe beats network storage for scratch/hot data. See [which file system to choose](https://github.com/stas00/ml-engineering/blob/master/storage/README.md#which-file-system-to-choose) and [local storage beats cloud storage](https://github.com/stas00/ml-engineering/blob/master/storage/README.md#local-storage-beats-cloud-storage). - **Benchmark IO before you trust it** (checkpoint save/load and dataset streaming are common stalls). See [storage benchmarks](https://github.com/stas00/ml-engineering/blob/master/storage/README.md#benchmarks). - **Gotchas that bite at scale:** you often get less usable capacity than sold; some clouds put backups on the same partition; always keep [checksums](https://github.com/stas00/ml-engineering/blob/master/storage/README.md#dont-forget-the-checksums). Clean up instead of overpaying - see [why pay for more storage](https://github.com/stas00/ml-engineering/blob/master/storage/README.md#why-pay-for-more-storage-when-you-can-easily-clean-it-up-instead). ## Network Full chapter: [Network](https://github.com/stas00/ml-engineering/blob/master/network/README.md). - **Inter-node speed can dominate the whole training's speed.** For sharded/parallel training, slow inter-node links stall everyone. Understand why before scaling out: [why inter-node speed is hugely important](https://github.com/stas00/ml-engineering/blob/master/network/README.md#understanding-why-inter-node-network-speed-is-of-a-huge-importance). - **Know the two speeds:** [intra-node](https://github.com/stas00/ml-engineering/blob/master/network/README.md#intra-node-networking) (NVLink/PCIe) vs [inter-node](https://github.com/stas00/ml-engineering/blob/master/network/README.md#inter-node-networking) (InfiniBand/RoCE/Ethernet), and [RDMA](https://github.com/stas00/ml-engineering/blob/master/network/README.md#rdma-networking). - **Benchmark real collective throughput,** not marketing numbers, with [`all_reduce_bench.py`](https://github.com/stas00/ml-engineering/blob/master/network/benchmarks/all_reduce_bench.py) (far simpler than nccl-tests). Mind the [important nuances](https://github.com/stas00/ml-engineering/blob/master/network/README.md#important-nuances) (e.g. unidirectional vs bidirectional, payload size, busbw vs algbw). - **When collectives hang or crawl:** see [network debug](https://github.com/stas00/ml-engineering/tree/master/network/debug) and [NCCL performance debug](https://github.com/stas00/ml-engineering/blob/master/debug/nccl-performance-debug.md). ## Orchestration & SLURM Full chapter: [Orchestration](https://github.com/stas00/ml-engineering/tree/master/orchestration) · [SLURM](https://github.com/stas00/ml-engineering/tree/master/orchestration/slurm) · [Kubernetes](https://github.com/stas00/ml-engineering/tree/master/orchestration/kubernetes). - **Verify the cluster before the big run:** every GPU on every node must talk to every other. Run [`torch-distributed-gpu-test.py`](https://github.com/stas00/ml-engineering/blob/master/debug/torch-distributed-gpu-test.py) across all nodes first. - **SLURM day-to-day:** the [users cheatsheet](https://github.com/stas00/ml-engineering/blob/master/orchestration/slurm/users.md) covers `sbatch`/`srun`/`salloc`, job arrays, dependencies, and inspecting the queue; keep the allocation and re-`srun` for fast debug iterations. - **Launchers** (torchrun/accelerate/deepspeed under SLURM): see [launchers](https://github.com/stas00/ml-engineering/tree/master/orchestration/slurm/launchers/README.md). ## Training: performance & memory Full chapter: [Performance](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md). - **Start from the checklist:** [how to improve speed and save memory](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#how-to-improve-speed-and-save-memory) enumerates the high-impact levers (parallelism choice, activation checkpointing/recomputation, offload, fused kernels, mixed precision). - **Fit the model:** know the [anatomy of memory usage](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#anatomy-of-models-memory-usage) (weights + grads + optimizer states + activations) and profile with [memory profiler tools](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#memory-profiler-tools). - **Free throughput wins:** keep tensor dims [divisible/aligned](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#vector-and-matrix-size-divisibility) for Tensor Cores, set [NUMA affinity](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#numa-affinity), stop the [DataLoader](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#dataloader) from starving the GPUs, try [`torch.compile`](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#torchcompile), and tame [automatic garbage collection](https://github.com/stas00/ml-engineering/blob/master/training/performance/README.md#automatic-garbage-collection) jitter in synchronized runs. - **Model parallelism** (DP/TP/PP/sequence/ZeRO) concepts and trade-offs: [model parallelism](https://github.com/stas00/ml-engineering/tree/master/training/model-parallelism). ## Training: stability (instabilities & NaN/Inf) Full chapter: [Instabilities](https://github.com/stas00/ml-engineering/blob/master/training/instabilities/README.md). - **Loss spikes/divergence are expected at scale.** Compare against public [training logbooks](https://github.com/stas00/ml-engineering/blob/master/training/instabilities/README.md#learning-from-training-logbooks) - your symptom is probably documented with a known mitigation. - **Numerical hygiene:** sane [weight init / STD](https://github.com/stas00/ml-engineering/blob/master/training/instabilities/README.md#std-init), watch for [numerical instabilities](https://github.com/stas00/ml-engineering/blob/master/training/instabilities/README.md#numerical-instabilities), and bad [data-batch × parameter-state combinations](https://github.com/stas00/ml-engineering/blob/master/training/instabilities/README.md#bad-combination-of-data-batch-and-model-parameter-state). - **Catch NaN/Inf early:** [underflow/overflow detection](https://github.com/stas00/ml-engineering/blob/master/debug/underflow_overflow.md) and PyTorch [tensor debugging](https://github.com/stas00/ml-engineering/blob/master/debug/pytorch.md#debugging-tensors) (fp16/bf16 range issues, per-tensor min/max/NaN scans). - **Reproducibility** for isolating a divergence: [reproducibility](https://github.com/stas00/ml-engineering/tree/master/training/reproducibility). ## Training: fault tolerance & checkpoints Full chapter: [Fault tolerance](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md) · [Checkpoints](https://github.com/stas00/ml-engineering/tree/master/training/checkpoints). - **Provision slack:** [always plan for more nodes than needed](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#always-plan-to-have-more-nodes-than-needed) so a dead node doesn't stop the run; [queue up multiple jobs](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#queue-up-multiple-training-jobs) for auto-continuation. - **Checkpoint often enough** that a crash costs minutes, not hours - see [frequent checkpoint saving](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#frequent-checkpoint-saving) and [prevention](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#prevention). - **Operator controls:** a [kill switch](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#kill-switch) to stop a doomed run cleanly and a [save switch](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#save-switch) to force a checkpoint on demand. - **Survive the cluster:** handle [forced job preemption](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#dealing-with-forced-job-preemption) and prefer [fixed over dynamic allocations](https://github.com/stas00/ml-engineering/blob/master/training/fault-tolerance/README.md#preferring-fixed-accelerator-allocations-to-dynamic-ones). ## Inference Full chapter: [Inference](https://github.com/stas00/ml-engineering/blob/master/inference/README.md). - **Speak the metrics:** TTFT, TPOT/ITL, throughput vs latency, and how batching trades them off. See [key inference performance metrics](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#key-inference-performance-metrics) and [concepts](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#concepts). - **Size the memory:** weights + KV cache + activations - see [anatomy of model's memory usage](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#anatomy-of-models-memory-usage). - **Cut cold starts:** [speeding up model loading time](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#speeding-up-model-loading-time); pick a serving stack from [inference frameworks](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#inference-frameworks) and validate with [benchmarks](https://github.com/stas00/ml-engineering/blob/master/inference/README.md#benchmarks). ## Debugging distributed / PyTorch at scale