Unit 3.6 · about 5 h
190 / 324Schedules, gradient accumulation and clipping
Schedules, batch size and clipping jointly set the update scale. Each one silently changes what the learning rate means, so they must be reasoned about as one mechanism.
- implement
- measure
You can now
- Implement linear warmup followed by cosine decay with the scheduler stepped in the right order.
- Implement gradient accumulation that matches a full-batch gradient up to floating-point order, and state the effective batch size.
- Apply global-norm clipping after unscaling in mixed precision, and run an LR range test to set bounds.
Covers
- Warmup, cosine schedules and schedule order
- Gradient accumulation and effective batch size
- Gradient clipping by global norm
- Learning-rate range tests
Study, in this order
6 steps
- Schedules and warmupD2L - optimization and regularization slice
Schedulers and Policies: factor, multi-factor, cosine with formula, and warmup
- Gradient clipping by global normD2L - optimization and regularization slice
Gradient Clipping: Lipschitz bound, the \(\min(1, \theta/\|g\|)\) projection, and its limits
Depth · Read the Gradient Clipping section only. Skip the RNN implementation around it.
- Clip API semanticsPyTorch - optimization and automatic mixed precision examples
clip_grad_norm_page: global-norm definition over concatenated per-parameter norms, in-place edit, returned total norm,error_if_nonfinite - Scaler order with clippingPyTorch - optimization and automatic mixed precision examples
Adding
GradScalerand Inspecting/modifying gradients:autocastforward, scale, backward, unscale, clip, step, update,zero_grad - LR range testCyclical Learning Rates for Training Neural Networks (Smith, 2017)
Abstract and range-test description: linearly increasing the rate for a few epochs to read reasonable bounds
- HOML: PyTorch schedulers and gradient clippingHands-On Machine Learning with Scikit-Learn and PyTorch
Chapter 11, Gradient Clipping (printed pp. 382-383) and Learning Rate Scheduling from Cosine Annealing through 1cycle Scheduling (pp. 400-405)
Depth · The warm-up loop on p. 402 calls
warmup_scheduler.step()at the start of each epoch. On PyTorch 2.14 that trains at 40, 70 and 100 percent of the rate and never at the 10 percent the text promises; step at the end of each epoch to get 10, 40, 70, 100.
Focus on
- Cosine decay sets \(\eta_t\) from a closed form with fixed \(T\). Stepping the scheduler per step versus per epoch changes the curve, so record which one the code does.
- Accumulation adds \(K\) microbatch gradients and steps once. Scaling each microbatch's mean loss by \(1/K\) reproduces the full-batch mean gradient when the microbatches are equal in size and nothing updates between them. Batch-norm statistics and dropout masks still differ, and the summation order changes, so agreement is to floating-point tolerance.
- Scaled gradients must be unscaled before clipping. Clipping scaled values applies the threshold in the wrong units; the recipe unscales first so
max_normkeeps its unscaled meaning. - The range test raises the rate linearly and reads bounds off the loss: flat start, steady fall, then divergence. It sets a search interval, not an optimum.
Reproduce closed-book
Traps
- Calling
scheduler.step()beforeoptimizer.step(). PyTorch warns because the first scheduled rate is skipped; assert the rate used on step one. Forgettingzero_gradat the start of each accumulation cycle silently adds the previous cycle's gradients. - Claiming accumulation equals large-batch training. Batch-norm running statistics update per microbatch and dropout draws differ, so the match is statistical, not bitwise.
- Treating warmup as a universal fix. The Tuning Playbook shows instabilities that need clipping or architecture change; warmup only shortens the risky early phase.
Practice
TrenTorch: Gradient Accumulation
Implement compute_output_head_gradient, accumulate_gradients and train_with_gradient_accumulation; the tests compare the accumulated update with one full-batch step
Average the microbatch gradients, hold the weights fixed until every microbatch is done, then step once. Then check your PyTorch loop, which divides each microbatch loss by \(K\), against a full-batch reference within floating-point tolerance.
TrenTorch: Warmup and Cosine Decay
Implement linear_warmup_lr, cosine_decay_lr and warmup_cosine_lr, with the decay step counted from the end of warmup
Check the closed form first: base 0.3, minimum 0.01, 20 decay steps, step 10 gives 0.155.
D2L Learning-Rate Scheduling: Exercises
Exercises 1 and 4 (best fixed-rate model, warmup length); run the cosine comparison of exercise 3 on Fashion-MNIST instead of ImageNet
Self-test
-
Cosine schedule \(\eta_0\) 0.3, \(\eta_T\) 0.01, \(T\) 20. What is \(\eta\) at \(t\) 10 and \(t\) 20?
\(\eta_t = 0.01 + 0.145(1+\cos(\pi t/20))\). At \(t=10\), \(\cos(\pi/2)=0\), \(\eta=0.155\). At \(t=20\), \(\cos(\pi)=-1\), \(\eta=0.01\). -
Four microbatches of 8 with mean-reduced losses give microbatch gradients \(a\), \(b\), \(c\) and \(d\). What update matches batch-32 SGD, and when does it fail to match?
One optimizer step along \((a + b + c + d)/4\). It fails if you add the four gradients without dividing by 4 (a step 4 times too large), if the microbatches differ in size, or if state such as batch-norm statistics updates per microbatch. -
Where does clipping go in a
GradScalerstep, and what breaks if it moves before unscaling?Afterscaler.unscale_and beforescaler.step. Before unscaling, gradients carry the dynamic scale factor, so a fixedmax_normthreshold means something different every iteration. -
Range-test loss falls slowly, then steeply, then explodes. How do you read it?
The flat start is below the useful range, the steep fall is the usable interval, and the explosion is the ceiling. Search between the fall and below the ceiling; the test gives bounds, not the best rate.
Completion criterion
Stop when
Your accumulated run matches the full-batch gradient within floating-point tolerance, your clipped AMP order passes the unscale check, and your range test names defensible bounds.
Mark passed, then open 3.7 Only mark this when you can do it without notes.Leads to
- 3.7 Objective checks, tiny-set overfit and shuffled-label controls Check the objective the schedule is optimizing before tuning further.
- 8.4 Precision formats, autocast and gradient scaling See the full mixed-precision scaler contract behind the clipping order.
Comes after
- 1.4 Inner products, norms, orthogonality and Gram-SchmidtApply vector norms to clipping; leave optimizer exercises there.
- 5.3 Adaptive optimizers, schedules, clipping and stoppingApply schedule clocks and norm clipping to training.
- 3.5 Optimizers from SGD to AdamWSchedule the learning rate and scale the batch around the optimizer you just built.
Resources for this unit
- D2L - optimization and regularization slice Use MLP implementation, dropout, deep generalization, numerical stability, optimization introduction, SGD, minibatch SGD, momentum, AdaGrad, RMSProp, Adam and learning-rate scheduler. Implement momentum, AdamW, initialization, dropout, one normalization, clipping and one schedule.
- PyTorch - optimization and automatic mixed precision examples Check optimizer, scheduler, gradient clipping and AMP API semantics. Preserve the correct order for zeroing, scaling, backward, unscaling, clipping, stepping and updating.
- Cyclical Learning Rates for Training Neural Networks (Smith, 2017) Read the abstract and the range test: increase the rate linearly for a few epochs and read usable bounds off the loss curve.
- Hands-On Machine Learning with Scikit-Learn and PyTorch Read the Chapter 11 sections named in the units for the PyTorch side of initialization, normalization, clipping and schedules.
Also in this section
- Understanding Deep Learning - fitting, gradients and regularization Use Chapters 6-9 and targeted parts of Chapter 20 for optimization, initialization, measurement, regularization and why training works. Re-derive the update rules you keep in the trainer.
- Deep-ML candidate - optimization and normalization problems After the platform admission check, use at most one of 146 Momentum, 169 AdamW, 197 Global-Norm Clipping, 109 LayerNorm or 115 BatchNorm as a changed-prompt drill.
- TrenTorch - training mechanics questions Solve only the questions the units name: Weight Initialization, Layer Normalization, Label Smoothing, Gradient Accumulation, and Warmup and Cosine Decay. Read the tests first; they run in the browser IDE.
- PyTorch GroupNorm reference Read how input channels are split into
num_groupsgroups and which axes each group normalizes; comparenum_groups1 againstnum_channels. - PyTorch RMSNorm reference Read the RMS formula and which axes
normalized_shapecovers; note the missing mean centering and bias. - PyTorch torch.optim.AdamW reference Read the AdamW algorithm box for the decoupled weight-decay step, the bias-corrected moments and the defaults, and follow the Decoupled Weight Decay Regularization link for the argument.
- PyTorch CrossEntropyLoss reference Read Shape, Parameters (
weight,ignore_index,reduction,label_smoothing) and both Examples; extend the class-index example withignore_indexand a classweight. - Deep Learning Tuning Playbook FAQ: debugging optimization failures Read How can optimization failures be debugged and mitigated (identifying unstable workloads, warmup, clipping) and Batch normalization implementation details.