774 lines
31 KiB
Python
774 lines
31 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import yaml
|
|
from accelerate import Accelerator
|
|
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
|
|
from torch.optim import AdamW
|
|
from transformers import get_cosine_schedule_with_warmup
|
|
|
|
from dots_tts.config import app as app_config
|
|
from dots_tts.data import builders as data_module
|
|
from dots_tts.models.dots_tts import model as dots_tts_model
|
|
from dots_tts.training import checkpoint as train_checkpoint
|
|
from dots_tts.training import losses as loss_ops
|
|
from dots_tts.training import utils as train_utils
|
|
from dots_tts.utils import util as util_module
|
|
|
|
_EMPTY_EPOCH_TOLERANCE = 32
|
|
_DEBUG_BATCH_LIMIT = 3
|
|
_DEBUG_GRAD_EARLY_STEP_LIMIT = 3
|
|
|
|
|
|
# region Training Step State
|
|
@dataclass(slots=True)
|
|
class _PreparedTrainingStep:
|
|
micro_batches: list[dict]
|
|
consumed_counts: list[int]
|
|
global_denominators: dict[str, float]
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _AccumulatedTrainingStep:
|
|
loss_totals: dict[str, float]
|
|
loss_denominators: dict[str, float]
|
|
source_loss_totals: dict[str, dict[str, float]]
|
|
source_loss_denominators: dict[str, dict[str, float]]
|
|
completed_optimizer_step: bool
|
|
grad_norm: torch.Tensor | None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _CompletedTrainingStep:
|
|
reduced_metrics: dict[str, float]
|
|
learning_rate: float
|
|
grad_norm_value: float
|
|
|
|
# endregion Training Step State
|
|
|
|
class DotsTtsTrainingRun:
|
|
# region Lifecycle
|
|
def __init__(self, cfg: app_config.AppConfig, *, debug_enabled: bool = False):
|
|
self.cfg = cfg
|
|
self.progress = train_utils.TrainProgress()
|
|
self.max_train_steps = int(cfg.train.max_train_steps)
|
|
self.grad_accumulation_steps = int(cfg.train.gradient_accumulation_steps)
|
|
self.last_validation_step: int | None = None
|
|
self.consecutive_empty_epochs = 0
|
|
self.saved_latest_checkpoint = False
|
|
self._last_log_step = 0
|
|
self._last_log_time = 0.0
|
|
self._debug_enabled = bool(debug_enabled)
|
|
self._debug_batch_count = 0
|
|
self._debug_audio_sample_rate = int(self.cfg.train_data.train_audio_sample_rate)
|
|
|
|
project_config = ProjectConfiguration(
|
|
project_dir=self.cfg.train.output_dir,
|
|
total_limit=self.cfg.train.max_checkpoints_to_keep,
|
|
)
|
|
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
|
|
self.accelerator = Accelerator(
|
|
kwargs_handlers=[ddp_kwargs],
|
|
gradient_accumulation_steps=self.grad_accumulation_steps,
|
|
log_with="tensorboard",
|
|
project_config=project_config,
|
|
step_scheduler_with_optimizer=False,
|
|
)
|
|
|
|
util_module.seed_everything(self.cfg.train.seed)
|
|
|
|
model = dots_tts_model.DotsTtsModel.from_pretrained(
|
|
self.cfg.train.pretrained_model_path
|
|
)
|
|
# model.set_cfg_droprate(
|
|
# cfg_droprate=self.cfg.train.cfg_droprate,
|
|
# xvec_drop_rate=self.cfg.train.xvec_drop_rate,
|
|
# )
|
|
optimizer = AdamW(
|
|
(param for param in model.parameters() if param.requires_grad),
|
|
lr=self.cfg.train.learning_rate,
|
|
weight_decay=self.cfg.train.weight_decay,
|
|
)
|
|
scheduler = get_cosine_schedule_with_warmup(
|
|
optimizer,
|
|
num_warmup_steps=self.cfg.train.warmup_steps,
|
|
num_training_steps=self.max_train_steps,
|
|
)
|
|
self.model, self.optimizer, self.scheduler = self.accelerator.prepare(
|
|
model,
|
|
optimizer,
|
|
scheduler,
|
|
)
|
|
self.unwrapped_model = self.accelerator.unwrap_model(self.model)
|
|
expected_sample_rate = int(self.unwrapped_model.config.vocoder.sample_rate)
|
|
expected_audio_samples_per_llm_token = (
|
|
int(self.unwrapped_model.hop_size) * int(self.unwrapped_model.config.patch_size)
|
|
)
|
|
if int(self.cfg.train_data.train_audio_sample_rate) != expected_sample_rate:
|
|
raise ValueError(
|
|
f"train_data.train_audio_sample_rate={int(self.cfg.train_data.train_audio_sample_rate)} "
|
|
f"does not match the pretrained model sample rate {expected_sample_rate}."
|
|
)
|
|
if (
|
|
int(self.cfg.train_data.audio_samples_per_llm_token)
|
|
!= expected_audio_samples_per_llm_token
|
|
):
|
|
raise ValueError(
|
|
"train_data.audio_samples_per_llm_token="
|
|
f"{int(self.cfg.train_data.audio_samples_per_llm_token)} "
|
|
"does not match the pretrained model audio token contract "
|
|
f"{expected_audio_samples_per_llm_token}."
|
|
)
|
|
if self.cfg.val_data is not None:
|
|
if int(self.cfg.val_data.train_audio_sample_rate) != expected_sample_rate:
|
|
raise ValueError(
|
|
f"val_data.train_audio_sample_rate={int(self.cfg.val_data.train_audio_sample_rate)} "
|
|
f"does not match the pretrained model sample rate {expected_sample_rate}."
|
|
)
|
|
if (
|
|
int(self.cfg.val_data.audio_samples_per_llm_token)
|
|
!= expected_audio_samples_per_llm_token
|
|
):
|
|
raise ValueError(
|
|
"val_data.audio_samples_per_llm_token="
|
|
f"{int(self.cfg.val_data.audio_samples_per_llm_token)} "
|
|
"does not match the pretrained model audio token contract "
|
|
f"{expected_audio_samples_per_llm_token}."
|
|
)
|
|
|
|
if self.accelerator.is_main_process:
|
|
total_params = sum(param.numel() for param in self.unwrapped_model.parameters())
|
|
trainable_params = sum(
|
|
param.numel()
|
|
for param in self.unwrapped_model.parameters()
|
|
if param.requires_grad
|
|
)
|
|
self.accelerator.print(f"Total parameters: {total_params:,}")
|
|
self.accelerator.print(f"Trainable parameters: {trainable_params:,}")
|
|
self.accelerator.print(
|
|
f"Distributed type: {self.accelerator.distributed_type}"
|
|
)
|
|
|
|
tokenizer = self.unwrapped_model.tokenizer
|
|
self.tokenizer = tokenizer
|
|
train_dataset = data_module.build_training_dataset(
|
|
self.cfg.train_data,
|
|
tokenizer=tokenizer,
|
|
seed=int(self.cfg.train.seed),
|
|
accelerator=self.accelerator,
|
|
)
|
|
self.train_loader = data_module.build_training_dataloader(
|
|
train_dataset,
|
|
self.cfg.train_data,
|
|
tokenizer=tokenizer,
|
|
)
|
|
|
|
self.val_loader = None
|
|
if (
|
|
self.cfg.train.eval_interval is not None
|
|
or self.cfg.train.run_eval_on_start
|
|
):
|
|
if self.cfg.val_data is None:
|
|
raise ValueError(
|
|
"Validation requires val_data when eval_interval or "
|
|
"run_eval_on_start is enabled."
|
|
)
|
|
validation_data_cfg = self.cfg.val_data.model_copy(deep=True)
|
|
validation_data_cfg.num_tokens_per_epoch = None
|
|
val_dataset = data_module.build_validation_dataset(
|
|
validation_data_cfg,
|
|
tokenizer=tokenizer,
|
|
seed=int(self.cfg.train.seed),
|
|
accelerator=self.accelerator,
|
|
)
|
|
self.val_loader = data_module.build_validation_dataloader(
|
|
val_dataset,
|
|
validation_data_cfg,
|
|
tokenizer=tokenizer,
|
|
)
|
|
|
|
self._resume_if_available()
|
|
self.train_loader.set_epoch(self.progress.epoch)
|
|
|
|
def run(self) -> int:
|
|
self.accelerator.init_trackers("dots_tts")
|
|
self._write_run_config()
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
|
|
try:
|
|
if self.cfg.train.run_eval_on_start:
|
|
self._run_validation()
|
|
self.last_validation_step = self.progress.global_step
|
|
|
|
self._last_log_step = self.progress.global_step
|
|
self._last_log_time = time.perf_counter()
|
|
|
|
while self.progress.global_step < self.max_train_steps:
|
|
self._run_training_step()
|
|
|
|
if (
|
|
self.cfg.train.eval_interval is not None
|
|
and self.val_loader is not None
|
|
and self.progress.global_step > 0
|
|
and self.last_validation_step != self.progress.global_step
|
|
):
|
|
self._run_validation()
|
|
|
|
if not self.saved_latest_checkpoint:
|
|
self._save_checkpoint(float(self.optimizer.param_groups[0]["lr"]))
|
|
return 0
|
|
finally:
|
|
try:
|
|
self._close_data_streams()
|
|
finally:
|
|
self.accelerator.end_training()
|
|
|
|
def _write_run_config(self) -> None:
|
|
if not bool(getattr(self.accelerator, "is_main_process", True)):
|
|
return
|
|
output_dir = Path(self.cfg.train.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
config_path = output_dir / "config.yml"
|
|
with config_path.open("w", encoding="utf-8") as fout:
|
|
yaml.safe_dump(
|
|
self.cfg.to_dict(),
|
|
fout,
|
|
sort_keys=False,
|
|
allow_unicode=True,
|
|
)
|
|
|
|
def _close_data_streams(self) -> None:
|
|
for loader_name in ("train_loader", "val_loader"):
|
|
loader = getattr(self, loader_name, None)
|
|
close = getattr(loader, "close", None)
|
|
if callable(close):
|
|
close()
|
|
setattr(self, loader_name, None)
|
|
|
|
def _resume_if_available(self) -> None:
|
|
try:
|
|
resume_dir = train_checkpoint.resolve_latest_train_checkpoint(
|
|
self.cfg.train.output_dir
|
|
)
|
|
except FileNotFoundError:
|
|
return
|
|
|
|
resume_state = train_checkpoint.load_train_checkpoint(
|
|
self.accelerator,
|
|
self.model,
|
|
self.optimizer,
|
|
self.progress,
|
|
resume_dir,
|
|
self.scheduler,
|
|
)
|
|
saved_max_train_steps = int(resume_state["scheduler_state"]["max_train_steps"])
|
|
if saved_max_train_steps != self.max_train_steps:
|
|
self.accelerator.print(
|
|
"Warning: resumed scheduler was saved with "
|
|
f"max_train_steps={saved_max_train_steps}, but current run uses "
|
|
f"{self.max_train_steps}."
|
|
)
|
|
|
|
self.train_loader.load_state_dict(resume_state["data_state"])
|
|
self.accelerator.print(
|
|
"Resumed training from "
|
|
f"{resume_dir} at step {self.progress.global_step}. "
|
|
"Restored committed data state. "
|
|
"In-memory prefetch and batching state is rebuilt on restart, so only "
|
|
"committed sample progress is resumed."
|
|
)
|
|
# endregion Lifecycle
|
|
|
|
# region Training Step Pipeline
|
|
def _run_training_step(self) -> None:
|
|
try:
|
|
self.model.train()
|
|
# Stage 1: collect one synchronized accumulation window and its
|
|
# normalization factors before touching model state.
|
|
prepared_step = self._prepare_training_step()
|
|
|
|
# Stage 2: run forward/backward over the prepared micro-batches and
|
|
# accumulate overall/source statistics for the completed optimizer step.
|
|
accumulated_step = self._accumulate_training_step(prepared_step)
|
|
|
|
# Stage 3: advance counters, reduce metrics, then trigger side effects
|
|
# (logging, validation, checkpointing) only after a real optimizer step.
|
|
self._apply_consumed_counts(prepared_step.consumed_counts)
|
|
if not accumulated_step.completed_optimizer_step:
|
|
return
|
|
completed_step = self._finalize_completed_training_step(accumulated_step)
|
|
if train_utils.should_log_training_step(
|
|
self.progress.global_step,
|
|
int(self.cfg.train.log_interval),
|
|
):
|
|
reduced_by_source = train_utils.reduce_source_metrics(
|
|
accumulated_step.source_loss_totals,
|
|
accumulated_step.source_loss_denominators,
|
|
device=self.accelerator.device,
|
|
loss_config=self.cfg.loss,
|
|
)
|
|
current_time = time.perf_counter()
|
|
report = train_utils.build_train_step_report(
|
|
completed_step.reduced_metrics,
|
|
learning_rate=completed_step.learning_rate,
|
|
grad_norm=completed_step.grad_norm_value,
|
|
current_time=current_time,
|
|
last_log_step=self._last_log_step,
|
|
last_log_time=self._last_log_time,
|
|
progress=self.progress,
|
|
max_train_steps=self.max_train_steps,
|
|
reduced_by_source=reduced_by_source,
|
|
)
|
|
self.accelerator.log(
|
|
report.log_values,
|
|
step=self.progress.global_step,
|
|
)
|
|
self.accelerator.print(report.console_line)
|
|
self._last_log_step = self.progress.global_step
|
|
self._last_log_time = current_time
|
|
|
|
if (
|
|
self.cfg.train.eval_interval is not None
|
|
and self.progress.global_step % self.cfg.train.eval_interval == 0
|
|
):
|
|
self._run_validation()
|
|
self.last_validation_step = self.progress.global_step
|
|
|
|
if self.progress.global_step % self.cfg.train.save_interval == 0:
|
|
self._save_checkpoint(completed_step.learning_rate)
|
|
self.saved_latest_checkpoint = True
|
|
except BaseException as exc:
|
|
train_utils.abort_on_out_of_memory(
|
|
exc,
|
|
stage="train",
|
|
batch=None,
|
|
progress=self.progress,
|
|
device=self.accelerator.device,
|
|
process_index=int(getattr(self.accelerator, "process_index", 0)),
|
|
num_processes=int(getattr(self.accelerator, "num_processes", 1)),
|
|
)
|
|
raise
|
|
|
|
def _prepare_training_step(self) -> _PreparedTrainingStep:
|
|
micro_batches: list[dict] = []
|
|
local_denominators: dict[str, float] = {}
|
|
|
|
while len(micro_batches) < self.grad_accumulation_steps:
|
|
batch, has_batch = self.train_loader.peek_batch()
|
|
if train_utils.any_rank_true(not has_batch, device=self.accelerator.device):
|
|
self._advance_epoch_after_empty_batch(has_local_batch=has_batch)
|
|
continue
|
|
|
|
self.consecutive_empty_epochs = 0
|
|
self.train_loader.commit_batch()
|
|
prepared_batch = self.unwrapped_model.prepare_training_batch(batch)
|
|
self._maybe_debug_training_batch(prepared_batch)
|
|
batch_denominators = loss_ops.to_host_named_scalars(
|
|
loss_ops.collapse_loss_masks(prepared_batch["loss_masks"])
|
|
)
|
|
if not local_denominators:
|
|
local_denominators = {name: 0.0 for name in batch_denominators}
|
|
loss_ops.accumulate_named_scalars_(local_denominators, batch_denominators)
|
|
micro_batches.append(prepared_batch)
|
|
|
|
consumed_counts = train_utils.sum_integer_counters_across_ranks(
|
|
[
|
|
sum(int(batch["input_ids_lengths"].sum().item()) for batch in micro_batches),
|
|
sum(int(batch["num_audio_tokens"].sum().item()) for batch in micro_batches),
|
|
sum(int(batch["num_text_tokens"].sum().item()) for batch in micro_batches),
|
|
],
|
|
device=self.accelerator.device,
|
|
)
|
|
global_denominators = loss_ops.sum_named_scalars_across_ranks(
|
|
local_denominators,
|
|
device=self.accelerator.device,
|
|
)
|
|
return _PreparedTrainingStep(
|
|
micro_batches=micro_batches,
|
|
consumed_counts=consumed_counts,
|
|
global_denominators=global_denominators,
|
|
)
|
|
|
|
def _advance_epoch_after_empty_batch(self, *, has_local_batch: bool) -> None:
|
|
if has_local_batch:
|
|
self.train_loader.discard_batch()
|
|
self.progress.epoch += 1
|
|
self.train_loader.set_epoch(self.progress.epoch)
|
|
self.consecutive_empty_epochs += 1
|
|
if self.consecutive_empty_epochs > _EMPTY_EPOCH_TOLERANCE:
|
|
raise RuntimeError(
|
|
"Unable to obtain a synchronized training batch across ranks. "
|
|
"Check shard assignment, dataset size, and filtering constraints."
|
|
)
|
|
|
|
def _accumulate_training_step(
|
|
self,
|
|
prepared_step: _PreparedTrainingStep,
|
|
) -> _AccumulatedTrainingStep:
|
|
accumulated_loss_totals: dict[str, float] = {}
|
|
accumulated_loss_denominators: dict[str, float] = {}
|
|
accumulated_source_loss_totals: dict[str, dict[str, float]] = {}
|
|
accumulated_source_loss_denominators: dict[str, dict[str, float]] = {}
|
|
completed_optimizer_step = False
|
|
grad_norm = None
|
|
|
|
for batch in prepared_step.micro_batches:
|
|
batch = train_utils.move_to_device(batch, self.accelerator.device)
|
|
with self.accelerator.accumulate(self.model):
|
|
with self.accelerator.autocast():
|
|
loss_terms = self.model(batch)
|
|
loss = loss_ops.compute_gradient_loss(
|
|
loss_terms,
|
|
global_normalizers=prepared_step.global_denominators,
|
|
loss_config=self.cfg.loss,
|
|
ddp_world_size=int(self.accelerator.num_processes),
|
|
gradient_accumulation_steps=self.grad_accumulation_steps,
|
|
)
|
|
|
|
batch_loss_totals, batch_loss_denominators = (
|
|
loss_ops.collapse_loss_terms(loss_terms)
|
|
)
|
|
batch_loss_totals = loss_ops.to_host_named_scalars(batch_loss_totals)
|
|
batch_loss_denominators = loss_ops.to_host_named_scalars(
|
|
batch_loss_denominators
|
|
)
|
|
if not accumulated_loss_totals:
|
|
accumulated_loss_totals = {name: 0.0 for name in batch_loss_totals}
|
|
accumulated_loss_denominators = {
|
|
name: 0.0 for name in batch_loss_denominators
|
|
}
|
|
loss_ops.accumulate_named_scalars_(
|
|
accumulated_loss_totals,
|
|
batch_loss_totals,
|
|
)
|
|
loss_ops.accumulate_named_scalars_(
|
|
accumulated_loss_denominators,
|
|
batch_loss_denominators,
|
|
)
|
|
|
|
batch_source_totals, batch_source_denominators = (
|
|
loss_ops.collapse_loss_terms_by_source(
|
|
loss_terms,
|
|
source_names=batch["source_names"],
|
|
)
|
|
)
|
|
loss_ops.accumulate_grouped_named_scalars_(
|
|
accumulated_source_loss_totals,
|
|
batch_source_totals,
|
|
)
|
|
loss_ops.accumulate_grouped_named_scalars_(
|
|
accumulated_source_loss_denominators,
|
|
batch_source_denominators,
|
|
)
|
|
|
|
self.accelerator.backward(loss)
|
|
if self.accelerator.sync_gradients:
|
|
grad_norm = self.accelerator.clip_grad_norm_(
|
|
self.model.parameters(),
|
|
self.cfg.train.grad_clip_norm,
|
|
)
|
|
self._maybe_print_gradient_debug(grad_norm)
|
|
self.optimizer.step()
|
|
completed_optimizer_step = (
|
|
not self.accelerator.optimizer_step_was_skipped
|
|
)
|
|
if completed_optimizer_step:
|
|
self.scheduler.step()
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
batch.clear()
|
|
|
|
return _AccumulatedTrainingStep(
|
|
loss_totals=accumulated_loss_totals,
|
|
loss_denominators=accumulated_loss_denominators,
|
|
source_loss_totals=accumulated_source_loss_totals,
|
|
source_loss_denominators=accumulated_source_loss_denominators,
|
|
completed_optimizer_step=completed_optimizer_step,
|
|
grad_norm=grad_norm,
|
|
)
|
|
|
|
def _apply_consumed_counts(self, consumed_counts: list[int]) -> None:
|
|
self.progress.total_tokens += consumed_counts[0]
|
|
self.progress.audio_tokens += consumed_counts[1]
|
|
self.progress.text_tokens += consumed_counts[2]
|
|
|
|
def _finalize_completed_training_step(
|
|
self,
|
|
accumulated_step: _AccumulatedTrainingStep,
|
|
) -> _CompletedTrainingStep:
|
|
if not accumulated_step.loss_totals or not accumulated_step.loss_denominators:
|
|
raise RuntimeError("Training step produced no accumulated loss totals.")
|
|
if all(
|
|
float(value) == 0.0 for value in accumulated_step.loss_denominators.values()
|
|
):
|
|
raise RuntimeError("Accumulated training step produced no loss statistics.")
|
|
|
|
self.progress.global_step += 1
|
|
self.saved_latest_checkpoint = False
|
|
|
|
reduced_totals = loss_ops.sum_named_scalars_across_ranks(
|
|
accumulated_step.loss_totals,
|
|
device=self.accelerator.device,
|
|
)
|
|
reduced_denominators = loss_ops.sum_named_scalars_across_ranks(
|
|
accumulated_step.loss_denominators,
|
|
device=self.accelerator.device,
|
|
)
|
|
reduced_metrics = loss_ops.reduce_loss_statistics(
|
|
reduced_totals,
|
|
reduced_denominators,
|
|
loss_config=self.cfg.loss,
|
|
)
|
|
learning_rate = float(self.optimizer.param_groups[0]["lr"])
|
|
grad_norm_value = (
|
|
math.nan
|
|
if accumulated_step.grad_norm is None
|
|
else float(accumulated_step.grad_norm.detach().float().item())
|
|
)
|
|
return _CompletedTrainingStep(
|
|
reduced_metrics=reduced_metrics,
|
|
learning_rate=learning_rate,
|
|
grad_norm_value=grad_norm_value,
|
|
)
|
|
# endregion Training Step Pipeline
|
|
|
|
# region Validation
|
|
def _run_validation(self) -> None:
|
|
try:
|
|
if self.val_loader is None:
|
|
raise ValueError(
|
|
"Validation requested, but validation loader was not initialized."
|
|
)
|
|
self.val_loader.set_epoch(0)
|
|
|
|
was_training = bool(self.model.training)
|
|
self.model.eval()
|
|
|
|
overall_loss_totals = None
|
|
overall_loss_denominators = None
|
|
source_loss_totals: dict[str, dict[str, float]] = {}
|
|
source_loss_denominators: dict[str, dict[str, float]] = {}
|
|
processed_batches = 0
|
|
|
|
# Collect rank-local partial sums using the same batch preparation and
|
|
# loss aggregation path as training.
|
|
with torch.no_grad():
|
|
for batch_idx, batch in enumerate(self.val_loader):
|
|
if (
|
|
self.cfg.train.max_eval_batches is not None
|
|
and batch_idx >= self.cfg.train.max_eval_batches
|
|
):
|
|
break
|
|
|
|
batch = self.unwrapped_model.prepare_training_batch(batch)
|
|
batch = train_utils.move_to_device(batch, self.accelerator.device)
|
|
|
|
with self.accelerator.autocast():
|
|
loss_terms = self.model(batch)
|
|
|
|
batch_loss_totals, batch_loss_denominators = (
|
|
loss_ops.collapse_loss_terms(loss_terms)
|
|
)
|
|
batch_loss_totals = loss_ops.to_host_named_scalars(batch_loss_totals)
|
|
batch_loss_denominators = loss_ops.to_host_named_scalars(
|
|
batch_loss_denominators
|
|
)
|
|
if overall_loss_totals is None:
|
|
overall_loss_totals = {name: 0.0 for name in batch_loss_totals}
|
|
overall_loss_denominators = {
|
|
name: 0.0 for name in batch_loss_denominators
|
|
}
|
|
loss_ops.accumulate_named_scalars_(
|
|
overall_loss_totals,
|
|
batch_loss_totals,
|
|
)
|
|
loss_ops.accumulate_named_scalars_(
|
|
overall_loss_denominators,
|
|
batch_loss_denominators,
|
|
)
|
|
|
|
batch_source_totals, batch_source_denominators = (
|
|
loss_ops.collapse_loss_terms_by_source(
|
|
loss_terms,
|
|
source_names=batch["source_names"],
|
|
)
|
|
)
|
|
loss_ops.accumulate_grouped_named_scalars_(
|
|
source_loss_totals,
|
|
batch_source_totals,
|
|
)
|
|
loss_ops.accumulate_grouped_named_scalars_(
|
|
source_loss_denominators,
|
|
batch_source_denominators,
|
|
)
|
|
processed_batches += 1
|
|
|
|
# Merge rank-local partial sums with tensor reductions only. Validation
|
|
# runs close to the training memory ceiling, so object collectives are
|
|
# not acceptable here because NCCL materializes pickled payloads on GPU.
|
|
processed_batches = train_utils.sum_integer_counters_across_ranks(
|
|
[processed_batches],
|
|
device=self.accelerator.device,
|
|
)[0]
|
|
overall_loss_totals = loss_ops.sum_named_scalars_across_ranks(
|
|
overall_loss_totals or {},
|
|
device=self.accelerator.device,
|
|
)
|
|
overall_loss_denominators = loss_ops.sum_named_scalars_across_ranks(
|
|
overall_loss_denominators or {},
|
|
device=self.accelerator.device,
|
|
)
|
|
source_loss_totals = loss_ops.sum_grouped_named_scalars_across_ranks(
|
|
source_loss_totals,
|
|
device=self.accelerator.device,
|
|
)
|
|
source_loss_denominators = (
|
|
loss_ops.sum_grouped_named_scalars_across_ranks(
|
|
source_loss_denominators,
|
|
device=self.accelerator.device,
|
|
)
|
|
)
|
|
|
|
if processed_batches <= 0:
|
|
raise RuntimeError(
|
|
"Validation produced no batches. Check validation data configuration."
|
|
)
|
|
if not overall_loss_totals or not overall_loss_denominators:
|
|
raise RuntimeError("Validation produced no aggregate loss totals.")
|
|
|
|
reduced_metrics = loss_ops.reduce_loss_statistics(
|
|
overall_loss_totals,
|
|
overall_loss_denominators,
|
|
loss_config=self.cfg.loss,
|
|
)
|
|
reduced_by_source = loss_ops.reduce_loss_statistics_by_source(
|
|
source_loss_totals,
|
|
source_loss_denominators,
|
|
loss_config=self.cfg.loss,
|
|
)
|
|
|
|
if was_training:
|
|
self.model.train()
|
|
|
|
self.accelerator.log(
|
|
train_utils.build_validation_log_dict(
|
|
reduced_metrics,
|
|
reduced_by_source=reduced_by_source,
|
|
),
|
|
step=self.progress.global_step,
|
|
)
|
|
self.accelerator.print(
|
|
train_utils.format_validation_line(
|
|
reduced_metrics,
|
|
global_step=self.progress.global_step,
|
|
reduced_by_source=reduced_by_source,
|
|
)
|
|
)
|
|
except BaseException as exc:
|
|
train_utils.abort_on_out_of_memory(
|
|
exc,
|
|
stage="validation",
|
|
batch=None,
|
|
progress=self.progress,
|
|
device=self.accelerator.device,
|
|
process_index=int(getattr(self.accelerator, "process_index", 0)),
|
|
num_processes=int(getattr(self.accelerator, "num_processes", 1)),
|
|
)
|
|
raise
|
|
# endregion Validation
|
|
|
|
# region Checkpointing
|
|
def _save_checkpoint(self, learning_rate: float) -> None:
|
|
train_checkpoint.save_train_checkpoint(
|
|
self.accelerator,
|
|
self.model,
|
|
self.optimizer,
|
|
self.progress,
|
|
self.cfg.train.output_dir,
|
|
self.cfg.train.max_checkpoints_to_keep,
|
|
self.train_loader.state_dict(),
|
|
{
|
|
"type": "transformers_cosine_with_warmup",
|
|
"global_step": int(self.progress.global_step),
|
|
"base_lr": float(self.cfg.train.learning_rate),
|
|
"current_lr": float(learning_rate),
|
|
"warmup_steps": int(self.cfg.train.warmup_steps),
|
|
"max_train_steps": int(self.max_train_steps),
|
|
"state_dict": self.scheduler.state_dict(),
|
|
},
|
|
)
|
|
# endregion Checkpointing
|
|
|
|
# region Debug Logging
|
|
def _maybe_debug_training_batch(self, batch: dict[str, object]) -> None:
|
|
if not bool(getattr(self, "_debug_enabled", False)):
|
|
return
|
|
if not bool(getattr(self.accelerator, "is_main_process", True)):
|
|
return
|
|
if self._debug_batch_count >= _DEBUG_BATCH_LIMIT:
|
|
return
|
|
|
|
batch_index = self._debug_batch_count
|
|
self._debug_batch_count += 1
|
|
for line in train_utils.build_data_debug_lines(
|
|
batch,
|
|
batch_index=batch_index,
|
|
tokenizer=self.tokenizer,
|
|
sample_rate=self._debug_audio_sample_rate,
|
|
):
|
|
self.accelerator.print(line)
|
|
|
|
def _maybe_print_gradient_debug(self, grad_norm: torch.Tensor | None) -> None:
|
|
if grad_norm is None:
|
|
return
|
|
if not train_utils.should_print_gradient_debug(
|
|
debug_enabled=bool(getattr(self, "_debug_enabled", False)),
|
|
is_main_process=bool(getattr(self.accelerator, "is_main_process", True)),
|
|
next_global_step=self.progress.global_step + 1,
|
|
log_interval=int(self.cfg.train.log_interval),
|
|
early_step_limit=_DEBUG_GRAD_EARLY_STEP_LIMIT,
|
|
):
|
|
return
|
|
for line in train_utils.build_gradient_debug_lines(
|
|
self.unwrapped_model,
|
|
global_step=self.progress.global_step + 1,
|
|
grad_norm=float(grad_norm.detach().float().item()),
|
|
grad_clip_norm=float(self.cfg.train.grad_clip_norm),
|
|
):
|
|
self.accelerator.print(line)
|
|
# endregion Debug Logging
|
|
|
|
|
|
# region CLI
|
|
def parse_args(argv=None):
|
|
parser = argparse.ArgumentParser(
|
|
description="Accelerate training entrypoint for dots.tts."
|
|
)
|
|
parser.add_argument("--config", default=app_config.DEFAULT_CONFIG_PATH)
|
|
parser.add_argument(
|
|
"--debug",
|
|
action="store_true",
|
|
help="Print training debug information.",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv=None):
|
|
args = parse_args(argv)
|
|
return DotsTtsTrainingRun(
|
|
app_config.load_config(args.config),
|
|
debug_enabled=args.debug,
|
|
).run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|
|
# endregion CLI
|