78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
"""Minimal training-loop helpers for course stage assembly.
|
|
|
|
The functions in this file are intentionally small. Lesson 14 asks the
|
|
learner to implement the training mechanics before wiring a full pretrain
|
|
script together.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Iterable
|
|
|
|
import torch
|
|
|
|
|
|
def cosine_lr(current_step: int, total_steps: int, base_lr: float) -> float:
|
|
"""TODO: implement MiniMind's cosine learning-rate schedule.
|
|
|
|
Align with: trainer/trainer_utils.py::get_lr
|
|
"""
|
|
return base_lr*(0.1+0.45*(1+math.cos(math.pi * current_step / total_steps)))
|
|
raise NotImplementedError("Implement in lesson 14.")
|
|
|
|
|
|
def scale_loss_for_accumulation(loss: torch.Tensor, accumulation_steps: int) -> torch.Tensor:
|
|
"""TODO: divide a micro-batch loss for gradient accumulation.
|
|
|
|
Align with: trainer/train_pretrain.py:35-40
|
|
"""
|
|
return loss/accumulation_steps
|
|
raise NotImplementedError("Implement in lesson 14.")
|
|
|
|
|
|
def should_step_optimizer(step: int, accumulation_steps: int) -> bool:
|
|
"""TODO: return whether this micro-step should call optimizer.step.
|
|
|
|
Align with: trainer/train_pretrain.py:42
|
|
"""
|
|
return step%accumulation_steps==0
|
|
raise NotImplementedError("Implement in lesson 14.")
|
|
|
|
|
|
def should_flush_tail(last_step: int, start_step: int, accumulation_steps: int) -> bool:
|
|
"""TODO: return whether a partial accumulation group remains at epoch end.
|
|
|
|
Align with: trainer/train_pretrain.py:75
|
|
"""
|
|
return last_step > start_step and last_step % args.accumulation_steps != 0:
|
|
raise NotImplementedError("Implement in lesson 14.")
|
|
|
|
|
|
def train_one_epoch(
|
|
model: torch.nn.Module,
|
|
loader: Iterable[Any],
|
|
optimizer: torch.optim.Optimizer,
|
|
*,
|
|
device: str | torch.device,
|
|
epochs: int,
|
|
epoch: int,
|
|
iters: int,
|
|
base_lr: float,
|
|
accumulation_steps: int,
|
|
grad_clip: float,
|
|
scaler: torch.cuda.amp.GradScaler | None = None,
|
|
autocast_ctx: Any | None = None,
|
|
start_step: int = 0,
|
|
) -> list[dict[str, float]]:
|
|
"""TODO: implement forward/loss/backward/optimizer.step loop.
|
|
|
|
Align with train_epoch functions in trainer/train_*.py, but omit DDP,
|
|
wandb, and full checkpoint resume in the first teaching version.
|
|
"""
|
|
raise NotImplementedError("Implement in lesson 14.")
|
|
|
|
|
|
def save_course_checkpoint(*args, **kwargs):
|
|
"""TODO: save teaching-version weights for the current stage."""
|
|
raise NotImplementedError("Implement during the first stage assembly.")
|