amachine.am_transformers.am_trainer_utils

View Source

 1import math
 2
 3def cosine_lr_with_warmup(
 4    step: int,
 5    warmup_steps: int,
 6    total_steps: int,
 7    peak_lr: float,
 8    min_lr: float,
 9) -> float:
10
11    if step < warmup_steps:
12        return peak_lr * (step + 1) / warmup_steps
13
14    if step >= total_steps:
15        return min_lr
16
17    progress = (step - warmup_steps) / (total_steps - warmup_steps)
18    
19    return min_lr + (peak_lr - min_lr) * 0.5 * (1.0 + math.cos(math.pi * progress))
20
21def wsd_lr_with_warmup(
22    step: int,
23    warmup_steps: int,
24    total_steps: int,
25    decay_steps: int | None,
26    decay_ratio: float,
27    peak_lr: float,
28    min_lr: float,
29) -> float:
30    """
31    Warmup -> Stable -> Decay.
32
33    - [0, warmup_steps):                 linear warmup to peak_lr
34    - [warmup_steps, total_steps-decay_steps): held at peak_lr
35    - [total_steps-decay_steps, total_steps):  cosine decay to min_lr
36
37    If `decay_steps` is None, it's derived as `decay_ratio * (total_steps -
38    warmup_steps)` every call — so if you extend `total_steps` (config.steps)
39    across resumes without setting an explicit decay_steps, the decay window
40    keeps rescaling to stay the last `decay_ratio` fraction of the (possibly
41    still-growing) run, rather than going stale. Set an explicit decay_steps
42    once you know exactly how many steps you want the decay to last.
43    """
44    if step < warmup_steps:
45        return peak_lr * (step + 1) / warmup_steps
46
47    effective_decay_steps = (
48        decay_steps if decay_steps is not None
49        else max(1, round(decay_ratio * (total_steps - warmup_steps)))
50    )
51
52    decay_start = total_steps - effective_decay_steps
53    if step < decay_start:
54        return peak_lr
55
56    progress = (step - decay_start) / effective_decay_steps
57    progress = min(max(progress, 0.0), 1.0)
58    return min_lr + 0.5 * (peak_lr - min_lr) * (1 + math.cos(math.pi * progress))
59
60
61def constant_lr_with_warmup(step: int, warmup_steps: int, peak_lr: float) -> float:
62    """Warmup then flat. Useful as a baseline/ablation against cosine and WSD."""
63    if step < warmup_steps:
64        return peak_lr * (step + 1) / warmup_steps
65    return peak_lr
66
67
68def get_lr(step: int, config, min_lr: float) -> float:
69    """Single dispatch point the Trainer calls. Add new schedules here."""
70    if config.lr_schedule == "cosine":
71        return cosine_lr_with_warmup(
72            step, config.warmup_steps, config.steps, config.lr, min_lr
73        )
74    elif config.lr_schedule == "wsd":
75        return wsd_lr_with_warmup(
76            step, config.warmup_steps, config.steps,
77            config.decay_steps, config.decay_ratio,
78            config.lr, min_lr,
79        )
80    elif config.lr_schedule == "constant":
81        return constant_lr_with_warmup(step, config.warmup_steps, config.lr)
82    else:
83        raise ValueError(f"Unknown lr_schedule: {config.lr_schedule}")

def cosine_lr_with_warmup( step: int, warmup_steps: int, total_steps: int, peak_lr: float, min_lr: float) -> float: View Source

 4def cosine_lr_with_warmup(
 5    step: int,
 6    warmup_steps: int,
 7    total_steps: int,
 8    peak_lr: float,
 9    min_lr: float,
10) -> float:
11
12    if step < warmup_steps:
13        return peak_lr * (step + 1) / warmup_steps
14
15    if step >= total_steps:
16        return min_lr
17
18    progress = (step - warmup_steps) / (total_steps - warmup_steps)
19    
20    return min_lr + (peak_lr - min_lr) * 0.5 * (1.0 + math.cos(math.pi * progress))

def wsd_lr_with_warmup( step: int, warmup_steps: int, total_steps: int, decay_steps: int | None, decay_ratio: float, peak_lr: float, min_lr: float) -> float: View Source

22def wsd_lr_with_warmup(
23    step: int,
24    warmup_steps: int,
25    total_steps: int,
26    decay_steps: int | None,
27    decay_ratio: float,
28    peak_lr: float,
29    min_lr: float,
30) -> float:
31    """
32    Warmup -> Stable -> Decay.
33
34    - [0, warmup_steps):                 linear warmup to peak_lr
35    - [warmup_steps, total_steps-decay_steps): held at peak_lr
36    - [total_steps-decay_steps, total_steps):  cosine decay to min_lr
37
38    If `decay_steps` is None, it's derived as `decay_ratio * (total_steps -
39    warmup_steps)` every call — so if you extend `total_steps` (config.steps)
40    across resumes without setting an explicit decay_steps, the decay window
41    keeps rescaling to stay the last `decay_ratio` fraction of the (possibly
42    still-growing) run, rather than going stale. Set an explicit decay_steps
43    once you know exactly how many steps you want the decay to last.
44    """
45    if step < warmup_steps:
46        return peak_lr * (step + 1) / warmup_steps
47
48    effective_decay_steps = (
49        decay_steps if decay_steps is not None
50        else max(1, round(decay_ratio * (total_steps - warmup_steps)))
51    )
52
53    decay_start = total_steps - effective_decay_steps
54    if step < decay_start:
55        return peak_lr
56
57    progress = (step - decay_start) / effective_decay_steps
58    progress = min(max(progress, 0.0), 1.0)
59    return min_lr + 0.5 * (peak_lr - min_lr) * (1 + math.cos(math.pi * progress))

Warmup -> Stable -> Decay.

[0, warmup_steps): linear warmup to peak_lr
[warmup_steps, total_steps-decay_steps): held at peak_lr
[total_steps-decay_steps, total_steps): cosine decay to min_lr

If decay_steps is None, it's derived as decay_ratio * (total_steps - warmup_steps) every call — so if you extend total_steps (config.steps) across resumes without setting an explicit decay_steps, the decay window keeps rescaling to stay the last decay_ratio fraction of the (possibly still-growing) run, rather than going stale. Set an explicit decay_steps once you know exactly how many steps you want the decay to last.

def constant_lr_with_warmup(step: int, warmup_steps: int, peak_lr: float) -> float: View Source

62def constant_lr_with_warmup(step: int, warmup_steps: int, peak_lr: float) -> float:
63    """Warmup then flat. Useful as a baseline/ablation against cosine and WSD."""
64    if step < warmup_steps:
65        return peak_lr * (step + 1) / warmup_steps
66    return peak_lr

Warmup then flat. Useful as a baseline/ablation against cosine and WSD.

def get_lr(step: int, config, min_lr: float) -> float: View Source

69def get_lr(step: int, config, min_lr: float) -> float:
70    """Single dispatch point the Trainer calls. Add new schedules here."""
71    if config.lr_schedule == "cosine":
72        return cosine_lr_with_warmup(
73            step, config.warmup_steps, config.steps, config.lr, min_lr
74        )
75    elif config.lr_schedule == "wsd":
76        return wsd_lr_with_warmup(
77            step, config.warmup_steps, config.steps,
78            config.decay_steps, config.decay_ratio,
79            config.lr, min_lr,
80        )
81    elif config.lr_schedule == "constant":
82        return constant_lr_with_warmup(step, config.warmup_steps, config.lr)
83    else:
84        raise ValueError(f"Unknown lr_schedule: {config.lr_schedule}")

Single dispatch point the Trainer calls. Add new schedules here.