651 lines
22 KiB
Python
651 lines
22 KiB
Python
"""Shared helpers for the dots_tts training entrypoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import os
|
|
import sys
|
|
import traceback
|
|
from collections import Counter
|
|
from dataclasses import dataclass, fields, is_dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from dots_tts.training import losses as loss_ops
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Training State
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TrainProgress:
|
|
"""Minimal progress counters that must survive checkpoint save/load."""
|
|
|
|
global_step: int = 0
|
|
epoch: int = 0
|
|
total_tokens: int = 0
|
|
audio_tokens: int = 0
|
|
text_tokens: int = 0
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TrainStepReport:
|
|
log_values: dict[str, float]
|
|
console_line: str
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Distributed Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def any_rank_true(flag: bool, *, device: torch.device) -> bool:
|
|
"""Return ``True`` if any distributed rank reports ``flag=True``."""
|
|
packed = torch.tensor(int(flag), device=device, dtype=torch.int32)
|
|
if dist.is_available() and dist.is_initialized():
|
|
dist.all_reduce(packed, op=dist.ReduceOp.MAX)
|
|
return bool(packed.item())
|
|
|
|
|
|
def sum_integer_counters_across_ranks(
|
|
values: list[int],
|
|
*,
|
|
device: torch.device,
|
|
) -> list[int]:
|
|
"""All-reduce integer counters and return their cross-rank sums."""
|
|
packed = torch.tensor(values, device=device, dtype=torch.int64)
|
|
if dist.is_available() and dist.is_initialized():
|
|
dist.all_reduce(packed, op=dist.ReduceOp.SUM)
|
|
return [int(value) for value in packed.tolist()]
|
|
|
|
|
|
def move_to_device(value, device):
|
|
"""Recursively move nested tensors/dataclasses onto ``device``."""
|
|
if isinstance(value, torch.Tensor):
|
|
return value.to(device, non_blocking=True)
|
|
if isinstance(value, dict):
|
|
return {key: move_to_device(item, device) for key, item in value.items()}
|
|
if isinstance(value, list):
|
|
return [move_to_device(item, device) for item in value]
|
|
if isinstance(value, tuple):
|
|
return tuple(move_to_device(item, device) for item in value)
|
|
if is_dataclass(value) and not isinstance(value, type):
|
|
return type(value)(
|
|
**{
|
|
field.name: move_to_device(getattr(value, field.name), device)
|
|
for field in fields(value)
|
|
}
|
|
)
|
|
return value
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Failure Handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def abort_on_out_of_memory(
|
|
exc: BaseException,
|
|
*,
|
|
stage: str,
|
|
batch: dict[str, object] | None,
|
|
progress: TrainProgress,
|
|
device: torch.device,
|
|
process_index: int,
|
|
num_processes: int,
|
|
) -> None:
|
|
if not _is_out_of_memory_error(exc):
|
|
return
|
|
|
|
message = (
|
|
"Fatal out-of-memory during training. "
|
|
f"stage={stage}, "
|
|
f"epoch={progress.epoch}, "
|
|
f"global_step={progress.global_step}, "
|
|
f"rank={process_index}/{num_processes}. "
|
|
f"{_build_batch_memory_summary(batch)}. "
|
|
f"{_build_cuda_memory_summary(device)}."
|
|
)
|
|
print(message, file=sys.stderr, flush=True)
|
|
traceback.print_exception(type(exc), exc, exc.__traceback__, file=sys.stderr)
|
|
sys.stderr.flush()
|
|
|
|
if num_processes > 1:
|
|
os._exit(1)
|
|
|
|
|
|
def _is_out_of_memory_error(exc: BaseException) -> bool:
|
|
oom_error_type = getattr(torch, "OutOfMemoryError", None)
|
|
if oom_error_type is not None and isinstance(exc, oom_error_type):
|
|
return True
|
|
if not isinstance(exc, RuntimeError):
|
|
return False
|
|
return "out of memory" in str(exc).lower()
|
|
|
|
|
|
def _build_batch_memory_summary(batch: dict[str, object] | None) -> str:
|
|
if not isinstance(batch, dict):
|
|
return "batch=unavailable"
|
|
|
|
fields = []
|
|
input_ids = batch.get("input_ids")
|
|
if isinstance(input_ids, torch.Tensor):
|
|
fields.append(f"input_ids_shape={tuple(input_ids.shape)}")
|
|
sample = batch.get("sample")
|
|
if isinstance(sample, torch.Tensor):
|
|
fields.append(f"sample_shape={tuple(sample.shape)}")
|
|
input_ids_lengths = batch.get("input_ids_lengths")
|
|
if isinstance(input_ids_lengths, torch.Tensor) and input_ids_lengths.numel() > 0:
|
|
fields.append(
|
|
f"max_input_ids_length={int(input_ids_lengths.max().detach().item())}"
|
|
)
|
|
num_audio_tokens = batch.get("num_audio_tokens")
|
|
if isinstance(num_audio_tokens, torch.Tensor) and num_audio_tokens.numel() > 0:
|
|
fields.append(f"max_audio_tokens={int(num_audio_tokens.max().detach().item())}")
|
|
num_text_tokens = batch.get("num_text_tokens")
|
|
if isinstance(num_text_tokens, torch.Tensor) and num_text_tokens.numel() > 0:
|
|
fields.append(f"max_text_tokens={int(num_text_tokens.max().detach().item())}")
|
|
return ", ".join(fields) if fields else "batch=unavailable"
|
|
|
|
|
|
def _build_cuda_memory_summary(device: torch.device) -> str:
|
|
if device.type != "cuda" or not torch.cuda.is_available():
|
|
return "device_memory=unavailable"
|
|
allocated = torch.cuda.memory_allocated(device) / (1024**3)
|
|
reserved = torch.cuda.memory_reserved(device) / (1024**3)
|
|
max_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
|
|
max_reserved = torch.cuda.max_memory_reserved(device) / (1024**3)
|
|
return (
|
|
f"device={device}, "
|
|
f"allocated_gb={allocated:.2f}, "
|
|
f"reserved_gb={reserved:.2f}, "
|
|
f"max_allocated_gb={max_allocated:.2f}, "
|
|
f"max_reserved_gb={max_reserved:.2f}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Debug Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def build_data_debug_lines(
|
|
batch: dict[str, object],
|
|
*,
|
|
batch_index: int,
|
|
tokenizer: Any,
|
|
sample_rate: int,
|
|
) -> list[str]:
|
|
input_ids = batch["input_ids"]
|
|
input_ids_lengths = batch["input_ids_lengths"]
|
|
sample = batch["sample"]
|
|
sample_lengths = batch["sample_lengths"]
|
|
num_audio_tokens = batch["num_audio_tokens"]
|
|
num_text_tokens = batch["num_text_tokens"]
|
|
|
|
if not isinstance(input_ids, torch.Tensor) or not isinstance(
|
|
input_ids_lengths, torch.Tensor
|
|
):
|
|
raise TypeError("Debug batch requires tensor input_ids and input_ids_lengths.")
|
|
if not isinstance(sample, torch.Tensor) or not isinstance(
|
|
sample_lengths, torch.Tensor
|
|
):
|
|
raise TypeError("Debug batch requires tensor sample and sample_lengths.")
|
|
if not isinstance(num_audio_tokens, torch.Tensor) or not isinstance(
|
|
num_text_tokens, torch.Tensor
|
|
):
|
|
raise TypeError(
|
|
"Debug batch requires tensor num_audio_tokens and num_text_tokens."
|
|
)
|
|
|
|
source_names = batch.get("source_names")
|
|
debug_lines = [
|
|
(
|
|
"[debug:data] "
|
|
f"batch_index={batch_index} "
|
|
f"batch_size={int(input_ids.size(0))} "
|
|
f"input_ids_shape={tuple(input_ids.shape)} "
|
|
f"sample_shape={tuple(sample.shape)} "
|
|
f"sample_rate={sample_rate} "
|
|
f"sources={dict(Counter(source_names or []))}"
|
|
),
|
|
(
|
|
"[debug:data] "
|
|
f"input_tokens(min/mean/max)={_format_tensor_triplet(input_ids_lengths)} "
|
|
f"text_tokens(min/mean/max)={_format_tensor_triplet(num_text_tokens)} "
|
|
f"audio_tokens(min/mean/max)={_format_tensor_triplet(num_audio_tokens)} "
|
|
f"audio_seconds(min/mean/max)={_format_audio_seconds_triplet(sample_lengths, sample_rate)}"
|
|
),
|
|
]
|
|
|
|
fbank = batch.get("fbank")
|
|
fbank_lengths = batch.get("fbank_lengths")
|
|
if isinstance(fbank, torch.Tensor):
|
|
debug_lines.append(
|
|
"[debug:data] "
|
|
f"fbank_shape={tuple(fbank.shape)} "
|
|
f"fbank_frames(min/mean/max)={_format_tensor_triplet(fbank_lengths)}"
|
|
)
|
|
|
|
loss_masks = batch.get("loss_masks")
|
|
if isinstance(loss_masks, dict):
|
|
debug_lines.append(
|
|
"[debug:data] "
|
|
"loss_masks="
|
|
+ ", ".join(
|
|
f"{name}:{_format_mask_density(mask)}"
|
|
for name, mask in sorted(loss_masks.items())
|
|
)
|
|
)
|
|
|
|
fids = batch.get("fids") or []
|
|
sample_count = min(int(input_ids.size(0)), 3)
|
|
for sample_idx in range(sample_count):
|
|
input_length = int(input_ids_lengths[sample_idx].item())
|
|
audio_length = int(sample_lengths[sample_idx].item())
|
|
fbank_shape = "unavailable"
|
|
if isinstance(fbank, torch.Tensor) and isinstance(fbank_lengths, torch.Tensor):
|
|
fbank_shape = (
|
|
f"({int(fbank_lengths[sample_idx].item())}, {int(fbank.size(-1))})"
|
|
)
|
|
debug_lines.append(
|
|
"[debug:data] "
|
|
f"sample_index={sample_idx} "
|
|
f"fid={str(fids[sample_idx]) if sample_idx < len(fids) else f'sample_{sample_idx:02d}'} "
|
|
f"source_name={source_names[sample_idx] if source_names else None} "
|
|
f"input_ids_shape=({input_length},) "
|
|
f"sample_shape=(1, {audio_length}) "
|
|
f"fbank_shape={fbank_shape} "
|
|
f"num_text_tokens={int(num_text_tokens[sample_idx].item())} "
|
|
f"num_audio_tokens={int(num_audio_tokens[sample_idx].item())} "
|
|
f"audio_seconds={audio_length / float(sample_rate):.2f} "
|
|
"text="
|
|
f"{tokenizer.decode(input_ids[sample_idx, :input_length].detach().cpu().tolist(), skip_special_tokens=False, clean_up_tokenization_spaces=False)!r}"
|
|
)
|
|
return debug_lines
|
|
|
|
|
|
def should_print_gradient_debug(
|
|
*,
|
|
debug_enabled: bool,
|
|
is_main_process: bool,
|
|
next_global_step: int,
|
|
log_interval: int,
|
|
early_step_limit: int,
|
|
) -> bool:
|
|
return bool(
|
|
debug_enabled
|
|
and is_main_process
|
|
and (
|
|
next_global_step <= early_step_limit
|
|
or next_global_step % log_interval == 0
|
|
)
|
|
)
|
|
|
|
|
|
def build_gradient_debug_lines(
|
|
model: torch.nn.Module,
|
|
*,
|
|
global_step: int,
|
|
grad_norm: float,
|
|
grad_clip_norm: float,
|
|
) -> list[str]:
|
|
top_param_candidates: list[tuple[str, float, float, float]] = []
|
|
nonfinite_grad_params: list[str] = []
|
|
nonfinite_param_count = 0
|
|
params_with_grad = 0
|
|
params_without_grad = 0
|
|
abs_sum = 0.0
|
|
abs_count = 0
|
|
max_abs_grad = 0.0
|
|
|
|
for name, parameter in model.named_parameters():
|
|
if not parameter.requires_grad:
|
|
continue
|
|
grad = parameter.grad
|
|
if grad is None:
|
|
params_without_grad += 1
|
|
continue
|
|
|
|
grad_tensor = grad.detach().float()
|
|
params_with_grad += 1
|
|
if not bool(torch.isfinite(grad_tensor).all().item()):
|
|
nonfinite_param_count += 1
|
|
if len(nonfinite_grad_params) < 8:
|
|
nonfinite_grad_params.append(name)
|
|
grad_abs = grad_tensor.abs()
|
|
param_norm = float(torch.linalg.vector_norm(grad_tensor).item())
|
|
param_max_abs = float(grad_abs.max().item())
|
|
param_mean_abs = float(grad_abs.mean().item())
|
|
max_abs_grad = max(max_abs_grad, param_max_abs)
|
|
abs_sum += float(grad_abs.sum().item())
|
|
abs_count += int(grad_abs.numel())
|
|
top_param_candidates.append((name, param_norm, param_max_abs, param_mean_abs))
|
|
|
|
mean_abs_grad = math.nan if abs_count == 0 else abs_sum / float(abs_count)
|
|
top_param_norms = sorted(
|
|
top_param_candidates,
|
|
key=lambda item: item[1],
|
|
reverse=True,
|
|
)[:6]
|
|
|
|
debug_lines = [
|
|
(
|
|
"[debug:grad] "
|
|
f"step={global_step} "
|
|
f"pre_clip_grad_norm={format_scalar(grad_norm)} "
|
|
f"clip_ratio={format_scalar(_safe_grad_clip_ratio(grad_norm, grad_clip_norm))} "
|
|
f"params_with_grad={params_with_grad} "
|
|
f"params_without_grad={params_without_grad} "
|
|
f"nonfinite_param_count={nonfinite_param_count} "
|
|
f"max_abs_grad={format_scalar(max_abs_grad)} "
|
|
f"mean_abs_grad={format_scalar(mean_abs_grad)}"
|
|
)
|
|
]
|
|
if top_param_norms:
|
|
debug_lines.append(
|
|
"[debug:grad] top_params="
|
|
+ ", ".join(
|
|
(
|
|
f"{name}:{param_norm:.4f}"
|
|
f"(max={param_max_abs:.4e},mean={param_mean_abs:.4e})"
|
|
)
|
|
for name, param_norm, param_max_abs, param_mean_abs in top_param_norms
|
|
)
|
|
)
|
|
if nonfinite_grad_params:
|
|
debug_lines.append(
|
|
"[debug:grad] nonfinite_params=" + ", ".join(nonfinite_grad_params)
|
|
)
|
|
return debug_lines
|
|
|
|
|
|
def _format_tensor_triplet(values: object) -> str:
|
|
if not isinstance(values, torch.Tensor) or values.numel() == 0:
|
|
return "n/a"
|
|
flattened = values.detach().cpu().to(torch.float32)
|
|
return (
|
|
f"{int(flattened.min().item())}/"
|
|
f"{flattened.mean().item():.2f}/"
|
|
f"{int(flattened.max().item())}"
|
|
)
|
|
|
|
|
|
def _format_audio_seconds_triplet(values: object, sample_rate: int) -> str:
|
|
if not isinstance(values, torch.Tensor) or values.numel() == 0:
|
|
return "n/a"
|
|
seconds = values.detach().cpu().to(torch.float32) / float(sample_rate)
|
|
return (
|
|
f"{seconds.min().item():.2f}/"
|
|
f"{seconds.mean().item():.2f}/"
|
|
f"{seconds.max().item():.2f}"
|
|
)
|
|
|
|
|
|
def _format_mask_density(mask: object) -> str:
|
|
if not isinstance(mask, torch.Tensor) or mask.numel() == 0:
|
|
return "n/a"
|
|
return f"{int(mask.detach().gt(0).sum().item())}/{int(mask.numel())}"
|
|
|
|
|
|
def _safe_grad_clip_ratio(grad_norm: float, grad_clip_norm: float) -> float:
|
|
if not math.isfinite(grad_norm):
|
|
return math.nan
|
|
return grad_norm / float(grad_clip_norm)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Step Reporting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def should_log_training_step(global_step: int, log_interval: int) -> bool:
|
|
return global_step % log_interval == 0
|
|
|
|
|
|
def reduce_source_metrics(
|
|
source_loss_totals: dict[str, dict[str, float]],
|
|
source_loss_denominators: dict[str, dict[str, float]],
|
|
*,
|
|
device: torch.device,
|
|
loss_config: Any,
|
|
) -> dict[str, dict[str, float]]:
|
|
reduced_source_totals = loss_ops.sum_grouped_named_scalars_across_ranks(
|
|
source_loss_totals,
|
|
device=device,
|
|
)
|
|
reduced_source_denominators = loss_ops.sum_grouped_named_scalars_across_ranks(
|
|
source_loss_denominators,
|
|
device=device,
|
|
)
|
|
return loss_ops.reduce_loss_statistics_by_source(
|
|
reduced_source_totals,
|
|
reduced_source_denominators,
|
|
loss_config=loss_config,
|
|
)
|
|
|
|
|
|
def build_train_step_report(
|
|
metrics: dict[str, Any],
|
|
*,
|
|
learning_rate: float,
|
|
grad_norm: float,
|
|
current_time: float,
|
|
last_log_step: int,
|
|
last_log_time: float,
|
|
progress: TrainProgress,
|
|
max_train_steps: int,
|
|
reduced_by_source: dict[str, dict[str, float]],
|
|
) -> TrainStepReport:
|
|
logged_steps = progress.global_step - last_log_step
|
|
elapsed = current_time - last_log_time
|
|
steps_per_second = (
|
|
math.nan
|
|
if logged_steps <= 0 or elapsed <= 0.0
|
|
else float(logged_steps) / elapsed
|
|
)
|
|
eta_seconds = (
|
|
math.nan
|
|
if not math.isfinite(steps_per_second) or steps_per_second <= 0.0
|
|
else float(max_train_steps - progress.global_step) / steps_per_second
|
|
)
|
|
return TrainStepReport(
|
|
log_values=build_train_log_dict(
|
|
metrics,
|
|
learning_rate=learning_rate,
|
|
grad_norm=grad_norm,
|
|
steps_per_second=steps_per_second,
|
|
eta_seconds=eta_seconds,
|
|
progress=progress,
|
|
reduced_by_source=reduced_by_source,
|
|
),
|
|
console_line=format_train_line(
|
|
metrics,
|
|
learning_rate=learning_rate,
|
|
grad_norm=grad_norm,
|
|
steps_per_second=steps_per_second,
|
|
eta_seconds=eta_seconds,
|
|
progress=progress,
|
|
max_train_steps=max_train_steps,
|
|
reduced_by_source=reduced_by_source,
|
|
),
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Formatting Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def flatten_config(values, parent_key="", sep="/"):
|
|
"""Flatten a nested config dict into ``path/to/key -> value`` pairs."""
|
|
items = []
|
|
for key, value in values.items():
|
|
new_key = f"{parent_key}{sep}{key}" if parent_key else key
|
|
if isinstance(value, dict):
|
|
items.extend(flatten_config(value, new_key, sep).items())
|
|
elif isinstance(value, (list, tuple)):
|
|
items.append((new_key, str(value)))
|
|
elif value is None:
|
|
items.append((new_key, "None"))
|
|
else:
|
|
items.append((new_key, value))
|
|
return dict(items)
|
|
|
|
|
|
def format_scalar(value: float) -> str:
|
|
"""Format a scalar for concise console logging."""
|
|
if not math.isfinite(value):
|
|
return "nan"
|
|
if float(value).is_integer():
|
|
return str(int(value))
|
|
return f"{value:.4f}"
|
|
|
|
|
|
def _format_eta(eta_seconds: float) -> str:
|
|
"""Render ETA seconds as ``HH:MM:SS`` or ``n/a``."""
|
|
if not math.isfinite(eta_seconds) or eta_seconds < 0.0:
|
|
return "n/a"
|
|
total_seconds = int(round(eta_seconds))
|
|
hours, remainder = divmod(total_seconds, 3600)
|
|
minutes, seconds = divmod(remainder, 60)
|
|
return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
|
|
|
|
|
|
def build_train_log_dict(
|
|
metrics: dict[str, Any],
|
|
*,
|
|
learning_rate: float,
|
|
grad_norm: float,
|
|
steps_per_second: float,
|
|
eta_seconds: float,
|
|
progress: TrainProgress,
|
|
reduced_by_source: dict[str, dict[str, Any]],
|
|
) -> dict[str, float]:
|
|
"""Build the flat metric dict sent to experiment trackers."""
|
|
log_dict = {
|
|
"train/epoch": float(progress.epoch),
|
|
"train/learning_rate": learning_rate,
|
|
"train/grad_norm": grad_norm,
|
|
"train/steps_per_second": steps_per_second,
|
|
"train/eta_seconds": eta_seconds,
|
|
"train/consumed_tokens": float(progress.total_tokens),
|
|
"train/consumed_audio_tokens": float(progress.audio_tokens),
|
|
"train/consumed_text_tokens": float(progress.text_tokens),
|
|
}
|
|
for name, value in metrics.items():
|
|
log_dict[f"train/{name}"] = float(value)
|
|
for source_name, source_metrics in reduced_by_source.items():
|
|
log_dict.update(
|
|
{
|
|
f"train/{source_name}/{name}": float(value)
|
|
for name, value in source_metrics.items()
|
|
}
|
|
)
|
|
return log_dict
|
|
|
|
|
|
def format_train_line(
|
|
metrics: dict[str, Any],
|
|
*,
|
|
learning_rate: float,
|
|
grad_norm: float,
|
|
steps_per_second: float,
|
|
eta_seconds: float,
|
|
progress: TrainProgress,
|
|
max_train_steps: int,
|
|
reduced_by_source: dict[str, dict[str, Any]],
|
|
) -> str:
|
|
"""Build a single human-readable console line for one training step."""
|
|
parts = [
|
|
f"iteration {progress.global_step}/{max_train_steps}",
|
|
f"epoch: {progress.epoch}",
|
|
f"consumed_tokens: {progress.total_tokens}",
|
|
f"consumed_audio_tokens: {progress.audio_tokens}",
|
|
f"consumed_text_tokens: {progress.text_tokens}",
|
|
f"learning_rate: {learning_rate:.2e}",
|
|
f"steps_per_second: {format_scalar(steps_per_second)}",
|
|
f"job_eta: {_format_eta(eta_seconds)}",
|
|
f"grad_norm: {format_scalar(grad_norm)}",
|
|
]
|
|
for name in sorted(name for name in metrics if name != "loss"):
|
|
parts.append(f"{name}: {format_scalar(float(metrics[name]))}")
|
|
if "loss" in metrics:
|
|
parts.append(f"loss: {format_scalar(float(metrics['loss']))}")
|
|
for source_name, source_metrics in reduced_by_source.items():
|
|
for name in sorted(name for name in source_metrics if name != "loss"):
|
|
parts.append(
|
|
f"{source_name}_{name}: {format_scalar(float(source_metrics[name]))}"
|
|
)
|
|
if "loss" in source_metrics:
|
|
parts.append(
|
|
f"{source_name}_loss: {format_scalar(float(source_metrics['loss']))}"
|
|
)
|
|
return " | ".join(parts)
|
|
|
|
|
|
def build_validation_log_dict(
|
|
metrics: dict[str, Any],
|
|
*,
|
|
reduced_by_source: dict[str, dict[str, Any]],
|
|
) -> dict[str, float]:
|
|
"""Build the flat validation metric dict sent to experiment trackers."""
|
|
log_dict = {f"val/{name}": float(value) for name, value in metrics.items()}
|
|
for source_name, source_metrics in reduced_by_source.items():
|
|
log_dict.update(
|
|
{
|
|
f"val/{source_name}/{name}": float(value)
|
|
for name, value in source_metrics.items()
|
|
}
|
|
)
|
|
return log_dict
|
|
|
|
|
|
def format_validation_line(
|
|
metrics: dict[str, Any],
|
|
*,
|
|
global_step: int,
|
|
reduced_by_source: dict[str, dict[str, Any]],
|
|
) -> str:
|
|
"""Build the console summary line printed after a validation pass."""
|
|
parts = [f"validation at iteration {global_step}"]
|
|
for name in sorted(name for name in metrics if name != "loss"):
|
|
parts.append(f"{name}: {format_scalar(float(metrics[name]))}")
|
|
if "loss" in metrics:
|
|
parts.append(f"loss: {format_scalar(float(metrics['loss']))}")
|
|
for source_name, source_metrics in reduced_by_source.items():
|
|
for name in sorted(name for name in source_metrics if name != "loss"):
|
|
parts.append(
|
|
f"{source_name}_{name}: {format_scalar(float(source_metrics[name]))}"
|
|
)
|
|
if "loss" in source_metrics:
|
|
parts.append(
|
|
f"{source_name}_loss: {format_scalar(float(source_metrics['loss']))}"
|
|
)
|
|
return " | ".join(parts)
|
|
|
|
|
|
__all__ = [
|
|
"TrainProgress",
|
|
"TrainStepReport",
|
|
"abort_on_out_of_memory",
|
|
"any_rank_true",
|
|
"build_data_debug_lines",
|
|
"build_gradient_debug_lines",
|
|
"build_train_step_report",
|
|
"build_train_log_dict",
|
|
"build_validation_log_dict",
|
|
"flatten_config",
|
|
"format_scalar",
|
|
"format_train_line",
|
|
"format_validation_line",
|
|
"move_to_device",
|
|
"reduce_source_metrics",
|
|
"should_log_training_step",
|
|
"should_print_gradient_debug",
|
|
"sum_integer_counters_across_ranks",
|
|
]
|