"""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])
"""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.