trainer.py

from __future__ import annotations

import json
import math
import time
from contextlib import nullcontext
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP

from .checkpointing import load_checkpoint, save_checkpoint
from .data import DataLoader
from .optim import CosineAnnealingLR, clip_gradients, create_optimizer, global_grad_norm


@dataclass(frozen=True)
class TrainingStage:
    start_iteration: int
    batch_size: int | None = None
    context_length: int | None = None
    attention_window_size: int | None = None
    override_attention_window_size: bool = False
    mtp_weights: tuple[float, ...] = ()

    def __post_init__(self) -> None:
        if self.start_iteration < 0:
            raise ValueError(f"start_iteration must be non-negative, got {self.start_iteration}.")
        if self.batch_size is not None and self.batch_size <= 0:
            raise ValueError(f"batch_size must be positive when provided, got {self.batch_size}.")
        if self.context_length is not None and self.context_length <= 0:
            raise ValueError(f"context_length must be positive when provided, got {self.context_length}.")
        if self.attention_window_size is not None and self.attention_window_size <= 0:
            raise ValueError(
                "attention_window_size must be positive when provided, "
                f"got {self.attention_window_size}."
            )
        if self.attention_window_size is not None and not self.override_attention_window_size:
            object.__setattr__(self, "override_attention_window_size", True)
        _normalize_mtp_weights(self.mtp_weights, field_name="mtp_weights", allow_empty=True)


@dataclass(frozen=True)
class RuntimeTrainingState:
    batch_size: int
    context_length: int
    attention_window_size: int | None
    mtp_weights: tuple[float, ...]


@dataclass(frozen=True)
class TrainingSchedule:
    batch_size: int
    context_length: int
    attention_window_size: int | None
    mtp_weights: tuple[float, ...]
    stages: tuple[TrainingStage, ...] = ()

    def resolve(self, completed_iterations: int) -> RuntimeTrainingState:
        batch_size = self.batch_size
        context_length = self.context_length
        attention_window_size = self.attention_window_size
        mtp_weights = self.mtp_weights
        for stage in self.stages:
            if completed_iterations < stage.start_iteration:
                break
            if stage.batch_size is not None:
                batch_size = stage.batch_size
            if stage.context_length is not None:
                context_length = stage.context_length
            if stage.override_attention_window_size:
                attention_window_size = stage.attention_window_size
            if stage.mtp_weights:
                mtp_weights = stage.mtp_weights
        return RuntimeTrainingState(
            batch_size=batch_size,
            context_length=context_length,
            attention_window_size=attention_window_size,
            mtp_weights=mtp_weights,
        )

    def max_context_length(self) -> int:
        max_context_length = self.context_length
        for stage in self.stages:
            if stage.context_length is not None:
                max_context_length = max(max_context_length, stage.context_length)
        return max_context_length

    def max_mtp_tokens(self) -> int:
        max_tokens = len(self.mtp_weights)
        for stage in self.stages:
            if stage.mtp_weights:
                max_tokens = max(max_tokens, len(stage.mtp_weights))
        return max_tokens


@dataclass
class TrainerConfig:
    batch_size: int
    context_length: int
    max_iters: int
    learning_rate: float = 3e-4
    min_learning_rate: float = 3e-5
    warmup_iters: int = 100
    lr_decay_iters: int | None = None
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    eps: float = 1e-8
    grad_clip: float | None = 1.0
    eval_interval: int = 100
    eval_iters: int = 10
    log_interval: int = 10
    checkpoint_interval: int = 500
    device: str = "cpu"
    dtype: str = "float32"
    seed: int = 1337
    output_dir: str | Path = "out"
    write_config: bool = True
    checkpoint_name: str = "latest.pt"
    best_checkpoint_name: str = "best.pt"
    gradient_accumulation_steps: int = 1
    compile_model: bool = False
    pin_memory: bool | None = None
    use_tf32: bool = True
    use_fused_optimizer: bool = True
    context_length_schedule: tuple[tuple[int, int], ...] = ()
    batch_size_schedule: tuple[tuple[int, int], ...] = ()
    mtp_weights: tuple[float, ...] = (1.0,)
    training_stages: tuple[TrainingStage, ...] = ()
    scale_lr_with_batch_size: bool = True
    loss_spike_threshold: float | None = 2.0
    grad_norm_spike_threshold: float | None = 2.0
    optimizer_name: str = "adamw"
    schedule_weight_decay_with_lr: bool = True
    cautious_weight_decay: bool = False
    document_boundary_sampling_probability: float = 0.5
    muon_momentum: float = 0.95
    muon_nesterov: bool = True
    muon_ns_steps: int = 5
    muon_eps: float = 1e-8

    def __post_init__(self) -> None:
        if self.batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {self.batch_size}.")
        if self.context_length <= 0:
            raise ValueError(f"context_length must be positive, got {self.context_length}.")
        if self.max_iters <= 0:
            raise ValueError(f"max_iters must be positive, got {self.max_iters}.")
        if self.learning_rate <= 0:
            raise ValueError(f"learning_rate must be positive, got {self.learning_rate}.")
        if self.min_learning_rate < 0:
            raise ValueError(f"min_learning_rate must be non-negative, got {self.min_learning_rate}.")
        if self.warmup_iters < 0:
            raise ValueError(f"warmup_iters must be non-negative, got {self.warmup_iters}.")
        if self.lr_decay_iters is not None and self.lr_decay_iters <= 0:
            raise ValueError(f"lr_decay_iters must be positive when provided, got {self.lr_decay_iters}.")
        if self.eval_interval <= 0:
            raise ValueError(f"eval_interval must be positive, got {self.eval_interval}.")
        if self.eval_iters <= 0:
            raise ValueError(f"eval_iters must be positive, got {self.eval_iters}.")
        if self.log_interval <= 0:
            raise ValueError(f"log_interval must be positive, got {self.log_interval}.")
        if self.checkpoint_interval <= 0:
            raise ValueError(f"checkpoint_interval must be positive, got {self.checkpoint_interval}.")
        if self.gradient_accumulation_steps <= 0:
            raise ValueError(
                "gradient_accumulation_steps must be positive, "
                f"got {self.gradient_accumulation_steps}."
            )
        if self.dtype not in {"float32", "float16", "bfloat16"}:
            raise ValueError(f"dtype must be one of 'float32', 'float16', or 'bfloat16', got {self.dtype!r}.")
        self.context_length_schedule = self._normalize_schedule(self.context_length_schedule, "context_length_schedule")
        self.batch_size_schedule = self._normalize_schedule(self.batch_size_schedule, "batch_size_schedule")
        self.mtp_weights = _normalize_mtp_weights(self.mtp_weights, field_name="mtp_weights")
        self.training_stages = self._normalize_training_stages(self.training_stages)
        if self.loss_spike_threshold is not None and self.loss_spike_threshold <= 1.0:
            raise ValueError(
                f"loss_spike_threshold must be greater than 1.0 when provided, got {self.loss_spike_threshold}."
            )
        if self.grad_norm_spike_threshold is not None and self.grad_norm_spike_threshold <= 1.0:
            raise ValueError(
                f"grad_norm_spike_threshold must be greater than 1.0 when provided, got {self.grad_norm_spike_threshold}."
            )
        if self.optimizer_name not in {"adamw", "muon"}:
            raise ValueError(f"optimizer_name must be 'adamw' or 'muon', got {self.optimizer_name!r}.")
        if not 0.0 <= self.document_boundary_sampling_probability <= 1.0:
            raise ValueError(
                "document_boundary_sampling_probability must be in [0, 1], "
                f"got {self.document_boundary_sampling_probability}."
            )
        if not 0 <= self.muon_momentum < 1:
            raise ValueError(f"muon_momentum must be in [0, 1), got {self.muon_momentum}.")
        if self.muon_ns_steps <= 0:
            raise ValueError(f"muon_ns_steps must be positive, got {self.muon_ns_steps}.")
        if self.muon_eps <= 0:
            raise ValueError(f"muon_eps must be positive, got {self.muon_eps}.")

    @property
    def torch_dtype(self) -> torch.dtype:
        mapping = {
            "float32": torch.float32,
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
        }
        return mapping[self.dtype]

    @staticmethod
    def _normalize_schedule(
        schedule: tuple[tuple[int, int], ...] | list[tuple[int, int]],
        field_name: str,
    ) -> tuple[tuple[int, int], ...]:
        normalized = tuple((int(step), int(value)) for step, value in schedule)
        for step, value in normalized:
            if step < 0:
                raise ValueError(f"{field_name} steps must be non-negative, got {step}.")
            if value <= 0:
                raise ValueError(f"{field_name} values must be positive, got {value}.")
        if tuple(sorted(normalized)) != normalized:
            raise ValueError(f"{field_name} must be sorted by step in ascending order.")
        return normalized

    @staticmethod
    def _normalize_training_stages(
        training_stages: tuple[TrainingStage, ...] | list[TrainingStage | dict[str, Any]],
    ) -> tuple[TrainingStage, ...]:
        normalized: list[TrainingStage] = []
        for stage in training_stages:
            if isinstance(stage, TrainingStage):
                normalized.append(stage)
                continue
            if isinstance(stage, dict):
                stage_kwargs = dict(stage)
                if "mtp_weights" in stage_kwargs and stage_kwargs["mtp_weights"] is not None:
                    stage_kwargs["mtp_weights"] = tuple(float(weight) for weight in stage_kwargs["mtp_weights"])
                normalized.append(TrainingStage(**stage_kwargs))
                continue
            raise TypeError(f"Unsupported training stage value {stage!r}.")
        if tuple(sorted(normalized, key=lambda stage: stage.start_iteration)) != tuple(normalized):
            raise ValueError("training_stages must be sorted by start_iteration in ascending order.")
        return tuple(normalized)


class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_dataset: npt.NDArray[np.integer] | np.memmap,
        valid_dataset: npt.NDArray[np.integer] | np.memmap,
        config: TrainerConfig,
        optimizer: torch.optim.Optimizer | None = None,
        train_document_ids: npt.NDArray[np.integer] | np.memmap | None = None,
        valid_document_ids: npt.NDArray[np.integer] | np.memmap | None = None,
    ):
        self.config = config
        self.device = torch.device(config.device)
        self.device_type = self.device.type
        self.distributed = dist.is_available() and dist.is_initialized()
        self.rank = dist.get_rank() if self.distributed else 0
        self.world_size = dist.get_world_size() if self.distributed else 1
        self.master_process = self.rank == 0

        if config.use_tf32 and self.device_type == "cuda":
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

        self.raw_model = model.to(self.device)
        self.model = self.raw_model
        self.train_dataset = train_dataset
        self.valid_dataset = valid_dataset
        self.iteration = 0
        self.best_val_loss = float("inf")
        self.history: list[dict[str, Any]] = []
        self.loss_ema: float | None = None
        self.grad_norm_ema: float | None = None

        self.training_schedule = _build_training_schedule(
            config,
            default_attention_window_size=getattr(self.raw_model, "attention_window_size", None),
        )
        scheduled_max_mtp_tokens = self.training_schedule.max_mtp_tokens()
        model_max_mtp_tokens = int(getattr(self.raw_model, "max_mtp_tokens", 1))
        if scheduled_max_mtp_tokens > model_max_mtp_tokens:
            raise ValueError(
                "Training schedule target horizon "
                f"{scheduled_max_mtp_tokens} exceeds model.max_mtp_tokens={model_max_mtp_tokens}. "
                "Increase max_mtp_tokens or reduce the configured MTP weights."
            )
        initial_runtime_state = self.training_schedule.resolve(0)
        self.base_batch_size = initial_runtime_state.batch_size
        self.current_batch_size = initial_runtime_state.batch_size
        self.current_context_length = initial_runtime_state.context_length
        self.current_attention_window_size = initial_runtime_state.attention_window_size
        self.current_mtp_weights = initial_runtime_state.mtp_weights
        self.current_target_horizon = len(initial_runtime_state.mtp_weights)
        self.base_effective_batch_size = (
            self.base_batch_size * config.gradient_accumulation_steps * self.world_size
        )
        self.raw_model.set_attention_window_size(self.current_attention_window_size)

        torch.manual_seed(config.seed + self.rank)
        np.random.seed(config.seed + self.rank)
        train_rng = np.random.default_rng(config.seed + self.rank)
        valid_rng = np.random.default_rng(config.seed + self.rank + 10_000)

        self.train_loader = DataLoader(
            dataset=train_dataset,
            document_ids=train_document_ids,
            batch_size=self.current_batch_size,
            context_length=self.current_context_length,
            device=self.device,
            rng=train_rng,
            pin_memory=config.pin_memory,
            document_boundary_sampling_probability=config.document_boundary_sampling_probability,
            target_horizon=self.current_target_horizon,
        )
        self.valid_loader = DataLoader(
            dataset=valid_dataset,
            document_ids=valid_document_ids,
            batch_size=self.current_batch_size,
            context_length=self.current_context_length,
            device=self.device,
            rng=valid_rng,
            pin_memory=config.pin_memory,
            document_boundary_sampling_probability=0.0,
            target_horizon=self.current_target_horizon,
        )

        self.optimizer = optimizer or create_optimizer(
            self.raw_model,
            optimizer_name=config.optimizer_name,
            lr=config.learning_rate,
            betas=(config.beta1, config.beta2),
            eps=config.eps,
            weight_decay=config.weight_decay,
            device_type=self.device_type,
            use_fused=config.use_fused_optimizer,
            muon_momentum=config.muon_momentum,
            muon_nesterov=config.muon_nesterov,
            muon_ns_steps=config.muon_ns_steps,
            muon_eps=config.muon_eps,
            cautious_weight_decay=config.cautious_weight_decay,
        )
        self._ensure_optimizer_group_metadata()
        self.scheduler = CosineAnnealingLR(
            max_learning_rate=config.learning_rate,
            min_learning_rate=config.min_learning_rate,
            warmup_iters=config.warmup_iters,
            cosine_cycle_iters=config.lr_decay_iters or config.max_iters,
        )

        if self.distributed:
            device_ids = [self.device.index] if self.device_type == "cuda" else None
            self.model = DDP(self.raw_model, device_ids=device_ids)

        self.training_model = self.model
        if config.compile_model and hasattr(torch, "compile"):
            self.training_model = torch.compile(self.model)

        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.checkpoint_path = self.output_dir / config.checkpoint_name
        self.best_checkpoint_path = self.output_dir / config.best_checkpoint_name
        self.config_path = self.output_dir / "trainer_config.json"
        self.metrics_path = self.output_dir / "metrics.jsonl"
        self._write_config()

        self.autocast_enabled = (
            self.device_type in {"cuda", "cpu"}
            and config.dtype in {"float16", "bfloat16"}
            and not (self.device_type == "cpu" and config.dtype == "float16")
        )
        self.grad_scaler = self._create_grad_scaler(enabled=self.device_type == "cuda" and config.dtype == "float16")

    def train(self, max_iters: int | None = None) -> list[dict[str, Any]]:
        target_iteration = max_iters if max_iters is not None else self.config.max_iters
        if target_iteration < self.iteration:
            raise ValueError(
                f"target iteration {target_iteration} is behind current iteration {self.iteration}."
            )

        while self.iteration < target_iteration:
            train_metrics = self.train_step()
            self._record_metrics(train_metrics)
            if self.master_process and (self.iteration % self.config.log_interval == 0 or self.iteration == 1):
                print(
                    f"[train] iter={self.iteration} "
                    f"loss={train_metrics['train_loss']:.4f} "
                    f"lr={train_metrics['lr']:.6f} "
                    f"tokens/sec={train_metrics['tokens_per_sec']:.1f}"
                )

            if self.iteration % self.config.eval_interval == 0 or self.iteration == target_iteration:
                eval_metrics = self.evaluate()
                self._record_metrics(eval_metrics)
                if self.master_process:
                    print(
                        f"[eval] iter={self.iteration} "
                        f"train_loss={eval_metrics['train_loss']:.4f} "
                        f"valid_loss={eval_metrics['valid_loss']:.4f}"
                    )
                val_loss = eval_metrics["valid_loss"]
                if val_loss < self.best_val_loss:
                    self.best_val_loss = val_loss
                    self.save_checkpoint(self.best_checkpoint_path)

            if self.iteration % self.config.checkpoint_interval == 0 or self.iteration == target_iteration:
                self.save_checkpoint(self.checkpoint_path)

        return self.history

    def train_step(self) -> dict[str, Any]:
        self._apply_runtime_schedules(self.iteration)
        self.training_model.train()
        lr = self._scheduled_learning_rate(self.iteration + 1)
        self._set_learning_rate(lr)
        self.optimizer.zero_grad(set_to_none=True)

        start_time = time.perf_counter()
        mean_ce_loss = 0.0
        mean_aux_loss = 0.0
        mean_z_loss = 0.0
        mean_total_loss = 0.0

        for micro_step in range(self.config.gradient_accumulation_steps):
            x, y, document_ids, loss_mask = self.train_loader.get_batch(return_document_ids=True)
            sync_context = (
                self.model.no_sync()
                if self.distributed and micro_step < self.config.gradient_accumulation_steps - 1
                else nullcontext()
            )
            with sync_context:
                with self._autocast_context():
                    if self.current_target_horizon == 1:
                        logits = self.training_model(x, document_ids=document_ids)
                        ce_loss = self._compute_next_token_loss(logits, y, loss_mask)
                    else:
                        hidden_states = self.training_model(x, document_ids=document_ids, return_hidden=True)
                        ce_loss = self._compute_mtp_loss(hidden_states, y, loss_mask)
                    moe_metrics = self._current_moe_metrics()
                    aux_loss = moe_metrics["aux_loss"].to(device=ce_loss.device, dtype=ce_loss.dtype)
                    z_loss = moe_metrics["z_loss"].to(device=ce_loss.device, dtype=ce_loss.dtype)
                    total_loss = ce_loss + aux_loss + z_loss
                    scaled_loss = total_loss / self.config.gradient_accumulation_steps

                if self.grad_scaler.is_enabled():
                    self.grad_scaler.scale(scaled_loss).backward()
                else:
                    scaled_loss.backward()

            mean_ce_loss += ce_loss.detach().item()
            mean_aux_loss += aux_loss.detach().item()
            mean_z_loss += z_loss.detach().item()
            mean_total_loss += total_loss.detach().item()

        if self.grad_scaler.is_enabled():
            self.grad_scaler.unscale_(self.optimizer)
        grad_norm_tensor = global_grad_norm(self.raw_model.parameters())
        grad_norm = self._reduce_scalar(float(grad_norm_tensor.item()))
        if self.grad_scaler.is_enabled():
            if self.config.grad_clip is not None:
                clip_gradients(self.raw_model.parameters(), self.config.grad_clip)
            self.grad_scaler.step(self.optimizer)
            self.grad_scaler.update()
        else:
            if self.config.grad_clip is not None:
                clip_gradients(self.raw_model.parameters(), self.config.grad_clip)
            self.optimizer.step()

        self.iteration += 1
        step_time = time.perf_counter() - start_time

        mean_ce_loss /= self.config.gradient_accumulation_steps
        mean_aux_loss /= self.config.gradient_accumulation_steps
        mean_z_loss /= self.config.gradient_accumulation_steps
        mean_total_loss /= self.config.gradient_accumulation_steps

        train_ce_loss = self._reduce_scalar(mean_ce_loss)
        train_aux_loss = self._reduce_scalar(mean_aux_loss)
        train_z_loss = self._reduce_scalar(mean_z_loss)
        train_loss = self._reduce_scalar(mean_total_loss)
        current_effective_batch_size = self._effective_batch_size(self.current_batch_size)
        tokens_per_iter = current_effective_batch_size * self.current_context_length
        tokens_per_sec = tokens_per_iter / max(step_time, 1e-12)
        loss_spike = self._detect_spike("loss", train_loss, self.config.loss_spike_threshold)
        grad_spike = self._detect_spike("grad_norm", grad_norm, self.config.grad_norm_spike_threshold)

        metrics = {
            "iter": self.iteration,
            "split": "train",
            "train_loss": train_loss,
            "train_ce_loss": train_ce_loss,
            "moe_aux_loss": train_aux_loss,
            "moe_z_loss": train_z_loss,
            "lr": float(lr),
            "step_time": float(step_time),
            "tokens_per_iter": int(tokens_per_iter),
            "tokens_per_sec": float(tokens_per_sec),
            "grad_norm": float(grad_norm),
            "weight_decay": float(self._current_weight_decay()),
            "batch_size": int(self.current_batch_size),
            "context_length": int(self.current_context_length),
            "attention_window_size": self.current_attention_window_size,
            "mtp_horizon": int(self.current_target_horizon),
            "mtp_weights": list(self.current_mtp_weights),
            "loss_spike": bool(loss_spike),
            "grad_norm_spike": bool(grad_spike),
        }
        moe_metrics = self._current_moe_metrics()
        metrics["moe_dropped_tokens"] = int(moe_metrics["dropped_tokens"])
        metrics["moe_tokens_dispatched"] = int(moe_metrics["tokens_dispatched"])
        return metrics

    @torch.no_grad()
    def evaluate(self) -> dict[str, Any]:
        self._apply_runtime_schedules(self.iteration)
        self.training_model.eval()
        train_metrics = self._estimate_loss(self.train_loader, self.config.eval_iters)
        valid_metrics = self._estimate_loss(self.valid_loader, self.config.eval_iters)
        return {
            "iter": self.iteration,
            "split": "eval",
            "train_loss": train_metrics["loss"],
            "valid_loss": valid_metrics["loss"],
            "train_ce_loss": train_metrics["ce_loss"],
            "valid_ce_loss": valid_metrics["ce_loss"],
            "train_moe_aux_loss": train_metrics["aux_loss"],
            "valid_moe_aux_loss": valid_metrics["aux_loss"],
            "train_moe_z_loss": train_metrics["z_loss"],
            "valid_moe_z_loss": valid_metrics["z_loss"],
        }

    def save_checkpoint(self, path: str | Path | None = None) -> Path:
        checkpoint_path = Path(path or self.checkpoint_path)
        if not self.master_process:
            return checkpoint_path
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        save_checkpoint(
            self.raw_model,
            self.optimizer,
            self.iteration,
            checkpoint_path,
            extra_state={"best_val_loss": self.best_val_loss},
        )
        return checkpoint_path

    def load_checkpoint(self, path: str | Path | None = None) -> int:
        checkpoint_path = Path(path or self.checkpoint_path)
        extra_state: dict[str, Any] = {}
        self.iteration = load_checkpoint(checkpoint_path, self.raw_model, self.optimizer, extra_state=extra_state)
        if "best_val_loss" in extra_state:
            self.best_val_loss = float(extra_state["best_val_loss"])
        return self.iteration

    def _estimate_loss(self, loader: DataLoader, num_batches: int) -> dict[str, float]:
        mean_ce_loss = 0.0
        mean_aux_loss = 0.0
        mean_z_loss = 0.0
        mean_total_loss = 0.0
        for _ in range(num_batches):
            x, y, document_ids, loss_mask = loader.get_batch(return_document_ids=True)
            with self._autocast_context():
                if self.current_target_horizon == 1:
                    logits = self.training_model(x, document_ids=document_ids)
                    ce_loss = self._compute_next_token_loss(logits, y, loss_mask)
                else:
                    hidden_states = self.training_model(x, document_ids=document_ids, return_hidden=True)
                    ce_loss = self._compute_mtp_loss(hidden_states, y, loss_mask)
                moe_metrics = self._current_moe_metrics()
                aux_loss = moe_metrics["aux_loss"].to(device=ce_loss.device, dtype=ce_loss.dtype)
                z_loss = moe_metrics["z_loss"].to(device=ce_loss.device, dtype=ce_loss.dtype)
                total_loss = ce_loss + aux_loss + z_loss
            mean_ce_loss += ce_loss.detach().item()
            mean_aux_loss += aux_loss.detach().item()
            mean_z_loss += z_loss.detach().item()
            mean_total_loss += total_loss.detach().item()
        return {
            "loss": self._reduce_scalar(mean_total_loss / num_batches),
            "ce_loss": self._reduce_scalar(mean_ce_loss / num_batches),
            "aux_loss": self._reduce_scalar(mean_aux_loss / num_batches),
            "z_loss": self._reduce_scalar(mean_z_loss / num_batches),
        }

    def _compute_next_token_loss(
        self,
        logits: torch.Tensor,
        targets: torch.Tensor,
        loss_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        token_losses = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="none")
        if loss_mask is None:
            return token_losses.mean()
        flat_mask = loss_mask.reshape(-1).to(device=token_losses.device, dtype=token_losses.dtype)
        normalizer = flat_mask.sum().clamp_min(1.0)
        return (token_losses * flat_mask).sum() / normalizer

    def _compute_mtp_loss(
        self,
        hidden_states: torch.Tensor,
        targets: torch.Tensor,
        loss_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if targets.ndim != 3:
            raise ValueError(f"MTP targets must have shape (batch, seq_len, horizon), got {tuple(targets.shape)}.")
        if targets.size(-1) != self.current_target_horizon:
            raise ValueError(
                f"MTP target horizon {targets.size(-1)} does not match current_target_horizon={self.current_target_horizon}."
            )
        if loss_mask is not None and loss_mask.shape != targets.shape:
            raise ValueError(
                f"loss_mask must match MTP target shape, got {tuple(loss_mask.shape)} and {tuple(targets.shape)}."
            )
        logits_per_head = self.raw_model.compute_mtp_logits(hidden_states, num_targets=self.current_target_horizon)
        total_loss = hidden_states.new_zeros(())
        total_weight = 0.0
        for horizon_idx, (weight, logits) in enumerate(zip(self.current_mtp_weights, logits_per_head)):
            if weight <= 0:
                continue
            target_slice = targets[:, :, horizon_idx]
            mask_slice = None if loss_mask is None else loss_mask[:, :, horizon_idx]
            loss = self._compute_next_token_loss(logits, target_slice, mask_slice)
            total_loss = total_loss + loss * float(weight)
            total_weight += float(weight)
        if total_weight == 0.0:
            raise ValueError("At least one MTP weight must be positive.")
        return total_loss / total_weight

    def _autocast_context(self):
        if not self.autocast_enabled:
            return nullcontext()
        return torch.autocast(device_type=self.device_type, dtype=self.config.torch_dtype)

    def _create_grad_scaler(self, *, enabled: bool):
        if hasattr(torch, "amp") and hasattr(torch.amp, "GradScaler"):
            return torch.amp.GradScaler("cuda", enabled=enabled)
        return torch.cuda.amp.GradScaler(enabled=enabled)

    def _scheduled_learning_rate(self, step: int) -> float:
        lr = self.scheduler.get_lr(step)
        if not self.config.scale_lr_with_batch_size:
            return lr
        current_effective_batch = self._effective_batch_size(self.current_batch_size)
        batch_ratio = current_effective_batch / max(self.base_effective_batch_size, 1)
        return lr * math.sqrt(batch_ratio)

    def _effective_batch_size(self, per_rank_batch_size: int) -> int:
        return per_rank_batch_size * self.config.gradient_accumulation_steps * self.world_size

    def _set_learning_rate(self, lr: float) -> None:
        for param_group in self.optimizer.param_groups:
            lr_scale = float(param_group.get("lr_scale", 1.0))
            param_group["lr"] = lr * lr_scale
            if self.config.schedule_weight_decay_with_lr:
                base_weight_decay = param_group["base_weight_decay"]
                scale = lr / max(self.config.learning_rate, 1e-12)
                param_group["weight_decay"] = float(base_weight_decay) * scale

    def _apply_runtime_schedules(self, completed_iterations: int) -> None:
        runtime_state = self.training_schedule.resolve(completed_iterations)
        if runtime_state.context_length != self.current_context_length:
            self.current_context_length = runtime_state.context_length
            self.train_loader.set_context_length(runtime_state.context_length)
            self.valid_loader.set_context_length(runtime_state.context_length)
        if runtime_state.batch_size != self.current_batch_size:
            self.current_batch_size = runtime_state.batch_size
            self.train_loader.set_batch_size(runtime_state.batch_size)
            self.valid_loader.set_batch_size(runtime_state.batch_size)
        if runtime_state.attention_window_size != self.current_attention_window_size:
            self.current_attention_window_size = runtime_state.attention_window_size
            self.raw_model.set_attention_window_size(runtime_state.attention_window_size)
        if runtime_state.mtp_weights != self.current_mtp_weights:
            self.current_mtp_weights = runtime_state.mtp_weights
            self.current_target_horizon = len(runtime_state.mtp_weights)
            self.train_loader.set_target_horizon(self.current_target_horizon)
            self.valid_loader.set_target_horizon(self.current_target_horizon)

    def _detect_spike(self, metric_name: str, value: float, threshold: float | None) -> bool:
        if threshold is None:
            return False
        ema_attr = f"{metric_name}_ema"
        ema_value = getattr(self, ema_attr)
        if ema_value is None:
            setattr(self, ema_attr, value)
            return False
        is_spike = value > ema_value * threshold
        setattr(self, ema_attr, 0.95 * ema_value + 0.05 * value)
        if is_spike and self.master_process:
            print(f"[warn] {metric_name} spike detected: value={value:.4f} ema={ema_value:.4f}")
        return is_spike

    def _record_metrics(self, metrics: dict[str, Any]) -> None:
        self.history.append(metrics)
        if not self.master_process:
            return
        with self.metrics_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(metrics) + "\n")

    def _write_config(self) -> None:
        if not self.config.write_config:
            return
        if not self.master_process:
            return
        serialized_config = asdict(self.config)
        serialized_config["output_dir"] = str(self.output_dir)
        with self.config_path.open("w", encoding="utf-8") as f:
            json.dump(serialized_config, f, indent=2)

    def _current_moe_metrics(self) -> dict[str, Any]:
        return getattr(self.raw_model, "last_moe_metrics", self._empty_moe_metrics())

    def _empty_moe_metrics(self) -> dict[str, Any]:
        zero = torch.zeros((), dtype=torch.float32, device=self.device)
        return {
            "aux_loss": zero,
            "z_loss": zero,
            "dropped_tokens": 0,
            "tokens_dispatched": 0,
        }

    def _reduce_scalar(self, value: float) -> float:
        if not self.distributed:
            return float(value)
        tensor = torch.tensor(value, device=self.device, dtype=torch.float32)
        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
        tensor /= self.world_size
        return float(tensor.item())

    def _current_weight_decay(self) -> float:
        return max(float(group.get("weight_decay", 0.0)) for group in self.optimizer.param_groups)

    def _ensure_optimizer_group_metadata(self) -> None:
        for param_group in self.optimizer.param_groups:
            if "base_weight_decay" not in param_group:
                param_group["base_weight_decay"] = float(param_group.get("weight_decay", 0.0))
            if "lr_scale" not in param_group:
                param_group["lr_scale"] = 1.0
            if "cautious_weight_decay" not in param_group:
                param_group["cautious_weight_decay"] = False


def _normalize_mtp_weights(
    mtp_weights: tuple[float, ...] | list[float],
    *,
    field_name: str,
    allow_empty: bool = False,
) -> tuple[float, ...]:
    normalized = tuple(float(weight) for weight in mtp_weights)
    if not allow_empty and not normalized:
        raise ValueError(f"{field_name} must contain at least one weight.")
    if any(weight < 0 for weight in normalized):
        raise ValueError(f"{field_name} must contain non-negative weights, got {normalized}.")
    if normalized and normalized[0] <= 0:
        raise ValueError(f"{field_name} must have a positive next-token weight, got {normalized}.")
    if normalized and sum(normalized) <= 0:
        raise ValueError(f"{field_name} must have a positive total weight, got {normalized}.")
    return normalized


def _build_training_schedule(
    config: TrainerConfig,
    *,
    default_attention_window_size: int | None,
) -> TrainingSchedule:
    stage_updates: dict[int, dict[str, Any]] = {}
    for start_iteration, value in config.context_length_schedule:
        stage_updates.setdefault(start_iteration, {})["context_length"] = value
    for start_iteration, value in config.batch_size_schedule:
        stage_updates.setdefault(start_iteration, {})["batch_size"] = value
    for stage in config.training_stages:
        stage_update = stage_updates.setdefault(stage.start_iteration, {})
        if stage.batch_size is not None:
            stage_update["batch_size"] = stage.batch_size
        if stage.context_length is not None:
            stage_update["context_length"] = stage.context_length
        if stage.override_attention_window_size:
            stage_update["override_attention_window_size"] = True
            stage_update["attention_window_size"] = stage.attention_window_size
        if stage.mtp_weights:
            stage_update["mtp_weights"] = stage.mtp_weights
    stages = tuple(
        TrainingStage(start_iteration=start_iteration, **stage_update)
        for start_iteration, stage_update in sorted(stage_updates.items())
    )
    return TrainingSchedule(
        batch_size=config.batch_size,
        context_length=config.context_length,
        attention_window_size=default_attention_window_size,
        mtp_weights=config.mtp_weights,
        stages=stages,
    )


__all__ = [
    "Trainer",
    "TrainerConfig",
    "TrainingSchedule",
    "TrainingStage",
]