src/trainer.py

from __future__ import annotations

import time
from dataclasses import replace
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

from .config import Config
from .data import TextDataset
from .diffusion import MultiMaskDiffusion
from .tokenizer import CharacterTokenizer
from .utils import build_model, resolve_device, save_checkpoint, save_json, set_seed


def compute_loss(
    model: torch.nn.Module,
    diffusion: MultiMaskDiffusion,
    x0: torch.Tensor,
    attention_mask: torch.Tensor,
    t: torch.Tensor,
    clean_token_start: int,
    aux_weight: float,
) -> tuple[torch.Tensor, dict[str, float], torch.Tensor]:
    xt = diffusion.q_sample(x0, t, attention_mask)
    logits, mask_logits = model(xt, t, attention_mask)
    masked = diffusion.is_mask(xt) & attention_mask
    clean = x0.ge(clean_token_start) & attention_mask
    target_local = (x0 - clean_token_start).clamp_min(0)

    if masked.any():
        reconstruction = F.cross_entropy(logits[masked], target_local[masked])
        accuracy = (
            logits[masked].argmax(dim=-1).eq(target_local[masked]).float().mean()
        )
    else:
        reconstruction = logits.sum() * 0.0
        accuracy = torch.zeros((), device=x0.device)

    if clean.any() and aux_weight > 0:
        mask_target = target_local.remainder(diffusion.num_masks)
        auxiliary = F.cross_entropy(mask_logits[clean], mask_target[clean])
    else:
        auxiliary = mask_logits.sum() * 0.0
    loss = reconstruction + aux_weight * auxiliary
    metrics = {
        "loss": float(loss.detach()),
        "reconstruction_loss": float(reconstruction.detach()),
        "auxiliary_loss": float(auxiliary.detach()),
        "masked_accuracy": float(accuracy.detach()),
        "masked_tokens": int(masked.sum().detach()),
    }
    return loss, metrics, xt


def _train_once(
    config: Config,
    tokenizer: CharacterTokenizer,
    train_texts: list[str],
    output_dir: Path,
) -> dict:
    set_seed(config.seed)
    device = resolve_device(config.device)
    dataset = TextDataset(train_texts, tokenizer, config.seq_len)
    loader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=True,
        drop_last=False,
        pin_memory=device.type == "cuda",
    )
    model = build_model(config, tokenizer).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        weight_decay=config.weight_decay,
    )
    diffusion = MultiMaskDiffusion(
        tokenizer.vocab_size,
        tokenizer.clean_token_start,
        tokenizer.num_masks,
        tokenizer.mask_token_start,
        tokenizer.pad_id,
    )
    history: list[dict] = []
    global_step = 0
    started = time.perf_counter()
    progress = tqdm(total=config.max_steps, desc=f"training on {device}")
    model.train()

    while global_step < config.max_steps:
        for batch in loader:
            x0 = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            # Avoid near-zero times so every batch usually contains masked positions.
            t = 0.05 + 0.95 * torch.rand(x0.shape[0], device=device)
            optimizer.zero_grad(set_to_none=True)
            loss, metrics, _ = compute_loss(
                model,
                diffusion,
                x0,
                attention_mask,
                t,
                tokenizer.clean_token_start,
                config.aux_weight,
            )
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
            optimizer.step()
            global_step += 1
            metrics["step"] = global_step
            history.append(metrics)
            progress.update(1)
            progress.set_postfix(loss=f"{metrics['loss']:.4f}")

            if global_step % config.log_every == 0:
                print(
                    f"step={global_step} loss={metrics['loss']:.4f} "
                    f"masked_acc={metrics['masked_accuracy']:.3f}"
                )
            if global_step % config.save_every == 0:
                save_checkpoint(
                    output_dir / f"model_step_{global_step}.pt",
                    model,
                    optimizer,
                    config,
                    tokenizer,
                    global_step,
                    metrics,
                )
            if global_step >= config.max_steps:
                break
    progress.close()
    elapsed = time.perf_counter() - started
    final_metrics = {
        **history[-1],
        "initial_loss": history[0]["loss"],
        "final_loss": history[-1]["loss"],
        "training_time_seconds": elapsed,
        "device": str(device),
        "parameters": sum(parameter.numel() for parameter in model.parameters()),
    }
    save_checkpoint(
        output_dir / "model.pt",
        model,
        optimizer,
        config,
        tokenizer,
        global_step,
        final_metrics,
    )
    config.save(output_dir / "config.json")
    tokenizer.save(output_dir / "vocab.json")
    save_json(history, output_dir / "training_log.json")
    return final_metrics


def train_with_fallbacks(
    config: Config,
    tokenizer: CharacterTokenizer,
    train_texts: list[str],
    output_dir: str | Path,
) -> tuple[Config, dict]:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    attempts = [config]
    if config.device in {"auto", "cuda"}:
        attempts.extend(
            [
                replace(config, batch_size=max(1, config.batch_size // 2)),
                replace(
                    config,
                    batch_size=max(1, config.batch_size // 2),
                    d_model=64,
                    nhead=min(config.nhead, 4),
                    dim_feedforward=128,
                ),
                replace(
                    config,
                    device="cpu",
                    batch_size=min(8, config.batch_size),
                    seq_len=min(32, config.seq_len),
                    d_model=64,
                    nhead=4,
                    dim_feedforward=128,
                    max_steps=min(300, config.max_steps),
                ),
            ]
        )
    last_error: Exception | None = None
    for attempt_index, attempt in enumerate(attempts, start=1):
        try:
            print(
                f"Attempt {attempt_index}/{len(attempts)}: "
                f"device={attempt.device}, batch={attempt.batch_size}, "
                f"seq_len={attempt.seq_len}, d_model={attempt.d_model}"
            )
            return attempt, _train_once(attempt, tokenizer, train_texts, output_dir)
        except torch.cuda.OutOfMemoryError as error:
            last_error = error
            print("CUDA OOM; applying the next automatic fallback.")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    if last_error is not None:
        raise last_error
    raise RuntimeError("No training attempt was executed")