Learning Rate Warmup for Stable LLM Training

This lesson begins with a small program that prints one number. The program first computes how many training steps a run will take, then computes how many of those steps should be used for warmup, and finally prints the warmup count. Warmup is the first phase of training, in which the learning rate is increased from a small starting value. So the program is answering a practical question: how many of the earliest training steps should be spent on that increase?

Keep that question in mind as you read. The printed value only makes sense once you know what it counts. It counts steps, not a percentage, even though the percentage is what the program uses to decide the count.

Every warmup calculation starts from the total number of training steps. In this program that total is the number of batches in the training loader multiplied by the number of epochs. An epoch is one full pass over the training data, and the loader splits that data into batches. If the loader yields a fixed number of batches per pass, then running several passes multiplies that count. The product is the total number of steps the run will take.

Written as a relationship, the total is the batch count times the epoch count, where B stands for batches per epoch and E for the number of epochs. The first line of the program is the foundation for everything after it. If the total is wrong, every value derived from it will be wrong too.

The source states that warmup steps are usually set between 0.1% and 20% of the total number of steps. The program uses the upper end of that range, 20%, and writes it as the decimal 0.2. Multiplying the total by 0.2 gives a fraction of the total. Steps are countable, so the fractional part is discarded by int(), which converts the result into a whole number of steps.

The count is not an end in itself. The source explains that the warmup steps are used to increase the initial learning rate. In the example it describes, the learning rate rises from 0.0001 to 0.01 across the first 27 training steps. The warmup count therefore sets how long that climb takes. A larger count spreads the increase over more steps, and a smaller count compresses it into fewer steps.

The lower end of the range can produce a surprise. Take a total of 500 steps and a warmup share of 0.1 percent. The direct calculation is 0.001 times 500, which equals 0.5. The integer conversion then discards the fraction and produces 0. A learner who expects at least one warmup step from a small percentage will be surprised, but the cause is the truncation, not an error in the multiplication.

When a printed warmup count looks wrong, do not start by questioning the percentage. Walk the calculation in order and check each intermediate value. The trace has three stages: the number of batches, the total after multiplying by epochs, and the warmup count after applying 0.2 and truncating. The first stage that disagrees with your expected value is the location of the problem.

The trace also shows where to look first. The printed 27 corresponds to a total near 135, so a mismatch usually traces back to the total rather than to the final line. This is why the most useful debugging habit with this program is to check the total before looking at the percentage logic.

A percentage describes a share of the run. A warmup count describes a number of steps. For a total of 135 steps, 20% is a share of the run, while the warmup count it produces is 27 steps. The 20 in the lesson describes the share; the 27 is the count the program prints. Mixing the two up leads to misreadings of the output.