| 470 | logging.info(f"Resumed training from step {global_step}") |
| 471 | |
| 472 | def lr_schedule(step: int): |
| 473 | if step < warmup_steps: |
| 474 | # Match JAX behavior: start from peak_lr / (warmup_steps + 1) |
| 475 | init_lr = peak_lr / (warmup_steps + 1) |
| 476 | return init_lr + (peak_lr - init_lr) * step / warmup_steps |
| 477 | # cosine decay |
| 478 | progress = min(1.0, (step - warmup_steps) / max(1, decay_steps - warmup_steps)) |
| 479 | cos = 0.5 * (1 + np.cos(np.pi * progress)) |
| 480 | return end_lr + (peak_lr - end_lr) * cos |
| 481 | |
| 482 | model.train() |
| 483 | start_time = time.time() |