dela
584f7e9e73
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
"""LR scale and token-horizon helpers for train_k3 / train_sft."""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
|
|
def lr_scale(
|
|
opt_step: int,
|
|
warmup: int,
|
|
total_opt: int,
|
|
min_ratio: float = 0.1,
|
|
) -> float:
|
|
"""Linear warmup (optimizer steps) then cosine down to ``min_ratio``.
|
|
|
|
``opt_step`` is 0-indexed at the optimizer update that is about to run.
|
|
"""
|
|
if warmup > 0 and opt_step < warmup:
|
|
return (opt_step + 1) / warmup
|
|
denom = max(total_opt - warmup - 1, 1)
|
|
progress = min(max(opt_step - warmup, 0) / denom, 1.0)
|
|
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
|
|
return min_ratio + (1.0 - min_ratio) * cosine
|
|
|
|
|
|
def tokens_per_micro(batch: int, seq_len: int) -> int:
|
|
return batch * seq_len
|
|
|
|
|
|
def total_opt_steps(
|
|
*,
|
|
max_tokens: int | None,
|
|
max_micro: int | None,
|
|
batch: int,
|
|
seq_len: int,
|
|
grad_acc: int,
|
|
) -> int:
|
|
"""Optimizer-step horizon used by cosine. At least 1."""
|
|
acc = max(grad_acc, 1)
|
|
candidates: list[int] = []
|
|
if max_tokens is not None and max_tokens > 0:
|
|
tpm = max(tokens_per_micro(batch, seq_len), 1)
|
|
candidates.append(math.ceil(max_tokens / (tpm * acc)))
|
|
if max_micro is not None and max_micro > 0:
|
|
candidates.append(math.ceil(max_micro / acc))
|
|
if not candidates:
|
|
return 1
|
|
return max(min(candidates), 1)
|