Why Optimization and Efficiency matters for Applied Scientists
As an Applied Scientist, you own models from prototype to production. Optimizing training and inference unlocks faster iteration, lower cost, and stable performance under real-world constraints. Mastering this skill lets you:
- Hit latency/throughput SLOs without sacrificing too much quality.
- Scale experiments and training to larger data and models.
- Cut cloud spend and time-to-result with smarter training loops.
- Debug bottlenecks quickly and prioritize impactful fixes.
Typical tasks this skill unlocks
- Budgeted hyperparameter search that converges fast.
- Mixed precision training (AMP) and safe loss scaling.
- Int8/FP16 inference, model distillation, and pruning choices.
- DistributedDataParallel training and data pipeline tuning.
- Profiling to eliminate compute, I/O, and memory hotspots.
- Designing cost–latency–quality tradeoffs and guardrails.
Practical learning path (milestones)
- Profile first: Measure baseline throughput, latency, GPU/CPU utilization, and memory. Identify top bottleneck before any change.
- Quick wins: Enable mixed precision, pin memory, prefetch data, and try gradient accumulation to fit larger batches.
- Budgeted HPO: Run a small random search with early stopping. Log results and rank by validation metric per cost.
- Quantize and distill for inference: Try FP16/Int8 variants and a compact student model trained from teacher logits.
- Distribute when needed: Move to multi-GPU with DDP and a distributed sampler. Monitor scaling efficiency.
- Scale experiments cleanly: Use fixed seeds, one-variable-at-a-time changes, and simple dashboards for cost/latency/quality curves.
Milestone checklist
- Have a repeatable baseline run with saved metrics.
- AMP enabled with stable loss scaling.
- Data loader is not the bottleneck (verified by profiler).
- Random search over 3–5 key hyperparameters completed.
- One quantized or distilled model meets latency SLO.
- DDP run achieves >70% scaling efficiency (where applicable).
Worked examples
1) Budgeted random search (PyTorch, minimal)
import random, math, time
import torch, torch.nn as nn, torch.optim as optim
def make_model(width):
return nn.Sequential(
nn.Linear(784, width), nn.ReLU(),
nn.Linear(width, 10)
)
def train_one_epoch(model, loader, opt, device):
model.train()
loss_fn = nn.CrossEntropyLoss()
total = 0
for x, y in loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
logits = model(x)
loss = loss_fn(logits, y)
loss.backward()
opt.step()
total += loss.item()
return total / max(1, len(loader))
def validate(model, loader, device):
model.eval()
correct = 0
count = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
logits = model(x)
pred = logits.argmax(dim=1)
correct += (pred == y).sum().item()
count += y.numel()
return correct / max(1, count)
# Hypers to search
search_space = {
'lr': lambda: 10 ** random.uniform(-4.5, -2.5),
'width': lambda: random.choice([64, 128, 256, 512]),
'wd': lambda: 10 ** random.uniform(-6, -2)
}
# dummy loaders assumed
train_loader, valid_loader = ..., ...
device = 'cuda' if torch.cuda.is_available() else 'cpu'
best = {'acc': 0.0}
BUDGET = 20 # trials
PATIENCE = 2 # early stop epochs without val improvement
for t in range(BUDGET):
lr = search_space['lr'](); width = search_space['width'](); wd = search_space['wd']()
model = make_model(width).to(device)
opt = optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
best_epoch_acc, bad = 0.0, 0
for epoch in range(10):
train_one_epoch(model, train_loader, opt, device)
acc = validate(model, valid_loader, device)
if acc > best_epoch_acc:
best_epoch_acc, bad = acc, 0
else:
bad += 1
if bad >= PATIENCE: # early stop under-performers
break
if best_epoch_acc > best['acc']:
best = {'acc': best_epoch_acc, 'lr': lr, 'width': width, 'wd': wd}
# save model weights as needed
print('Best trial:', best)
2) Mixed precision training with safe loss scaling
import torch
from torch.cuda.amp import autocast, GradScaler
model = ...; optimizer = ...; loss_fn = ...
device = 'cuda'
model.to(device)
scaler = GradScaler(enabled=True)
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad(set_to_none=True)
with autocast():
logits = model(x)
loss = loss_fn(logits, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
Tip: Keep reductions and softmax-cross-entropy inside autocast; keep layer norm in FP32 if instability appears.
3) Fit larger batches: gradient accumulation
acc_steps = 4
optimizer.zero_grad(set_to_none=True)
for i, (x, y) in enumerate(train_loader):
x, y = x.to(device), y.to(device)
with autocast():
loss = loss_fn(model(x), y) / acc_steps
scaler.scale(loss).backward()
if (i + 1) % acc_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
4) Basic knowledge distillation (teacher → student)
import torch.nn.functional as F
temperature = 2.0
alpha = 0.7 # weight on KD vs hard labels
teacher.eval()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
with torch.no_grad():
t_logits = teacher(x) / temperature
s_logits = student(x) / temperature
kd_loss = F.kl_div(
F.log_softmax(s_logits, dim=1),
F.softmax(t_logits, dim=1),
reduction='batchmean'
) * (temperature ** 2)
ce_loss = F.cross_entropy(student(x), y)
loss = alpha * kd_loss + (1 - alpha) * ce_loss
loss.backward(); optimizer.step(); optimizer.zero_grad()
5) Profile to find data bottlenecks
import time
start = time.time()
for _ in range(5):
for x, y in train_loader:
break
print('Time to first batch (cold):', time.time() - start)
# Torch profiler (PyTorch >=1.8)
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
for i, (x, y) in enumerate(train_loader):
x, y = x.to(device), y.to(device)
out = model(x)
loss = loss_fn(out, y)
loss.backward()
if i > 20:
break
print(prof.key_averages().table(sort_by='self_cuda_time_total'))
Use the table to confirm whether kernels (compute) or DataLoader/CPU are dominant.
6) Minimal DistributedDataParallel (DDP) skeleton
# Launched via: torchrun --nproc_per_node=4 train_ddp.py
import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
rank = int(os.environ['RANK']); world = int(os.environ['WORLD_SIZE'])
dist.init_process_group('nccl')
device = rank % torch.cuda.device_count()
torch.cuda.set_device(device)
model = ... .to(device)
model = DDP(model, device_ids=[device])
sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True)
loader = torch.utils.data.DataLoader(dataset, batch_size=..., sampler=sampler, num_workers=..., pin_memory=True)
for epoch in range(epochs):
sampler.set_epoch(epoch)
for x, y in loader:
x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
loss = loss_fn(model(x), y)
loss.backward(); optimizer.step(); optimizer.zero_grad()
dist.destroy_process_group()
Drills and exercises
- Turn on AMP and measure speedup and memory savings vs FP32 (report both).
- Increase DataLoader num_workers and pin_memory; record step time before/after.
- Run a 15-trial random search on lr, weight_decay, and dropout; plot best metric vs trial number.
- Quantize a trained model to FP16 or Int8; measure accuracy drop and latency change.
- Distill a small student; compare accuracy–latency to the teacher.
- Run a 2×GPU DDP job; compute scaling efficiency = (throughput_2gpus)/(2×throughput_1gpu).
Common mistakes and quick debugging tips
- Optimizing blindly: Always profile first; chase the top bottleneck only.
- AMP instability: If loss becomes NaN, ensure scaler.step/update are used, reduce LR, or keep sensitive layers in FP32.
- Data loader stalls: Small num_workers, heavy CPU transforms, or slow storage. Fix with prefetching, caching, pin_memory, or converting to faster formats.
- Distributed under-utilization: Missing DistributedSampler or imbalanced batches causes stragglers.
- Over-sweeping hypers: Start narrow. Log-scale for LR/WD. Early-stop bad trials to save budget.
- Quantization surprises: Post-training Int8 may need calibration data; per-channel quant can help with accuracy.
Debugging mini-tasks
- Plot GPU utilization over time; if low with high CPU, tune DataLoader.
- Check P95 latency and batch size; micro-batches may increase tail latency.
- Verify deterministic seeds when comparing runs.
Mini project: Meet a latency SLO without losing more than 1% accuracy
Goal: Take an existing model that passes accuracy requirements, then bring P95 inference latency under a target (e.g., 80 ms) with ≤1% absolute accuracy drop.
- Baseline: Measure P50/P95 latency and accuracy. Save metrics and model hash.
- Quick wins: Enable FP16 or Int8 where supported. Batch or pre-allocate tensors to avoid overheads.
- Model shrink: Distill to a smaller student (or reduce width/depth). Re-measure accuracy and latency.
- Data path: Preprocess features to a compact format and warm up the model to reduce cold-start.
- Guardrails: Establish a simple acceptance test that checks latency and accuracy thresholds on each run.
Acceptance checklist
- P95 latency ≤ target on test workload.
- Accuracy drop ≤ 1% absolute vs baseline.
- Reproducible script that prints both metrics.
- Short write-up explaining tradeoffs taken.
Practical projects (apply multiple subskills)
- Train-to-Serve pipeline: Build a notebook + script that trains with AMP, runs HPO for LR/WD, exports FP16 weights, and benchmarks inference latency.
- Throughput sprint: Start with a slow training loop; after profiling, implement 3 optimizations (data, compute, memory) and report speedup with ablations.
- Scaling sweep: Measure validation loss vs model size on 3 widths and 3 dataset subsizes. Fit a simple scaling curve and discuss diminishing returns.
Cost–Latency–Quality tradeoffs
Define your budget and SLOs first. Draw a simple curve: x-axis = cost/latency, y-axis = quality. Choose the knee point, add guardrails (min accuracy, max P95 latency), and document the rationale.
Mini task: Draw your tradeoff curve
- Collect (latency, accuracy) for 3 model variants.
- Mark feasible region under SLOs.
- Pick the knee point and justify in one sentence.
Scaling experiments
Scale carefully: change one factor at a time, fix seeds, and keep runs short but comparable. Track throughput, time-to-accuracy, and cost per improvement.
- Report scaling efficiency when adding GPUs: efficiency = speedup / GPUs.
- Use smaller proxies (subset data, fewer steps) to rank ideas quickly before full runs.
Subskills
- Hyperparameter Optimization: Run budgeted sweeps with early stopping and log-scale sampling for LR/WD.
- Efficient Training Techniques: Apply AMP, gradient accumulation, checkpointing, and data prefetching.
- Mixed Precision And Quantization Basics: Use FP16/BF16 for training; apply FP16/Int8 for inference.
- Distillation Basics: Train a compact student using teacher logits and temperature scaling.
- Distributed Training Basics: Launch DDP with proper initialization and distributed sampling.
- Profiling And Bottleneck Analysis: Use timers/profilers to find and fix the top bottleneck first.
- Cost Latency Quality Tradeoffs: Frame SLOs and choose the best knee point on the tradeoff curve.
- Scaling Experiments: Design controlled sweeps and interpret scaling efficiency.
Who this is for
Applied Scientists and ML engineers who train, optimize, and ship models to production, and anyone responsible for meeting accuracy and latency SLOs.
Prerequisites
- Comfort with Python and a DL framework (e.g., PyTorch).
- Basic training loops, validation metrics, and GPU fundamentals.
- Familiarity with batch sizes, learning rates, and regularization.
Next steps
- Finish the drills and mini project above.
- Take the skill exam to check mastery. Everyone can take it; logged-in users get saved progress.
- Apply one technique to an active work project and measure impact with a before/after report.