14.4 Tuning Models for Production Training & Serving
Key Takeaways
- Quantization lowers numeric precision (for example, to 8-bit or 4-bit weights) to cut memory and latency at a possible accuracy cost that must be measured.
- Knowledge distillation trains a smaller student model to mimic a larger teacher, lowering serving cost and latency.
- For TensorFlow models, the optimized TensorFlow runtime and server-side request batching can reduce serving cost and latency without code changes.
- LLM serving optimizations on Agent Platform and GKE include continuous batching, PagedAttention, prefix caching, speculative decoding, tensor parallelism, and KV-cache quantization.
- Production training efficiency comes from mixed precision, efficient input pipelines, checkpointing, warm-starting from previous models, and right-sized distributed training.
The exam guide ends its serving section with tuning ML models for training and serving in production. Here "tuning" means performance engineering: fitting production latency, throughput, memory, and cost budgets without losing more accuracy than the business accepts.
Start with Measurements
- Define service-level objectives (SLOs): for example, p95 latency under 50 ms at 1,500 QPS, cost under $0.20 per 1,000 predictions, and AUC drop under 0.5 points.
- Profile where time goes: request parsing, feature lookup, preprocessing, model forward pass, postprocessing, network.
- Load-test candidates on staging endpoints with production-like payloads.
- Change one thing at a time, and re-check accuracy on the evaluation set after every optimization.
Model-Level Optimizations
| Technique | What it does | Typical gain | Risk |
|---|---|---|---|
| Quantization | Store and compute with lower precision (float16, int8, 4-bit) | Smaller memory footprint, higher throughput, possibly fewer GPUs | Accuracy loss, especially at 4-bit, so evaluate |
| Pruning | Remove low-importance weights or neurons | Smaller models, sometimes faster | Needs fine-tuning. Speedups depend on hardware support |
| Distillation | Train a small student to mimic a large teacher | Large latency and cost reductions | Student may lose quality on hard cases |
| Architecture choice | Smaller variants, fewer layers, efficient architectures | Fundamental latency improvement | Requires retraining |
| Feature reduction | Drop low-value, expensive-to-compute features | Faster preprocessing and lookups | Slight accuracy change |
| Compilation | Graph optimization (XLA, TensorRT) for target hardware | Faster kernels | Build complexity, hardware-specific artifacts |
Example from GKE's LLM guidance: quantizing Llama-2 13B with GPTQ, or Gemma 7B with AWQ, lets each model serve on one L4 GPU instead of two.
Serving-Level Optimizations
Predictive models
- Optimized TensorFlow runtime: a drop-in serving image with lower cost and latency than open-source TensorFlow Serving, and no code changes.
- Server-side batching: TensorFlow prebuilt containers read
config/batching_parameters_confignext to the SavedModel. Batching raises GPU utilization and throughput at a small latency cost. - Right hardware: trees on CPU, deep models on GPU or TPU (Section 14.3).
- Caching: cache predictions for repeated inputs (such as popular product pages) and cache static lookups in the container.
- Precompute: move expensive features to Feature Store or batch jobs.
- Co-hosting and scale to zero for low-traffic models, to cut idle cost.
LLM serving
| Technique | Effect |
|---|---|
| Continuous batching (vLLM) | Combines incoming requests into running batches to maximize GPU utilization and throughput |
| PagedAttention (vLLM) | Manages KV-cache memory efficiently for high concurrency |
| Prefix caching | Reuses computation for shared prompt prefixes, such as the same long document or system prompt, and lowers time to first token. Agent Platform's vLLM adds host-memory prefix caching |
| Speculative decoding | A fast draft proposes tokens that the main model verifies, lowering time per output token |
| Tensor parallelism | Splits a large model across GPUs so it fits and serves faster |
| KV-cache quantization (for example, FP8) | Shrinks KV-cache memory, allowing larger batches. Can reduce accuracy |
| Weight quantization (AWQ, GPTQ) | Fits models on fewer or smaller GPUs |
For Google's managed Gemini models, the equivalent levers are model tier, context caching, output limits, thinking level, and consumption options (Chapter 4).
Training-Side Tuning for Production Retraining
Retraining pipelines run repeatedly, so efficiency compounds:
| Lever | Benefit |
|---|---|
| Mixed precision training | Faster steps and lower memory on modern GPUs and TPUs |
| Efficient input pipelines (sharded files, parallel reads, prefetch) | Keeps accelerators busy |
| Warm starting from the previous model's weights or checkpoint | Fewer epochs for incremental retraining on new data |
| Right-sized distributed training and Reduction Server | Shorter wall-clock time without wasted accelerators |
| Spot VMs with checkpointing | Lower cost for fault-tolerant retraining |
| Pipeline caching (Chapter 15) | Skips unchanged steps such as preprocessing on unchanged data |
| Hyperparameter search reuse | Narrow the search space from previous studies instead of starting fresh |
Trading Accuracy for Latency and Cost
Use a structured comparison:
| Candidate | AUC | p95 latency | Cost / 1K predictions | Meets SLO? |
|---|---|---|---|---|
| Baseline DNN on CPU | 0.912 | 140 ms | $0.31 | No (latency) |
| Same DNN on L4 with batching | 0.912 | 38 ms | $0.22 | No (cost) |
| Distilled student on CPU | 0.906 | 22 ms | $0.08 | Yes |
| int8-quantized DNN on L4 | 0.909 | 25 ms | $0.15 | Yes |
Several options may meet the SLOs. Choose based on the accuracy the business needs and on operational simplicity. Here, the distilled student gives the biggest cost saving for a 0.006 AUC drop.
Guarding Against Regressions
Optimizations can quietly change model behavior. Build them into the release process:
- Keep the unoptimized model as a reference, and compare predictions on a golden dataset. Flag rows where outputs differ beyond tolerance.
- Evaluate per slice (region, device, language) as well as overall, because quantization errors can concentrate in rare segments.
- Roll out optimized models with a canary (Chapter 13), and watch both latency and prediction distributions.
- Record the optimization settings (precision, batch configuration, runtime image) as part of the model version in Model Registry.
Exam Traps
- Buying more GPUs when profiling shows preprocessing or feature lookups dominate latency.
- Deploying a quantized model without re-evaluating accuracy, including per-slice results.
- Optimizing average latency when the SLO is p95 or p99.
- Forgetting that batching adds queueing delay. Tune batch timeouts for your latency budget.
A DNN meets accuracy targets but has a p95 latency of 180 ms against a 40 ms SLO on CPU replicas. Profiling shows the forward pass dominates. Which option is most likely to meet the SLO with limited accuracy loss?
A chat application sends the same 20,000-token policy document as context before every user question to a self-deployed open LLM, and time to first token is high. Which serving optimization targets this directly?
An LLM needs two L4 GPUs per replica because its weights don't fit on one. The team wants to cut serving cost and accepts a small, measured accuracy trade-off. What should they try first?