LLMs-from-scratch/ch05/04_learning_rate_schedulers
..
README.md

Adding Bells and Whistles to the Training Loop

The main chapter used a relatively simple training function to keep the code readable and fit Chapter 5 within the page limits. Optionally, we can add a linear warm-up, a cosine decay schedule, and gradient clipping to improve the training stability and convergence.

You can find the code for this more sophisticated training function in Appendix D: Adding Bells and Whistles to the Training Loop.