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.

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.
• Watch the logbooks. Others have already hit your instability; training chronicles document the loss spikes and the fixes. See LLM/VLM chronicles.

Compute / accelerators

Full chapter: Compute · Accelerators.

• The number that matters is achievable, not theoretical. Real matmul FLOPS are well below spec. Measure Maximum Achievable Matmul FLOPS with mamf-finder.py; compute MFU (model FLOPS utilization) to compare setups. See the most important thing to understand and what accelerator characteristics we care for.
• Memory capacity + bandwidth often bind before FLOPS. Check the accelerator memory size and speed table and the TFLOPS comparison table when choosing hardware.
• Don't forget power and cooling - throttling silently caps sustained throughput. See power and cooling.
• CPU and CPU memory matter too (dataloading, offload, pinned memory). See CPU and CPU memory.

Storage (IO)

Full chapter: Storage.

• 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 and local storage beats cloud storage.
• Benchmark IO before you trust it (checkpoint save/load and dataset streaming are common stalls). See storage 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. Clean up instead of overpaying - see why pay for more storage.

Network

Full chapter: Network.

• 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.
• Know the two speeds: intra-node (NVLink/PCIe) vs inter-node (InfiniBand/RoCE/Ethernet), and RDMA.
• Benchmark real collective throughput, not marketing numbers, with all_reduce_bench.py (far simpler than nccl-tests). Mind the important nuances (e.g. unidirectional vs bidirectional, payload size, busbw vs algbw).
• When collectives hang or crawl: see network debug and NCCL performance debug.

Orchestration & SLURM

Full chapter: Orchestration · SLURM · Kubernetes.

• Verify the cluster before the big run: every GPU on every node must talk to every other. Run torch-distributed-gpu-test.py across all nodes first.
• SLURM day-to-day: the users cheatsheet 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.

Training: performance & memory

Full chapter: Performance.

• Start from the checklist: 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 (weights + grads + optimizer states + activations) and profile with memory profiler tools.
• Free throughput wins: keep tensor dims divisible/aligned for Tensor Cores, set NUMA affinity, stop the DataLoader from starving the GPUs, try torch.compile, and tame automatic garbage collection jitter in synchronized runs.
• Model parallelism (DP/TP/PP/sequence/ZeRO) concepts and trade-offs: model parallelism.

Training: stability (instabilities & NaN/Inf)

Full chapter: Instabilities.

• Loss spikes/divergence are expected at scale. Compare against public training logbooks - your symptom is probably documented with a known mitigation.
• Numerical hygiene: sane weight init / STD, watch for numerical instabilities, and bad data-batch × parameter-state combinations.
• Catch NaN/Inf early: underflow/overflow detection and PyTorch tensor debugging (fp16/bf16 range issues, per-tensor min/max/NaN scans).
• Reproducibility for isolating a divergence: reproducibility.

Training: fault tolerance & checkpoints

Full chapter: Fault tolerance · Checkpoints.

• Provision slack: always plan for more nodes than needed so a dead node doesn't stop the run; queue up multiple jobs for auto-continuation.
• Checkpoint often enough that a crash costs minutes, not hours - see frequent checkpoint saving and prevention.
• Operator controls: a kill switch to stop a doomed run cleanly and a save switch to force a checkpoint on demand.
• Survive the cluster: handle forced job preemption and prefer fixed over dynamic allocations.

Inference

Full chapter: Inference.

• Speak the metrics: TTFT, TPOT/ITL, throughput vs latency, and how batching trades them off. See key inference performance metrics and concepts.
• Size the memory: weights + KV cache + activations - see anatomy of model's memory usage.
• Cut cold starts: speeding up model loading time; pick a serving stack from inference frameworks and validate with benchmarks.

Debugging distributed / PyTorch at scale

Full chapter: Debugging · PyTorch.

• Iterate cheaply: shrink to tiny models/tokenizers/datasets and fast debug of PyTorch models.
• Multi-node/multi-GPU hang or deadlock (the classic scale bug):
  1. Rule out comms with torch-distributed-gpu-test.py.
  2. Dump every rank's Python stack at once with py-spy; ranks stuck at different collectives reveal the desync. See diagnosing crashes, hangs and tracing execution and hanging solutions.
  3. Set NCCL_DEBUG=INFO; escalate to NCCL performance debug.
• Cryptic CUDA errors: CUDA_LAUNCH_BLOCKING=1 for a real traceback (CUDA is async). See The Art of Debugging - PyTorch.
• OOM / memory: memory usage; fragmentation via PYTORCH_CUDA_ALLOC_CONF.
• NaN/Inf: underflow and overflow detection and debugging tensors.
• GPU-specific faults: troubleshooting NVIDIA GPUs.

Testing

Full chapter: Testing.

• Run tests surgically (select, parametrize, repeat, control output/parallelism): running tests.
• Write robust tests (fixtures, temp dirs, RNG control for reproducibility, distributed tests): writing tests.
• When a test misbehaves: debugging tests.

Key tools

Pick the fix by symptom

Notes for AI agents

• Diagnose before optimizing. Identify which resource (compute/memory/network/storage/dataloader) is the actual bottleneck with a measurement; don't tune blindly.
• Prefer measured numbers over spec sheets. Use the provided benchmark scripts on the target hardware/software stack before recommending changes or purchases.
• Assume failures at scale. For any long/large run, verify checkpointing, restart, spare capacity, and a kill switch exist before worrying about peak speed.
• Verify the cluster first. Run the distributed connectivity test before blaming model code for a multi-node problem.
• Reuse the community's hard-won lessons. Check the training logbooks for known instabilities and fixes before re-deriving them.
• Read the linked chapter section before applying a recipe - each has worked examples, exact commands, caveats, and scripts.
• For deep single-process/tool debugging (gdb, strace, py-spy, cProfile, core files), use the companion skill: The Art of Debugging.