LR Scheduler

mediumPROFine-Tuning MechanicsLLM 2: Training & Alignment

Implement a warmup + cosine decay learning rate scheduler.

The schedule has two phases:

Warmup (steps 0 to warmup_steps-1): Linear warmup from 0 to max_lr $lr = \text{max\_lr} \times \frac{\text{step}}{\text{warmup\_steps}}$
Cosine decay (steps warmup_steps to total_steps-1): Decay from max_lr to min_lr $lr = \text{min\_lr} + 0.5 \times (\text{max\_lr} - \text{min\_lr}) \times (1 + \cos(\pi \times \frac{\text{step} - \text{warmup\_steps}}{\text{total\_steps} - \text{warmup\_steps}}))$

Input: warmup_steps total_steps max_lr min_lr Output: Learning rate for each step (0 to total_steps-1), one per line, rounded to 6 decimal places.

Example

Input

2 6 0.001 0.0001

Output

Reasoning

The warmup phase starts from step 0 to 1 (since warmup_steps = 2), with a linear increase from 0 to max_lr = 0.001. At step 0, the learning rate is $0.001 \times \frac{0}{2} = 0.000000$ , and at step 1, it's $0.001 \times \frac{1}{2} = 0.000500$ .
From step 2 onwards, the cosine decay phase begins, with the learning rate calculated as $0.0001 + 0.5 \times (0.001 - 0.0001) \times (1 + \cos(\pi \times \frac{\text{step} - 2}{6 - 2}))$ .
Applying this formula for steps 2 to 5 gives the remaining learning rates: at step 2, $lr = 0.001000$ (since $\cos(0) = 1$ ), at step 3, $lr = 0.000775$ , at step 4, $lr = 0.000325$ , and at step 5, $lr = 0.000100$ .
These calculated learning rates are then rounded to 6 decimal places and output, one per line, resulting in the given sample output.

Constraints

0 < warmup_steps < total_steps
min_lr < max_lr
Round to 6 decimal places

Learn this in

Fine-tuningLLM Study Plan Text ClassificationNLP Study Plan

More from LLM 2: Training & Alignment

RLHF Reward Shapinghard Implicit Reward from DPOmedium PPO Clipped Objectivemedium DPO Lossmedium KL Divergence Penaltyeasy

Concept primerPREVIEW

Background Knowledge

The problem involves implementing a learning rate scheduler, which is a crucial component in machine learning and deep learning. A learning rate scheduler is used to adjust the learning rate during the training process of a model. The learning rate determines how quickly the model learns from the training data. A high learning rate can lead to fast convergence but may also cause the model to overshoot the optimal solution, while a low learning rate can lead to slow convergence.

The problem specifies a warmup and cosine decay schedule. The warmup phase is used to gradually increase the learning rate from 0 to a maximum value, allowing the model to start learning from the data without overshooting. The cosine decay phase is used to decay the learning rate from the maximum value to a minimum value, allowing the model to fine-tune its parameters and converge to the optimal solution. The cosine decay schedule is a type of annealing schedule, which is commonly used in optimization problems to gradually reduce the learning rate.

The mathematical formulas provided in the problem description are used to calculate the learning rate at each step. The warmup phase uses a linear interpolation formula to increase the learning rate from 0 to the maximum value, while the cosine decay phase uses a cosine function to decay the learning rate from the maximum value to the minimum value. Understanding these formulas and how they are used to calculate the learning rate is essential to solving the problem.

Algorithm/Approach

The general approach to solving this problem involves implementing a loop that iterates over each step from 0 to the total number of steps. At each step, the algorithm checks whether it is in the warmup phase or the cosine decay phase and calculates the learning rate accordingly. The algorithm uses the provided formulas to calculate the learning rate at each step and rounds the result to 6 decimal places.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results