Profiling training phases and verifying optimizations

Optimization without a phase breakdown optimizes the wrong line. The profile decides whether the step is data-starved, launch-bound, compute-bound or sync-bound, and the revalidation proves the faster run still learns.

measure

You can now

Annotate a training step and attribute time to data wait, transfer, forward, backward and optimizer phases.
Read self versus total time and memory, and split shared operators by input shape.
Revalidate learning and resume behavior after any optimization, and explain one model-versus-measured discrepancy.

Covers

Profiler annotations and phase timing
Data wait, transfer, forward, backward and optimizer decomposition
Correctness and learning checks after optimization

Study, in this order

4 steps

PyTorch Profiler recipe
PyTorch - reliability and performance recipes

Step 3 execution time with activities, record_function and key_averages self versus total; group_by_input_shape with record_shapes; Step 4 memory with profile_memory; Step 5 chrome trace export; Step 7 schedule, on_trace_ready and prof.step for long jobs

Done
Tuning guide: data flow and zeroing
PyTorch - reliability and performance recipes

Enable asynchronous data loading and augmentation with num_workers and pin_memory; Use parameter.grad = None instead of zeroing; Avoid unnecessary CPU-GPU synchronization

Done
Profiler inside a training loop
Learn PyTorch - profiler walkthrough

Adjust training function to track results with SummaryWriter: train_step with tensorboard_trace_handler, non_blocking transfers, autocast forward, and backward plus step inside the profiled range

Depth · Transfer the annotation pattern into your own compact model. Do not keep the tutorial's EfficientNet code, and do not copy its fp16 autocast loop, which has no GradScaler.

Done
Transformer Training FLOP Estimator
TensorTonic - mixed-precision memory and transformer FLOPs

Estimate one training step from forward matrix multiplications and a supplied forward attention cost

Done

Focus on

Annotate before measuring: one range per phase with record_function, plus DataLoader timing outside the step, so data wait cannot hide inside forward.
Self time excludes children and total includes them. A big total with small self points at the children, not at the named line.
Predict first, then profile. Write the expected slowest phase and the expected memory peak before running, and explain the residual in writing.
Every optimization revalidates learning: loss, gradients, controls and resume. A faster run that learns differently is a defect, not a win.

Reproduce closed-book

On CPU, profile one annotated step of a small MLP with forward, backward and optimizer ranges. Show the train_step total covers its children, export the chrome trace to a nonempty file, and write one sentence attributing the largest phase.

Traps

Profiling with stacks, anomaly detection or gradcheck on and calling it a measurement. Debug tooling distorts the run; measure the clean configuration.
Tuning the fastest phase. Rank phases first; a 2x win on 5 percent of the step is noise against an unexamined data wait.
Skipping learning checks after AMP, compile or fusion. Numerics move silently; rerun controls and resume checks before keeping the change.

Practice

TensorTonic: Transformer Training FLOP Estimator

Estimate one training step from forward matrix multiplications and a supplied forward attention cost

Solve it, then reconcile your estimate with the compact model's profiled step and write down every included and excluded term.

Self-test

key_averages shows forward total 55 percent but self 13 percent. Where is the rest?
One aten operator appears twice with different shapes. Which flags produce that split?
The memory table shows aten::empty with large self memory and convolutions with none. What does that mean?
After AMP, compile and a larger batch, loss matches but step time barely moves. Name the next two checks.

Completion criterion

Stop when

Your annotated profile attributes every phase, your prediction names the slowest one before the run, and one optimization ships with passing learning and resume checks.

Mark passed, then open 1.1 Only mark this when you can do it without notes.

Leads to

3.7 Objective checks, tiny-set overfit and shuffled-label controls Reuse the overfit and shuffled-label controls as learning checks.

Comes after

7.7 Architecture trade-offs and attention correctness testsProfile the verified block phase by phase.
8.2 Streams, asynchronous copies and honest timingApply event discipline inside annotated profiled phases.
8.4 Precision formats, autocast and gradient scalingMeasure whether narrower math actually moved the bottleneck.
8.5 Overflow, underflow and finding the first invalid valueRecord invalid events alongside phase timings in profiles.
8.6 Activation checkpointing and recomputationAttribute the recompute cost to the backward phase in a profile.
8.8 `torch.compile`: capture, guards, graph breaks and dynamic shapesMeasure compiled versus eager inside profiled phases.

Resources for this unit

PyTorch - reliability and performance recipes
PyTorchprimarycourse
Use the Profiler, Automatic Mixed Precision, model loading and Benchmark recipes, then the Performance Tuning Guide. Apply each to one compact model from an earlier section and keep a CPU fallback when CUDA is unavailable.
TensorTonic - mixed-precision memory and transformer FLOPs
TensorTonicsupportpractice
Solve only Mixed-Precision Training Memory Accountant and Transformer Training FLOP Estimator. Reconcile both with the compact model's profiler and document every included and excluded term.
Learn PyTorch - profiler walkthrough
Daniel Bourkeoptionalimplementation
Use notebook 10 only if the official profiler route remains unclear. Transfer the relevant annotations into the compact model instead of keeping the tutorial code.

Also in this section

PyTorch - compiler troubleshooting
PyTorchrequiredreference
Read graph breaks, guards, recompilations and logging. Compile one stable fragment, intentionally trigger a graph break or recompilation, inspect it, and compare eager and compiled outputs and gradients.
PyTorch - checkpointing, numerical accuracy and CUDA semantics
PyTorchrequiredreference
Use the official user-guide entries for torch.utils.checkpoint, numerical accuracy, CUDA semantics, reproducibility and serialization. Predict memory categories before measuring and define resume equivalence before running.
Deep Learning Training Debugger
Roadmap-Finalreferenceassessment
Use the failure-injection, measurement and evidence rules as the deeper standard. The full Tier 2 project is a later route and is not added to this stage milestone.
TrenTorch - precision, checkpointing and resume
TrenTorchsupportpractice
Solve Loss Scaling, FP16/BF16 Representable Range, Gradient Checkpointing and Resume-from-checkpoint in NumPy, then reproduce each behavior with the PyTorch API in the compact model.
Python os - `os.replace`
Python Software Foundationreferencereference
Read os.replace for the atomic-rename step of checkpoint writes: save to a temp file in the same directory, flush and fsync, then replace the final name.
Hands-On Machine Learning with Scikit-Learn and PyTorch
Aurelien Geron (O'Reilly, 2025)supporttheory
Read Appendix B, Reduced Precision Models and Mixed-Precision Training, for why float16 training underflows and what fp32 primary weights and loss scaling fix.