Keep it fed: loaders and async copies
A fast training step needs an equally deliberate supply chain from storage to pinned memory to the device.
Training repeatedly OOMs and misses its nightly completion window.
This chapter moves from separate wait from compute to verify sustained overlap. We will make the idea concrete in a small code example before placing it back inside a production system with 24 GB device and nightly completion.
Use a profiler timeline after warm-up. Healthy training shows copies prepared ahead of compute rather than wide blank gaps before every step.
Verify the behavior illustrated in fed_loader.py, then measure decision utility, slice reliability, and system latency.
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
An expensive GPU that waits for the next batch is a very fast machine doing nothing.
The input pipeline spans storage reads, decoding, augmentation, collation, host memory, and the host-to-device copy. If any stage is slower than the training step, gaps appear between GPU kernels. More data-loader workers help only when CPU preparation is the bottleneck; they can hurt when they contend for storage or duplicate large process state.
Pinned host memory enables DMA transfers, and a non-blocking copy lets the host enqueue work without waiting. Real overlap also requires the copy and compute to use compatible CUDA streams and enough queued work. Measure batch wait time and device utilization before changing worker count, prefetch depth, or batch size.
- HBM bandwidth
- arithmetic intensity
- pinned memory
- FP16/BF16
Follow the reasoning before touching the code.
- 01
Separate wait from compute
Time how long iteration spends obtaining a batch, copying it, and running forward/backward. A single end-to-end throughput number cannot identify starvation.
- 02
Tune one stage
Increase workers until preparation no longer dominates, then test pinned memory and non-blocking copies. Watch RAM, file handles, and storage throughput.
- 03
Verify sustained overlap
Use a profiler timeline after warm-up. Healthy training shows copies prepared ahead of compute rather than wide blank gaps before every step.
import torch
from torch.utils.data import DataLoader, TensorDataset
features = torch.randn(4096, 32)
targets = torch.randint(0, 2, (4096,))
dataset = TensorDataset(features, targets)
loader = DataLoader(dataset, batch_size=256, num_workers=2,
pin_memory=True, persistent_workers=True)
for host_x, host_y in loader:
device_x = host_x.to('cuda', non_blocking=True)
device_y = host_y.to('cuda', non_blocking=True)
assert device_x.shape[0] == device_y.shape[0]Prepare batches in pinned host memory and enqueue asynchronous device copies.
What to notice
- `non_blocking=True` permits an asynchronous copy because the loader produced pinned tensors; it does not by itself prove copy/compute overlap.
- `persistent_workers=True` avoids recreating worker processes each epoch but keeps their memory allocated between epochs.
YOUR TURN Pause and predict before revealing the answer.
GPU utilization is low, but increasing loader workers from 4 to 16 makes throughput worse. What should you inspect next?
Need a hint?
Workers share resources outside the GPU.
Check your reasoning
Inspect storage latency and bandwidth, CPU saturation, augmentation cost, worker memory, and batch wait time. The workers may be competing for the same disk or oversubscribing CPU cores.
Lessons 1 in this course.
Connect the chapter to the course’s quantitative spine.
A fast training step needs an equally deliberate supply chain from storage to pinned memory to the device. The roofline bound below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Roofline bound
| Symbol | Meaning / shape / unit |
|---|---|
P_peak | peak compute |
I | arithmetic intensity |
B_mem | memory bandwidth |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: P_peak, I, B_mem.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c12(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
GPU memory ledger
Change microbatch size and calculate a simplified training-memory budget.
Assumption: 3.2 GB weights+grads+optimizer, 0.34 GB activations per sample, 24 GB device, 10% reserve.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| block 1 | 1.09 GB activations | resident |
| block 2 | 2.18 GB activations | resident |
| block 3 | 3.26 GB activations | resident |
| block 4 | 4.35 GB activations | resident |
| block 5 | 5.44 GB activations | resident |
Trace the complete operating path.
Carry the chapter’s central idea into the system: A fast training step needs an equally deliberate supply chain from storage to pinned memory to the device. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Separate wait from compute” into a written invariant and a failing test.
- 02
Promote fed_loader.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Verify sustained overlap” as the release checkpoint: Use a profiler timeline after warm-up. Healthy training shows copies prepared ahead of compute rather than wide blank gaps before every step.
- 04
Serve or materialize behind explicit 24 GB device and nightly completion.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Single-GPU optimization plan
Profile and optimize a correctness-locked trainer without sacrificing model quality.
- Profiler diagnosis
- Mixed precision parity
- Memory plan
- Throughput/cost report
Read primary material with a purpose.
Return to the opening failure.
Record hardware/software, shapes, precision, seeds, throughput method, peak memory, and cost per sample.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce a profiled, mixed-precision, memory-budgeted training loop.