Hardware Optimization for ML

Make a trained model run fast and cheap in production — choose the right chip (CPU vs GPU vs TPU vs accelerators), beat the memory-bandwidth bottleneck, batch for throughput, drop precision (FP32 → FP16 → INT8), fuse and compile the graph with TensorRT/ONNX Runtime, and profile to find the real bottleneck instead of guessing.

Learn Hardware Optimization for ML in our free AI & Machine Learning course — a beginner-friendly interactive lesson with worked examples, a practice…

Part of the free AI & Machine Learning course at LearnCodingFast — hands-on lessons with examples you run in your browser, plus practice exercises and a quick quiz.

Moving goods isn't about owning the fastest engine — it's about picking the right vehicle and loading it well. Hardware optimization is the same: the model is the cargo, and your job is to deliver it as fast and cheaply as possible.

And before you change anything, you check the traffic report : profiling tells you whether you're stuck on the road (bandwidth), idling at the depot (data loading), or actually limited by the engine (compute). Optimizing the wrong thing is just driving faster down the wrong street.

A neural network is mostly one operation done billions of times: multiply a matrix of inputs by a matrix of weights . Which chip you run that on changes your speed by orders of magnitude, because the chips are built for different jobs.

Beginners assume a model is slow because the chip can't do enough maths (not enough FLOPs , floating-point operations per second). For large-model inference the opposite is usually true: the chip is starving , waiting for weights to arrive from memory. This is the memory-bandwidth bottleneck .

Here's why. To generate one token, an LLM has to read every weight in the model from memory, but only does a tiny bit of arithmetic with each one. So the limit isn't compute — it's how many bytes per second you can stream out of memory (the bandwidth ). A useful back-of-the-envelope estimate for single-stream speed is:

Two big consequences follow. First, shrinking the model (lower precision, Section 4) directly raises speed because there are fewer bytes to read. Second, batching (Section 3) helps because you read each weight once and reuse it for the whole batch — turning a bandwidth-bound job into a compute-bound one that keeps the cores busy.

Running one input at a time wastes the hardware: you pay the cost of reading the weights but only process a single item. Batching stacks many inputs and runs them together, so the expensive weight read is shared across the whole batch. Two metrics matter, and they pull in opposite directions:

The model is simple: a batch costs a fixed overhead (kernel launch, reading weights) plus a small per-item compute cost. Throughput is batch / time_per_batch . It climbs as the batch grows — until the hardware saturates, after which throughput flattens while latency keeps rising. The job is to pick the batch size that maximizes throughput within your latency budget.

Every weight is stored in a number format, and the format's bit width decides how many bytes you move and how fast the hardware runs. Cutting precision is the single easiest large speedup in inference.

Full precision. The training default, the most accurate, the slowest and largest. Keep it only for numerically sensitive layers.

Half the memory and bandwidth, ~2x faster on Tensor Cores, almost no accuracy loss. The sensible default for inference.

A quarter of the memory, often 2-4x faster. Needs calibration and an accuracy check, but huge for serving at scale.

Because speed scales with bytes read, going FP32 → FP16 roughly doubles tokens/sec and FP16 → INT8 roughly doubles it again. The catch is accuracy: lower precision rounds more, so always benchmark the quantized model on your task before shipping. You'll feel this effect directly in Your Turn #2.

A model is a graph of small operations — a matrix multiply, then add a bias, then a ReLU activation. Run naively, each operation launches its own GPU kernel (a unit of GPU work) and writes its result back to slow memory before the next one reads it. Those round-trips dominate the time.

Operator fusion merges a chain of ops into a single kernel that keeps the intermediate values in fast registers — e.g. matmul + bias + ReLU becomes one fused kernel, eliminating the memory round-trips. You rarely write fusion by hand; graph compilers do it for you:

The usual pipeline: PyTorch/TF model → export to ONNX → compile with the right backend for your hardware → deploy. The two worked examples below show torch.compile and a TensorRT FP16 export.

The most expensive mistake in this whole lesson is optimizing by guessing. The bottleneck is rarely where you assume — a slow model is often waiting on data loading or Python overhead, not compute. Profile first.

For GPU work, NVIDIA Nsight Systems shows whether the GPU is busy or idle and whether transfers overlap compute. Low GPU utilization with slow inference almost always means the GPU is waiting : tiny batches, data-loading stalls, or per-call Python overhead. Fix what the profiler points at — then measure again.

One line wraps your model and gives you fused kernels with no accuracy change. The first call compiles (slow); every call after runs the optimized graph.

The production pipeline in code: export to ONNX, then build a TensorRT engine with FP16 enabled. The FP16 flag is the speed switch that turns on half-precision kernels.

This is the batching trade-off with no libraries — just a per-batch overhead, a per-item cost, and the throughput each batch size delivers. Read the comments, then press run and watch throughput climb as the batch grows.

Fill in the two blanks marked ___ so the loop computes throughput ( items / time ) and remembers the batch size that wins. Check your output against the # ✅ Expected output comment.

Fill in the two blanks so the loop turns parameter count plus bytes-per-weight into model size, then estimates tokens/sec from memory bandwidth. Watch the numbers double as precision halves.

These five mistakes sink most first attempts at optimizing inference:

Serving a deep model on a CPU and wondering why each request takes seconds — the CPU does those huge matrix multiplies a tiny chunk at a time.

✅ Fix: move the model to a GPU ( model.to("cuda") ) and the inputs with it. Keep the CPU for tiny models, branchy logic, and preprocessing.

Looping over inputs one at a time in a data pipeline, leaving the GPU 90% idle while it reads the weights again for every single item.

✅ Fix: batch the inputs so the weight read is amortized. Use the largest batch that fits memory for offline work; use continuous batching for serving.

Pushing the batch size or sequence length until the GPU throws CUDA out of memory — activations and the KV cache grow with batch and length, not just the weights.

✅ Fix: lower the batch size, drop precision (FP16/INT8) to free VRAM, use gradient/activation checkpointing or paged attention, and leave headroom for the KV cache.

Rewriting model code for a week to speed it up, when a profile would have shown the GPU was simply starved by a slow data loader.

✅ Fix: profile before touching anything (PyTorch profiler, Nsight Systems), fix the single biggest hotspot, then measure again. Never optimize blind.

Quantizing straight to INT8 (or INT4) to chase speed and shipping it without checking — accuracy quietly collapses on maths, code, or rare classes.

✅ Fix: step down gradually (FP16 → INT8), calibrate properly, and benchmark accuracy on your real task at each step. Keep sensitive layers in higher precision.

Now combine throughput and latency with only a comment outline — no filled-in logic. Find the largest batch whose latency fits a 100 ms budget, and report its throughput.

Lesson 40 complete — you can make a model fly in production!

You can match a workload to the right chip (CPU, GPU, TPU, accelerator), explain why memory bandwidth — not FLOPs — caps LLM inference, batch inputs to maximize throughput and pick the best batch size, trade precision (FP32 → FP16 → INT8) for speed, fuse and compile graphs with torch.compile/ONNX Runtime/TensorRT, and profile to find the real bottleneck instead of guessing — while avoiding CPU-for-GPU work, batch-size-1 stalls, memory overflow, blind optimization, and precision dropped too low.

🚀 Up next: Distributed Training — scale a model across many GPUs and machines when one chip is no longer enough.

Practice quiz

• Browse all free courses
• Cheat sheets & quizzes
• Coding dictionary