tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-07-10 14:59:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 56660c690b into e569b59f92

Browse Source
This commit is contained in:
Dipesh Babu 2026-02-11 03:18:21 +08:00 committed by GitHub
commit c04a9a4e24
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 21 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
scripts/base_train.py
View File

@ -346,15 +346,28 @@ print0(f"Total training FLOPs estimate: {num_flops_per_token * total_tokens:e}")
# Learning rate schedule (linear warmup, constant, linear warmdown) # Learning rate schedule (linear warmup, constant, linear warmdown)
def get_lr_multiplier(it): def get_lr_multiplier(it):
warmup_iters = round(args.warmup_ratio * num_iterations) # Note: optimizer steps run for it in [0, num_iterations-1]
warmdown_iters = round(args.warmdown_ratio * num_iterations) warmup_iters = int(round(args.warmup_ratio * num_iterations))
if it < warmup_iters: warmdown_iters = int(round(args.warmdown_ratio * num_iterations))
# Warmup (avoid division by zero when warmup_iters == 0)
if warmup_iters > 0 and it < warmup_iters:
return (it + 1) / warmup_iters return (it + 1) / warmup_iters
elif it <= num_iterations - warmdown_iters:
return 1.0 # Warmdown should cover the last `warmdown_iters` optimizer steps:
else: # it in [num_iterations - warmdown_iters, num_iterations - 1]
progress = (num_iterations - it) / warmdown_iters if warmdown_iters > 0:
return progress * 1.0 + (1 - progress) * args.final_lr_frac warmdown_start = num_iterations - warmdown_iters
# Ensure warmdown doesn't start before warmup ends (prevents overlap weirdness)
warmdown_start = max(warmdown_start, warmup_iters)
if it >= warmdown_start:
# progress: 1.0 at warmdown_start, 0.0 at last optimizer step (num_iterations - 1)
span = max(1, (num_iterations - 1) - warmdown_start) # denom >= 1
progress = (num_iterations - 1 - it) / span
return progress * 1.0 + (1.0 - progress) * args.final_lr_frac
return 1.0
# Momentum scheduler for Muon optimizer (warms up to 0.95 over the first 300 steps) # Momentum scheduler for Muon optimizer (warms up to 0.95 over the first 300 steps)
def get_muon_momentum(it): def get_muon_momentum(it):