Blame for src/nida/train/loop.py

"""The trainer: flow-matching training loop with resumable, bit-identical state.

Randomness. One ``cfg.seed`` drives everything: Python/NumPy/torch global generators
(model init) are seeded with it, and two private generators are spawned from
``SeedSequence(seed)``: a NumPy ``Generator`` for data randomness (target spans, CFG
dropout) and a CPU torch ``Generator`` for flow randomness (``t`` and the noise ``x0``,
drawn on CPU so a run is device independent up to float error). Evaluation uses its own
fixed generators and never touches the training streams. All of them, the global RNGs and
the sampler position are checkpointed, so ``resume`` reproduces the uninterrupted run
exactly (tested).

One optimizer step = ``grad_accum_steps`` micro-batches; the loss of micro-batch ``i`` is
weighted by ``target_frames_i / total_target_frames`` so the step equals one big batch.
Losses/reductions run in float32; ``bfloat16`` autocast is used on CUDA only.

E12 (``prompt_pairing.cross_p > 0``, :mod:`nida.train.pairing`): the *training* stream draws,
per item and before the target spans and the dropout, three uniforms from the data generator
(always ``3 * batch`` of them, whatever the items are) to decide whether the item's prompt
comes from another segment of its group. With ``cross_p == 0`` nothing is drawn and the
batches are bit-identical to those of earlier versions. Evaluation always keeps the in-clip
prompt (``val_loss`` stays comparable across runs).

``overfit_one_batch`` (ladder step 2, docs/04): the first batch of epoch 0 is collated once
(its target span fixed) and reused every step; ``t`` and the noise still vary.
"""

from __future__ import annotations

import contextlib
import json
import logging
import random
import signal
import threading
import time
from collections.abc import Iterator, Mapping
from dataclasses import asdict, dataclass, field, replace
from pathlib import Path
from types import FrameType, TracebackType
from typing import cast

import numpy as np
import torch
from torch import Tensor

from nida.model.align import even_char_phase
from nida.model.dit import DUR_FILM_PREFIX, Conditions, DiT, ModelConfig
from nida.model.flow import interpolate, masked_mse, sample_t, target_velocity
from nida.train.checkpoint import (
    CheckpointError,
    capture_rng,
    latest_checkpoint,
    load_checkpoint,
    restore_rng,
    save_checkpoint,
)
from nida.train.collate import Batch, CollateSpec, collate
from nida.train.config import ConfigError, TrainConfig
from nida.train.data import (
    EMOTIONS,
    N_LOUDNESS_BUCKETS,
    N_RATE_BUCKETS,
    LatentDataset,
    LatentStats,
    compute_latent_stats,
)
from nida.train.optim import EMA, build_optimizer, build_scheduler
from nida.train.pairing import (
    GroupIndex,
    PairingCounts,
    build_paired_inputs,
    draw_pairing_uniforms,
    match_speaker_dirs,
    reserved_lengths,
    speaker_group_keys,
)
from nida.train.sampler import BatchStream, BudgetError, FrameBudgetSampler, SamplerPosition
from nida.utils.io import atomic_write_json, atomic_write_text

__all__ = [
    "JsonlLog",
    "NonFiniteError",
    "ResumeError",
    "StepMetrics",
    "StopFlag",
    "TrainResult",
    "Trainer",
    "peak_memory_gb",
    "seed_everything",
    "sigterm_handler",
    "truncate_jsonl_after_step",
]

log = logging.getLogger(__name__)

METRICS_NAME = "metrics.jsonl"
CHECKPOINT_DIR = "checkpoints"
#: Holds the single checkpoint with the lowest EMA validation loss.
BEST_DIR = "best"
#: Rewritten at every log step; its age tells a watcher whether the run is alive.
HEARTBEAT_NAME = "heartbeat.json"


class NonFiniteError(FloatingPointError):
    """The loss or the gradient became NaN/Inf; training stops instead of continuing."""


class ResumeError(RuntimeError):
    """The checkpoint cannot be resumed with this config / dataset."""


class StopFlag:
    """Cooperative stop request (set by the SIGTERM handler, polled between steps)."""

    def __init__(self) -> None:
        self._event = threading.Event()

    def set(self) -> None:
        self._event.set()

    def is_set(self) -> bool:
        return self._event.is_set()


@contextlib.contextmanager
def sigterm_handler(flag: StopFlag) -> Iterator[None]:
    """Route SIGTERM to ``flag.set()`` for the duration of the block (main thread only).

    The handler only sets a flag; the trainer saves at the next step boundary.
    """
    if threading.current_thread() is not threading.main_thread():
        log.warning("not in the main thread: SIGTERM handler not installed")
        yield
        return

    def handler(signum: int, frame: FrameType | None) -> None:
        flag.set()

    previous = signal.signal(signal.SIGTERM, handler)
    try:
        yield
    finally:
        signal.signal(signal.SIGTERM, previous if previous is not None else signal.SIG_DFL)


def seed_everything(seed: int) -> None:
    """Seed the global Python, NumPy and torch (CPU + CUDA) generators."""
    random.seed(seed)
    np.random.seed(seed % (2**32))
    torch.manual_seed(seed)


class JsonlLog:
    """Append-only JSON-lines log, flushed after every record."""

    def __init__(self, path: Path) -> None:
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        self._file = path.open("a", encoding="utf-8")

    def write(self, record: Mapping[str, object]) -> None:
        self._file.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
        self._file.flush()

    def close(self) -> None:
        self._file.close()

    def __enter__(self) -> JsonlLog:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.close()


def truncate_jsonl_after_step(path: Path, step: int) -> int:
    """Keep only records with ``step <= step`` (records without a step are kept).

    Used on resume so replayed steps are not logged twice. Unparseable lines (a torn last
    write after a kill) are dropped with a warning. Returns the number of removed lines.
    """
    if not path.exists():
        return 0
    kept: list[str] = []
    removed = 0
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            log.warning("dropping unparseable line from %s: %r", path.name, line[:80])
            removed += 1
            continue
        record_step = record.get("step") if isinstance(record, dict) else None
        if isinstance(record_step, int) and record_step > step:
            removed += 1
        else:
            kept.append(line)
    atomic_write_text(path, "".join(f"{line}\n" for line in kept))
    return removed


#: The only model-config change :meth:`Trainer.init_from` accepts: field -> (parent, child).
_ALLOWED_INIT_DIFF: dict[str, tuple[object, object]] = {"dur_cond": ("none", "emotion")}


def model_config_diff(parent: ModelConfig, child: ModelConfig) -> dict[str, tuple[object, object]]:
    """Fields that differ between two resolved model configs: ``{name: (parent, child)}``."""
    a, b = asdict(parent), asdict(child)
    return {name: (a[name], b[name]) for name in a if a[name] != b[name]}


def _allowed_diff_record() -> dict[str, list[str]]:
    """The accepted difference as stored in checkpoints / ``run_info.json``."""
    return {name: [str(v) for v in pair] for name, pair in _ALLOWED_INIT_DIFF.items()}


def _check_added_keys(path: Path, present: set[str], expected: set[str], added: set[str]) -> None:
    """The checkpoint's state-dict keys must be the model's minus exactly ``added``.

    Raises:
        CheckpointError: any other key is missing, or the checkpoint has an unknown key, or
            it already holds some of ``added``.
    """
    missing, unexpected = expected - present, present - expected
    if missing != added or unexpected:
        raise CheckpointError(
            f"{path}: model keys differ beyond the new {sorted(added)} parameters "
            f"(missing {sorted(missing - added)[:5]}, not missing but expected absent "
            f"{sorted(added - missing)[:5]}, unexpected {sorted(unexpected)[:5]})"
        )


class TempoSums:
    """Per-emotion sums for the validation tempo metric ``val_dur_tempo_<emotion>``.

    For every real character of an item with a known emotion: ``log(1 + d_hat)`` under the
    item's true emotion/intensity minus ``log(1 + d_hat)`` under the null emotion (the
    emotion-averaged rhythm). A positive mean = slower than average, negative = faster;
    exactly 0 for ``dur_cond="none"`` (the head ignores the emotion).
    """

    def __init__(self, null_id: int) -> None:
        self.null_id = null_id
        self.diff_sum: dict[int, float] = {}
        self.num_chars: dict[int, int] = {}

    @torch.no_grad()
    def add(self, model: DiT, c: Conditions) -> None:
        """Accumulate one batch of conditions (aligned model, no dropout)."""
        assert c.char_mask is not None
        null = replace(
            c,
            emotion=torch.full_like(c.emotion, self.null_id),
            intensity_known=torch.zeros_like(c.intensity_known),
        )
        diff = model.log_durations(c).float() - model.log_durations(null).float()
        zero = torch.zeros((), device=diff.device)
        per_item = torch.where(c.char_mask, diff, zero).sum(dim=1).cpu().tolist()
        chars = c.char_mask.sum(dim=1).cpu().tolist()
        for emotion, total, n in zip(c.emotion.cpu().tolist(), per_item, chars, strict=True):
            if emotion == self.null_id or n == 0:
                continue
            self.diff_sum[emotion] = self.diff_sum.get(emotion, 0.0) + float(total)
            self.num_chars[emotion] = self.num_chars.get(emotion, 0) + int(n)

    def metrics(self) -> dict[str, float]:
        """``val_dur_tempo_<emotion>`` for every emotion seen (mean over its characters)."""
        return {
            f"val_dur_tempo_{EMOTIONS[e]}": self.diff_sum[e] / self.num_chars[e]
            for e in sorted(self.diff_sum)
        }


def peak_memory_gb(device: torch.device) -> float | None:
    """Peak tensor memory allocated on ``device`` so far in GB (1e9 bytes); ``None`` unless CUDA."""
    if device.type != "cuda":
        return None
    return float(torch.cuda.max_memory_allocated(device)) / 1e9


def _num_chars(batch: Batch) -> int:
    """Real characters of an aligned batch (0 when it has no character mask)."""
    mask = batch.conditions.char_mask
    return 0 if mask is None else int(mask.sum())


@dataclass(frozen=True)
class StepMetrics:
    """Metrics of one optimizer step (``step`` = number of steps done including this one)."""

    step: int
    loss: float
    lr: float
    grad_norm: float
    frames_per_s: float
    epoch: int
    batch_items: int
    real_frames: int
    target_frames: int
    #: Aligned text path only (E11): the two terms of ``loss`` (``flow + w · dur``).
    flow_loss: float | None = None
    dur_loss: float | None = None
    #: E12 (``prompt_pairing.cross_p > 0``): ``paired / items`` of this step, and
    #: ``fallbacks / (paired + fallbacks)``: of the items selected for a cross-segment prompt,
    #: the fraction that kept the in-clip prompt (reasons: ``nida.train.pairing``).
    cross_frac: float | None = None
    cross_fallback_frac: float | None = None
    #: Peak GPU memory allocated by tensors since the start of the process, in GB (1e9 bytes);
    #: ``None`` on CPU (and then absent from the record).
    max_mem_gb: float | None = None

    def to_record(self) -> dict[str, object]:
        record: dict[str, object] = {
            "kind": "train",
            "step": self.step,
            "loss": self.loss,
            "lr": self.lr,
            "grad_norm": self.grad_norm,
            "frames_per_s": self.frames_per_s,
            "epoch": self.epoch,
            "batch_items": self.batch_items,
            "real_frames": self.real_frames,
            "target_frames": self.target_frames,
        }
        if self.flow_loss is not None:
            record["flow_loss"] = self.flow_loss
        if self.dur_loss is not None:
            record["dur_loss"] = self.dur_loss
        if self.cross_frac is not None:
            record["cross_frac"] = self.cross_frac
        if self.cross_fallback_frac is not None:
            record["cross_fallback_frac"] = self.cross_fallback_frac
        if self.max_mem_gb is not None:
            record["max_mem_gb"] = self.max_mem_gb
        return record


@dataclass
class TrainResult:
    """Outcome of :meth:`Trainer.run`."""

    steps_done: int
    completed: bool
    stopped_by_signal: bool
    last_loss: float | None
    mean_frames_per_s: float | None
    resumed_from_step: int | None
    checkpoints: list[str] = field(default_factory=list)
    evals: list[dict[str, float]] = field(default_factory=list)
    best_checkpoint: str | None = None


@dataclass(frozen=True)
class _CrossEval:
    """Fixed batches (validation item indices) and speaker index of the cross-segment
    validation pass."""

    groups: GroupIndex
    batches: list[list[int]]


class Trainer:
    """Flow-matching trainer (see module docstring).

    Args:
        cfg: validated training config.
        train_data: accepted training segments (must be non-empty).
        out_dir: run workspace; checkpoints go to ``out_dir/checkpoints``, the metrics log
            to ``out_dir/metrics.jsonl``.
        val_data: optional validation segments (evaluated every ``cfg.eval_every`` steps).
        device: defaults to CPU. Autocast (bf16) is used only on CUDA.
        git_commit: recorded in every checkpoint.
        stop: stop flag polled between steps.

    Raises:
        ValueError: dataset incompatible with the model config (latent dim, buckets).
        BudgetError: an item longer than ``batch.frame_budget``.
    """

    def __init__(
        self,
        cfg: TrainConfig,
        train_data: LatentDataset,
        out_dir: Path,
        *,
        val_data: LatentDataset | None = None,
        device: torch.device | None = None,
        git_commit: str | None = None,
        stop: StopFlag | None = None,
    ) -> None:
        self.cfg = cfg
        self.model_cfg = cfg.model_config()
        if len(train_data) == 0:
            raise ValueError("training dataset is empty")
        if train_data.latent_dim != self.model_cfg.latent_dim:
            raise ValueError(
                f"dataset latent_dim {train_data.latent_dim} != model {self.model_cfg.latent_dim}"
            )
        if (self.model_cfg.n_rate_buckets, self.model_cfg.n_loudness_buckets) != (
            N_RATE_BUCKETS,
            N_LOUDNESS_BUCKETS,
        ):
            raise ValueError("model bucket counts do not match the data format")
        self.train_data = train_data
        self.val_data = val_data if val_data is not None and len(val_data) > 0 else None
        # Computed on the training split only; validation uses the same transform.
        self.latent_stats = compute_latent_stats(train_data)
        self.out_dir = out_dir
        self.checkpoint_dir = out_dir / CHECKPOINT_DIR
        self.metrics_path = out_dir / METRICS_NAME
        self.device = device if device is not None else torch.device("cpu")
        self.git_commit = git_commit
        self.stop = stop if stop is not None else StopFlag()
        self.use_amp = self.device.type == "cuda"

        seed_everything(cfg.seed)
        self.model = DiT(self.model_cfg).to(self.device)
        self.optimizer = build_optimizer(self.model, cfg)
        self.scheduler = build_scheduler(self.optimizer, cfg)
        self.ema = EMA(self.model, decay=cfg.ema.decay, warmup=cfg.ema.warmup)

        pairing = cfg.prompt_pairing
        self._groups: GroupIndex | None = None
        #: Training directories that ``prompt_pairing.speaker_dirs`` marks as speaker ids.
        self.speaker_dirs = match_speaker_dirs(train_data.data_dirs, pairing.speaker_dirs)
        for pattern in self.speaker_dirs.unmatched_patterns:
            log.warning("prompt_pairing.speaker_dirs pattern %r matches no data directory", pattern)
        #: E12 pairing outcomes of the whole run (checkpointed; for logs and reports).
        self.pairing_totals = PairingCounts()
        self._step_pairing = PairingCounts()
        train_lengths = train_data.lengths
        if pairing.enabled:
            self._groups = self._build_group_index(train_data)
        if pairing.reserve_frames:
            # Reserve room for the prompt of a composite item; identical for the control arm.
            # Scope "speaker_dirs": only items that can ever be paired (those of speaker
            # directories; partners come from the same speaker, hence the same directory).
            train_lengths = reserved_lengths(
                train_lengths,
                reserve_frames=pairing.reserve_frames,
                eligible=(
                    None
                    if pairing.reserve_scope == "all"
                    else [
                        key is not None for key in speaker_group_keys(train_data, self.speaker_dirs)
                    ]
                ),
            )
        try:
            self.sampler = FrameBudgetSampler(
                train_lengths,
                frame_budget=cfg.batch.frame_budget,
                seed=cfg.seed,
                pool_size=cfg.batch.bucket_pool_size,
                max_items=cfg.batch.max_items,
                stream=0,
            )
        except BudgetError as exc:
            if not pairing.reserve_frames:
                raise
            raise BudgetError(
                f"{exc} (lengths include prompt_pairing.reserve_frames={pairing.reserve_frames})"
            ) from exc
        self.stream = BatchStream(self.sampler)
        self._eval_sampler: FrameBudgetSampler | None = None
        if self.val_data is not None:
            self._eval_sampler = FrameBudgetSampler(
                self.val_data.lengths,
                frame_budget=cfg.batch.frame_budget,
                seed=cfg.seed,
                pool_size=cfg.batch.bucket_pool_size,
                max_items=cfg.batch.max_items,
                stream=1,
            )

        self._cross_eval = self._build_cross_eval()
        data_seq, flow_seq, overfit_seq = np.random.SeedSequence(cfg.seed).spawn(3)
        self.data_rng = np.random.default_rng(data_seq)
        self.flow_gen = torch.Generator().manual_seed(_seed_int(flow_seq))
        self._train_spec = CollateSpec.from_config(cfg, self.model_cfg, dropout=True)
        self._eval_spec = CollateSpec.from_config(cfg, self.model_cfg, dropout=False)

        self.step = 0
        #: Lowest EMA validation loss seen so far; its checkpoint is kept in ``best/``.
        self.best_val_loss: float | None = None
        self.resumed_from_step: int | None = None
        #: Checkpoint whose weights started this run (:meth:`init_from`), if any.
        self.init_from_path: str | None = None
        #: Model-config difference to that checkpoint (``{"dur_cond": ["none", "emotion"]}``),
        #: ``None`` when the configs are identical or there was no ``init_from``.
        self.init_from_model_diff: dict[str, list[str]] | None = None
        #: True when ``latent_stats`` come from that checkpoint rather than this run's data.
        self.latent_stats_inherited = False
        self._last_saved_step: int | None = None
        self._overfit_batch: Batch | None = None
        if cfg.overfit_one_batch:
            first = self.sampler.epoch_batches(0)[0]
            batch = self._collate(
                self.train_data, first, self._train_spec, np.random.default_rng(overfit_seq)
            )
            self._overfit_batch = batch.to(self.device)

    # -- data -----------------------------------------------------------------------------

    def _build_group_index(self, data: LatentDataset) -> GroupIndex:
        """Speaker index for cross-segment prompts; refuses data that cannot use them."""
        no_units = sum(1 for s in data.segments if not s.units)
        if no_units:
            raise ValueError(
                f"prompt_pairing.cross_p > 0 needs the unit timeline of every training item, "
                f"{no_units} have none (load the dataset with timings)"
            )
        groups = GroupIndex(speaker_group_keys(data, self.speaker_dirs))
        if groups.num_eligible_items == 0:
            raise ValueError(
                "prompt_pairing.cross_p > 0 but no data directory matching "
                f"speaker_dirs={list(self.cfg.prompt_pairing.speaker_dirs)} (matched: "
                f"{list(self.speaker_dirs.matched)}) has a speaker with two items"
            )
        log.info(
            "prompt pairing: cross_p=%g, speaker dirs %s, %d speakers, %d of %d items have a "
            "partner (%d items are in speaker directories)",
            self.cfg.prompt_pairing.cross_p,
            list(self.speaker_dirs.matched),
            groups.num_groups,
            groups.num_eligible_items,
            len(data),
            groups.num_speaker_items,
        )
        return groups

    def _build_cross_eval(self) -> _CrossEval | None:
        """Cross-segment validation pass (``val_loss*_cross``): the validation items of the
        speaker directories that have a speaker with two items, in their own deterministic
        batches. ``None`` (with a log note) when it cannot run."""
        pairing = self.cfg.prompt_pairing
        data = self.val_data
        if data is None or not pairing.speaker_dirs:
            return None
        if self.model_cfg.text_path != "aligned":
            log.info("no cross-segment validation: it needs the aligned text path")
            return None
        if any(not s.units for s in data.segments):
            log.info("no cross-segment validation: validation items have no unit timeline")
            return None
        speakers = match_speaker_dirs(data.data_dirs, pairing.speaker_dirs)
        groups = GroupIndex(speaker_group_keys(data, speakers))
        items = [i for i in range(len(data)) if len(groups.members(i)) >= 2]
        if not items:
            log.info(
                "no cross-segment validation: no validation speaker in %s has two items",
                list(speakers.matched),
            )
            return None
        lengths = [data[i].num_frames + pairing.max_prompt_frames for i in items]
        try:
            sampler = FrameBudgetSampler(
                lengths,
                frame_budget=self.cfg.batch.frame_budget,
                seed=self.cfg.seed,
                pool_size=self.cfg.batch.bucket_pool_size,
                max_items=self.cfg.batch.max_items,
                stream=2,
            )
        except BudgetError as exc:
            raise BudgetError(
                f"{exc} (cross-segment validation lengths include max_prompt_frames="
                f"{pairing.max_prompt_frames})"
            ) from exc
        batches = [[items[k] for k in b] for b in sampler.epoch_batches(0)]
        return _CrossEval(groups, batches[: self.cfg.eval_max_batches])

    def _collate(
        self,
        data: LatentDataset,
        indices: list[int],
        spec: CollateSpec,
        rng: np.random.Generator,
    ) -> Batch:
        segments = [data[i] for i in indices]
        latents = [self.latent_stats.normalize(data.load_latents(i)) for i in indices]
        return collate(segments, latents, spec, rng)

    def _next_batch(self) -> Batch:
        if self._overfit_batch is not None:
            return self._overfit_batch
        indices = self.stream.next_batch()
        if self._groups is not None:
            return self._collate_paired(indices).to(self.device)
        return self._collate(self.train_data, indices, self._train_spec, self.data_rng).to(
            self.device
        )

    def _collate_paired(self, indices: list[int]) -> Batch:
        """Training batch with cross-segment prompts (E12): the pairing uniforms are drawn
        first, then :func:`collate` draws the spans and the dropout from the same generator."""
        assert self._groups is not None
        draws = draw_pairing_uniforms(self.data_rng, len(indices))
        inputs = build_paired_inputs(
            self.train_data,
            indices,
            self._groups,
            self.cfg.prompt_pairing,
            draws,
            self.latent_stats.normalize,
        )
        self._step_pairing.add(inputs.counts)
        self.pairing_totals.add(inputs.counts)
        return collate(
            inputs.segments,
            inputs.latents,
            self._train_spec,
            self.data_rng,
            target_spans=inputs.target_spans,
        )

    # -- one step -------------------------------------------------------------------------

    def _batch_loss(self, batch: Batch, generator: torch.Generator) -> Tensor:
        """Training objective of ``batch``: flow loss (+ ``dur_loss_weight`` × duration loss
        on the aligned text path); float32 scalar."""
        flow, dur = self._batch_losses(batch, generator)
        if dur is None:
            return flow
        return flow + self.model.cfg.dur_loss_weight * dur

    def _batch_losses(
        self, batch: Batch, generator: torch.Generator
    ) -> tuple[Tensor, Tensor | None]:
        """Masked flow-matching MSE of ``batch`` and, on the aligned text path, the
        duration L1 ``mean_i |l_hat_i − log(1 + d̄_i)|`` over real characters (float32 scalars).

        ``t`` and ``x0`` are drawn from ``generator`` on CPU (``t`` first, then ``x0``).
        """
        x1 = batch.x1
        t = sample_t(
            x1.shape[0],
            kind=self.cfg.t_sampling.kind,
            mean=self.cfg.t_sampling.mean,
            std=self.cfg.t_sampling.std,
            generator=generator,
        ).to(self.device)
        x0 = torch.randn(x1.shape, generator=generator).to(self.device)
        x_t = interpolate(x0, x1, t)
        velocity = target_velocity(x0, x1)
        mask = batch.loss_mask & batch.conditions.frame_mask
        if self.model.cfg.text_path != "aligned":
            with torch.autocast(
                device_type=self.device.type, dtype=torch.bfloat16, enabled=self.use_amp
            ):
                pred = self.model(x_t, t, batch.conditions)
            return masked_mse(pred.float(), velocity.float(), mask), None
        char_mask, target = batch.conditions.char_mask, batch.char_dur_target
        if char_mask is None or target is None:
            raise ValueError("aligned text path: the batch has no duration targets")
        with torch.autocast(
            device_type=self.device.type, dtype=torch.bfloat16, enabled=self.use_amp
        ):
            pred, log_dur = self.model(x_t, t, batch.conditions, return_aux=True)
        flow = masked_mse(pred.float(), velocity.float(), mask)
        abs_err = torch.where(
            char_mask, (log_dur.float() - target.float()).abs(), torch.zeros((), device=x1.device)
        )
        dur = abs_err.sum() / char_mask.sum().clamp(min=1)
        return flow, dur

    def train_step(self) -> StepMetrics:
        """One optimizer step (``grad_accum_steps`` micro-batches)."""
        started = time.perf_counter()
        self.model.train()
        self._step_pairing = PairingCounts()
        batches = [self._next_batch() for _ in range(self.cfg.batch.grad_accum_steps)]
        total_target = sum(b.num_target_frames for b in batches)
        self.optimizer.zero_grad(set_to_none=True)
        loss_value = 0.0
        aligned = self.model.cfg.text_path == "aligned"
        flow_value: float | None = 0.0 if aligned else None
        dur_value: float | None = 0.0 if aligned else None
        total_chars = sum(_num_chars(b) for b in batches) if aligned else 0
        for batch in batches:
            if aligned:
                # Each term is weighted by its own unit across micro-batches: the flow loss by
                # target frames, the duration loss (a per-character mean) by characters.
                flow, dur = self._batch_losses(batch, self.flow_gen)
                assert dur is not None and flow_value is not None and dur_value is not None
                w_flow = batch.num_target_frames / total_target
                w_dur = _num_chars(batch) / total_chars
                loss = flow * w_flow + self.model.cfg.dur_loss_weight * dur * w_dur
                if not bool(torch.isfinite(loss)):
                    raise NonFiniteError(f"non-finite loss at step {self.step + 1}")
                loss.backward()  # type: ignore[no-untyped-call]  # torch stub gap
                loss_value += float(loss.detach())
                flow_value += float(flow.detach()) * w_flow
                dur_value += float(dur.detach()) * w_dur
                continue
            loss = self._batch_loss(batch, self.flow_gen)
            if not bool(torch.isfinite(loss)):
                raise NonFiniteError(f"non-finite loss at step {self.step + 1}")
            weight = batch.num_target_frames / total_target
            (loss * weight).backward()  # type: ignore[no-untyped-call]  # torch stub gap
            loss_value += float(loss.detach()) * weight
        params = [p for p in self.model.parameters() if p.requires_grad]
        try:
            grad_norm = float(
                torch.nn.utils.clip_grad_norm_(
                    params, self.cfg.optim.grad_clip, error_if_nonfinite=True
                )
            )
        except RuntimeError as exc:
            raise NonFiniteError(f"non-finite gradient at step {self.step + 1}: {exc}") from exc
        lr = float(self.optimizer.param_groups[0]["lr"])
        self.optimizer.step()
        self.scheduler.step()
        self.ema.update(self.model)
        self.step += 1
        real_frames = sum(b.num_real_frames for b in batches)
        elapsed = max(time.perf_counter() - started, 1e-9)
        return StepMetrics(
            step=self.step,
            loss=loss_value,
            lr=lr,
            grad_norm=grad_norm,
            frames_per_s=real_frames / elapsed,
            epoch=self.stream.position.epoch,
            batch_items=sum(b.num_items for b in batches),
            real_frames=real_frames,
            target_frames=total_target,
            flow_loss=flow_value,
            dur_loss=dur_value,
            cross_frac=self._step_pairing.cross_frac if self._groups is not None else None,
            cross_fallback_frac=(
                self._step_pairing.fallback_frac if self._groups is not None else None
            ),
            max_mem_gb=peak_memory_gb(self.device),
        )

    # -- evaluation -----------------------------------------------------------------------

    @torch.no_grad()
    def evaluate(self) -> dict[str, float] | None:
        """Validation loss with raw and EMA weights, or ``None`` without validation data.

        Fixed batches, fixed target spans, fixed ``t``/noise (own generators), no dropout:
        successive evaluations are comparable and the training RNG streams are untouched.
        """
        if self.val_data is None or self._eval_sampler is None:
            return None
        batches = self._eval_sampler.epoch_batches(0)[: self.cfg.eval_max_batches]
        seed = _seed_int(np.random.SeedSequence([self.cfg.seed, 1]))
        result: dict[str, float] = {"step": float(self.step), "num_batches": float(len(batches))}
        self.model.eval()
        # ``val_loss_ema_notext`` drops the text of every item (same batches, spans and
        # noise): its gap to ``val_loss_ema`` measures how much the model uses the text.
        passes = [
            ("val_loss", False, False, False),
            ("val_loss_ema", True, False, False),
            ("val_loss_ema_notext", True, True, False),
        ]
        if self.model.cfg.text_path == "aligned":
            # Inference spaces characters evenly inside a unit (E11 §7); training uses the
            # predicted division. This pass measures the inference geometry.
            passes.append(("val_loss_ema_evenphase", True, False, True))
        for label, use_ema, no_text, even_phase in passes:
            rng = np.random.default_rng(seed)
            generator = torch.Generator().manual_seed(seed)
            total, count = 0.0, 0
            dur_total, dur_batches = 0.0, 0
            tempo = TempoSums(self.model_cfg.n_emotions)
            with self.ema.swapped_in(self.model) if use_ema else contextlib.nullcontext():
                for indices in batches:
                    batch = self._collate(self.val_data, indices, self._eval_spec, rng).to(
                        self.device
                    )
                    if label == "val_loss_ema" and self.model_cfg.text_path == "aligned":
                        tempo.add(self.model, batch.conditions)
                    if no_text:
                        batch = replace(
                            batch,
                            conditions=replace(
                                batch.conditions,
                                drop_text=torch.ones_like(batch.conditions.drop_text),
                            ),
                        )
                    if even_phase:
                        c = batch.conditions
                        assert c.char_unit is not None and c.unit_first_char is not None
                        assert c.char_mask is not None
                        phase = even_char_phase(c.char_unit, c.unit_first_char, c.char_mask)
                        batch = replace(batch, conditions=replace(c, char_phase=phase))
                    n = batch.num_target_frames
                    flow, dur = self._batch_losses(batch, generator)
                    total += float(flow) * n
                    count += n
                    if dur is not None:
                        dur_total += float(dur)
                        dur_batches += 1
            # Flow loss only: comparable across text paths (the duration loss is reported
            # separately as ``<label>_dur``).
            result[label] = total / count
            if dur_batches:
                result[f"{label}_dur"] = dur_total / dur_batches
            result.update(tempo.metrics())
        if self._cross_eval is not None:
            result.update(self._evaluate_cross(self._cross_eval, self.val_data))
        self.model.train()
        return result

    def _evaluate_cross(self, cross: _CrossEval, data: LatentDataset) -> dict[str, float]:
        """``val_loss_cross`` / ``val_loss_ema_cross`` (E12): the flow loss on the target clip
        of validation items whose prompt is another segment of the same speaker, i.e. the
        inference situation. Deterministic and independent of the training streams: fixed
        batches, a pairing generator seeded afresh for every pass (raw and EMA weights see the
        same pairs), the evaluation noise seed, no dropout. Items that cannot be paired are
        left out of the loss and counted (``val_cross_items`` / ``val_cross_fallbacks``)."""
        seed = _seed_int(np.random.SeedSequence([self.cfg.seed, 1]))
        pair_seed = _seed_int(np.random.SeedSequence([self.cfg.seed, 2]))
        spec = self.cfg.prompt_pairing
        out: dict[str, float] = {}
        for label, use_ema in (("val_loss_cross", False), ("val_loss_ema_cross", True)):
            pair_rng = np.random.default_rng(pair_seed)
            rng = np.random.default_rng(seed)
            generator = torch.Generator().manual_seed(seed)
            total, count, paired, fallbacks = 0.0, 0, 0, 0
            dur_total, dur_batches = 0.0, 0
            with self.ema.swapped_in(self.model) if use_ema else contextlib.nullcontext():
                for indices in cross.batches:
                    draws = draw_pairing_uniforms(pair_rng, len(indices))
                    inputs = build_paired_inputs(
                        data,
                        indices,
                        cross.groups,
                        spec,
                        draws,
                        self.latent_stats.normalize,
                        select_all=True,
                    )
                    fallbacks += inputs.counts.num_fallbacks
                    keep = [k for k, span in enumerate(inputs.target_spans) if span is not None]
                    if not keep:
                        continue
                    spans = [inputs.target_spans[k] for k in keep]
                    batch = collate(
                        [inputs.segments[k] for k in keep],
                        [inputs.latents[k] for k in keep],
                        self._eval_spec,
                        rng,
                        target_spans=spans,
                    ).to(self.device)
                    paired += len(keep)
                    n = batch.num_target_frames
                    flow, dur = self._batch_losses(batch, generator)
                    total += float(flow) * n
                    count += n
                    if dur is not None:
                        dur_total += float(dur)
                        dur_batches += 1
            if count == 0:
                log.warning("%s: no validation item could be paired", label)
                continue
            out[label] = total / count
            if dur_batches:
                out[f"{label}_dur"] = dur_total / dur_batches
            out["val_cross_items"] = float(paired)
            out["val_cross_fallbacks"] = float(fallbacks)
        return out

    # -- checkpointing --------------------------------------------------------------------

    def state_dict(self) -> dict[str, object]:
        """Everything needed to continue bit-identically (see :mod:`nida.train.checkpoint`)."""
        position = self.stream.position
        return {
            "step": self.step,
            "model": self.model.state_dict(),
            "ema": self.ema.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "sampler": {"epoch": position.epoch, "batch_index": position.batch_index},
            "rng": {
                "global": capture_rng(),
                "data": self.data_rng.bit_generator.state,
                "flow": self.flow_gen.get_state(),
            },
            "config": self.cfg.to_dict(),
            "seed": self.cfg.seed,
            "git_commit": self.git_commit,
            "dataset_fingerprint": self.train_data.fingerprint(),
            "latent_stats": self.latent_stats.to_dict(),
            "best_val_loss": self.best_val_loss,
            "init_from": self.init_from_path,
            "init_from_model_diff": self.init_from_model_diff,
            "latent_stats_inherited": self.latent_stats_inherited,
            "pairing": self.pairing_totals.to_dict(),
        }

    def save(self) -> Path:
        path = save_checkpoint(
            self.checkpoint_dir, self.step, self.state_dict(), self.cfg.keep_last
        )
        self._last_saved_step = self.step
        log.info("saved %s", path)
        return path

    def resume(self, path: Path) -> None:
        """Load a checkpoint into this trainer.

        Raises:
            ResumeError: config (other than logging/eval/save settings) or dataset differ.
            CheckpointError: unreadable or incomplete checkpoint.
        """
        ckpt = load_checkpoint(path)
        try:
            saved_cfg = TrainConfig.from_mapping(cast("Mapping[str, object]", ckpt["config"]))
        except (KeyError, ConfigError) as exc:
            raise CheckpointError(f"{path}: bad config in checkpoint: {exc}") from exc
        diff = self.cfg.resume_diff(saved_cfg)
        if diff:
            raise ResumeError(f"config differs from the checkpoint's in {diff}; cannot resume")
        if ckpt.get("dataset_fingerprint") != self.train_data.fingerprint():
            raise ResumeError("training data differ from the checkpoint's (fingerprint mismatch)")
        try:
            saved_stats = LatentStats.from_mapping(ckpt["latent_stats"])
        except (KeyError, ValueError) as exc:
            raise CheckpointError(f"{path}: missing or bad latent_stats: {exc}") from exc
        inherited = ckpt.get("latent_stats_inherited", False) is True
        # Same data (fingerprint above), but the stats are recomputed: allow float noise
        # from another numpy build / CPU, then use the checkpoint's exact values. Stats a
        # fine-tune inherited (init_from) belong to the parent run's data, not to this one.
        if not inherited and not saved_stats.close_to(self.latent_stats):
            raise ResumeError("latent normalisation stats differ from the checkpoint's")
        self.latent_stats = saved_stats
        self.latent_stats_inherited = inherited
        init_from = ckpt.get("init_from")
        self.init_from_path = init_from if isinstance(init_from, str) else None
        saved_diff = ckpt.get("init_from_model_diff")
        if saved_diff is not None and saved_diff != _allowed_diff_record():
            raise CheckpointError(f"{path}: unexpected init_from_model_diff {saved_diff!r}")
        self.init_from_model_diff = cast("dict[str, list[str]] | None", saved_diff)
        try:
            step = cast("int", ckpt["step"])
            sampler = cast("Mapping[str, int]", ckpt["sampler"])
            rng = cast("Mapping[str, object]", ckpt["rng"])
            self.model.load_state_dict(cast("dict[str, Tensor]", ckpt["model"]))
            self.optimizer.load_state_dict(cast("dict[str, object]", ckpt["optimizer"]))
            self.scheduler.load_state_dict(cast("dict[str, object]", ckpt["scheduler"]))
            self.ema.load_state_dict(cast("dict[str, object]", ckpt["ema"]))
            self.stream.set_position(SamplerPosition(sampler["epoch"], sampler["batch_index"]))
            restore_rng(cast("Mapping[str, object]", rng["global"]))
            self.data_rng.bit_generator.state = cast("dict[str, object]", rng["data"])
            self.flow_gen.set_state(cast("Tensor", rng["flow"]))
            best = ckpt["best_val_loss"]
            self.best_val_loss = None if best is None else float(cast("float", best))
        except (KeyError, TypeError) as exc:
            raise CheckpointError(f"{path}: incomplete checkpoint: {exc!r}") from exc
        if step > self.cfg.steps:
            raise ResumeError(f"checkpoint step {step} is beyond cfg.steps={self.cfg.steps}")
        try:
            # Checkpoints written before the counters existed start from zero.
            totals = ckpt.get("pairing")
            self.pairing_totals = (
                PairingCounts() if totals is None else PairingCounts.from_mapping(totals)
            )
        except ValueError as exc:
            raise CheckpointError(f"{path}: bad pairing counters: {exc}") from exc
        self.step = step
        self.resumed_from_step = step
        self._last_saved_step = step
        log.info("resumed from %s at step %d", path, step)

    def init_from(self, path: Path) -> None:
        """Start a new run (fine-tune) from another run's weights.

        Model weights, EMA weights (with their update count, so the EMA keeps its decay)
        and the latent normalisation stats come from ``path``; the optimizer, scheduler,
        sampler, RNG streams and step counter stay fresh, so ``cfg``'s own schedule (warmup,
        cosine) runs from step 0 on this run's data. The stats are inherited because the
        model's input/output space is defined by them; recomputing them on new data would
        shift every latent the model sees.

        The model configs (compared after resolving the overrides, so an explicit default
        equals an absent one) must be identical, with exactly one exception: parent
        ``dur_cond="none"`` -> this run ``dur_cond="emotion"`` (E13). Then the checkpoint
        must hold every other key and lack exactly the new ``dur_film.*`` parameters, which
        keep their (zero-FiLM) initialisation in the model and in the EMA;
        ``init_from_model_diff`` records the difference.

        Raises:
            ResumeError: the checkpoint's model config differs from ``cfg.model`` (other
                than the allowed ``dur_cond`` step), or this trainer has already taken a
                step, or ``cfg.overfit_one_batch`` is set.
            CheckpointError: unreadable or incomplete checkpoint, or (``dur_cond`` step) a
                missing or unexpected key besides the new parameters.
        """
        if self.step != 0:
            raise ResumeError(f"init_from needs a fresh trainer (at step {self.step})")
        if self.cfg.overfit_one_batch:
            # The fixed batch was normalised with this run's own stats in __init__.
            raise ResumeError("init_from cannot be combined with overfit_one_batch")
        ckpt = load_checkpoint(path)
        try:
            saved_cfg = TrainConfig.from_mapping(cast("Mapping[str, object]", ckpt["config"]))
        except (KeyError, ConfigError) as exc:
            raise CheckpointError(f"{path}: bad config in checkpoint: {exc}") from exc
        diff = model_config_diff(saved_cfg.model_config(), self.model_cfg)
        if diff and diff != _ALLOWED_INIT_DIFF:
            raise ResumeError(
                f"{path}: model config differs from this run's (fields {sorted(diff)}); "
                "cannot init_from"
            )
        try:
            stats = LatentStats.from_mapping(ckpt["latent_stats"])
        except (KeyError, ValueError) as exc:
            raise CheckpointError(f"{path}: missing or bad latent_stats: {exc}") from exc
        # Parameters the parent cannot have (empty for identical configs).
        added_model = {k for k in self.model.state_dict() if k.startswith(DUR_FILM_PREFIX)}
        added_ema = {k for k in self.ema.shadow if k.startswith(DUR_FILM_PREFIX)}
        if not diff:
            added_model, added_ema = set(), set()
        try:
            model_state = cast("dict[str, Tensor]", ckpt["model"])
            if diff:
                _check_added_keys(path, set(model_state), set(self.model.state_dict()), added_model)
            self.model.load_state_dict(model_state, strict=not diff)
            self.ema.load_state_dict(
                cast("dict[str, object]", ckpt["ema"]), allow_missing=added_ema
            )
        except (KeyError, TypeError, ValueError, RuntimeError) as exc:
            if isinstance(exc, CheckpointError):
                raise
            # The trainer may be half-loaded now; the caller aborts the run.
            raise CheckpointError(f"{path}: incomplete or mismatched checkpoint: {exc!r}") from exc
        self.latent_stats = stats
        self.latent_stats_inherited = True
        self.init_from_path = str(path.resolve())
        self.init_from_model_diff = _allowed_diff_record() if diff else None
        log.info("initialised weights from %s (step %s)", path, ckpt.get("step"))

    def resume_latest(self) -> Path | None:
        """Resume from the newest checkpoint in the workspace, if any."""
        path = latest_checkpoint(self.checkpoint_dir)
        if path is not None:
            self.resume(path)
        return path

    # -- the loop -------------------------------------------------------------------------

    def run(self, *, handle_sigterm: bool = False) -> TrainResult:
        """Train until ``cfg.steps`` or a stop request; always leaves a checkpoint behind.

        A stop request (SIGTERM with ``handle_sigterm``, or ``stop.set()``) is honoured at
        the next step boundary: the current step finishes, a checkpoint is written, and the
        result has ``stopped_by_signal=True``.
        """
        cfg = self.cfg
        removed = truncate_jsonl_after_step(self.metrics_path, self.step)
        if removed:
            log.warning("removed %d replayed/torn log lines after step %d", removed, self.step)
        result = TrainResult(
            steps_done=self.step,
            completed=False,
            stopped_by_signal=False,
            last_loss=None,
            mean_frames_per_s=None,
            resumed_from_step=self.resumed_from_step,
        )
        fps: list[float] = []
        signals = sigterm_handler(self.stop) if handle_sigterm else contextlib.nullcontext()
        with signals, JsonlLog(self.metrics_path) as metrics_log:
            while self.step < cfg.steps and not self.stop.is_set():
                metrics = self.train_step()
                result.last_loss = metrics.loss
                fps.append(metrics.frames_per_s)
                if metrics.step % cfg.log_every == 0 or metrics.step == cfg.steps:
                    metrics_log.write(metrics.to_record())
                    atomic_write_json(
                        self.out_dir / HEARTBEAT_NAME,
                        {
                            "step": metrics.step,
                            "total_steps": cfg.steps,
                            "loss": metrics.loss,
                            "best_val_loss": self.best_val_loss,
                            "time_unix": time.time(),
                        },
                    )
                    log.info(
                        "step %d loss %.5f lr %.3e grad_norm %.3f %.0f frames/s",
                        metrics.step,
                        metrics.loss,
                        metrics.lr,
                        metrics.grad_norm,
                        metrics.frames_per_s,
                    )
                if cfg.eval_every is not None and metrics.step % cfg.eval_every == 0:
                    self._log_eval(metrics_log, result)
                if metrics.step % cfg.save_every == 0:
                    result.checkpoints.append(str(self.save()))
            result.completed = self.step >= cfg.steps
            result.stopped_by_signal = self.stop.is_set() and not result.completed
            if (
                result.completed
                and cfg.eval_every is not None
                and not (result.evals and result.evals[-1]["step"] == float(self.step))
            ):
                self._log_eval(metrics_log, result)
            if self._last_saved_step != self.step:
                result.checkpoints.append(str(self.save()))
        result.steps_done = self.step
        result.mean_frames_per_s = float(np.mean(fps)) if fps else None
        return result

    def _log_eval(self, metrics_log: JsonlLog, result: TrainResult) -> None:
        evaluation = self.evaluate()
        if evaluation is None:
            return
        result.evals.append(evaluation)
        metrics_log.write({"kind": "eval", **evaluation, "step": self.step})
        current = evaluation["val_loss_ema"]
        if self.best_val_loss is None or current < self.best_val_loss:
            self.best_val_loss = current
            path = save_checkpoint(self.out_dir / BEST_DIR, self.step, self.state_dict(), 1)
            result.best_checkpoint = str(path)
            log.info("new best val_loss_ema %.5f -> %s", current, path)
        log.info(
            "eval step %d val_loss %.5f val_loss_ema %.5f (no text %.5f)",
            self.step,
            evaluation["val_loss"],
            evaluation["val_loss_ema"],
            evaluation["val_loss_ema_notext"],
        )
        if "val_loss_ema_cross" in evaluation:
            log.info(
                "eval step %d cross-segment prompt: val_loss %.5f val_loss_ema %.5f (%d items)",
                self.step,
                evaluation["val_loss_cross"],
                evaluation["val_loss_ema_cross"],
                int(evaluation["val_cross_items"]),
            )


def _seed_int(seq: np.random.SeedSequence) -> int:
    """A 64-bit integer seed derived from a ``SeedSequence`` (for ``torch.Generator``)."""
    return int(seq.generate_state(1, dtype=np.uint64)[0])