Function train

train.py:178–287

(cfg: dict)

# ──────────────────────────────────────────────────────────────────────────────

def train(cfg: dict) -> None:
    torch.manual_seed(cfg["seed"])
    device = get_device()
    print(f"Using device: {device}")

    # ── tokenizer ──────────────────────────────────────────────────────────
    print("Loading tokenizer …")
    tokenizer = AutoTokenizer.from_pretrained(
        cfg["tokenizer_name"],
        use_fast=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # ── data ───────────────────────────────────────────────────────────────
    dataloader = build_dataloader(tokenizer, cfg)

    # ── model ──────────────────────────────────────────────────────────────
    model = build_model(tokenizer, cfg)
    model.to(device)
    print(f"Parameters: {count_parameters(model):,}")

    # ── optimiser + schedule ───────────────────────────────────────────────
    # Do not apply weight decay to LayerNorm / bias parameters
    decay_params = [p for n, p in model.named_parameters() if p.ndim >= 2]
    nodecay_params = [p for n, p in model.named_parameters() if p.ndim < 2]
    optim_groups = [
        {"params": decay_params, "weight_decay": cfg["weight_decay"]},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(optim_groups, lr=cfg["learning_rate"])
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=cfg["warmup_steps"],
        num_training_steps=cfg["max_steps"],
    )

    # ── W&B (optional) ─────────────────────────────────────────────────────
    if cfg["use_wandb"]:
        import wandb
        wandb.init(project=cfg["wandb_project"], name=cfg["run_name"], config=cfg)

    # ── checkpoint dir ─────────────────────────────────────────────────────
    Path(cfg["checkpoint_dir"]).mkdir(parents=True, exist_ok=True)

    # ── training ───────────────────────────────────────────────────────────
    model.train()
    step = 0
    optimizer.zero_grad()
    autocast_ctx = get_autocast_ctx(device)
    data_iter = iter(dataloader)

    print("Starting training …")
    while step < cfg["max_steps"]:
        # Gradient accumulation
        accumulated_loss = 0.0
        for _ in range(cfg["gradient_accumulation_steps"]):
            try: