| import argparse |
| import logging |
| import math |
| import os |
| import shutil |
| import time |
| from copy import deepcopy |
| from glob import glob |
|
|
| import bitsandbytes as bnb |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from accelerate import Accelerator |
| from accelerate.utils import set_seed |
| from transformers import get_cosine_schedule_with_warmup |
| from torchao.float8 import Float8LinearConfig, convert_to_float8_training |
|
|
| from dataset import build_tts_dataloader, extract_speaker_embs_from_batch, load_titanet |
|
|
| from discriminator import ConformerDiscirminator |
| from model_transformer import load_config, model_from_config |
|
|
| logging.getLogger("tilelang.cache.kernel_cache").setLevel(logging.ERROR) |
|
|
| torch.set_float32_matmul_precision("high") |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
| logger = logging.getLogger(__name__) |
|
|
| FP8_SKIP_EXACT = { |
| "in_proj", |
| "out_proj", |
| "cond_proj", |
| "span_mask_proj", |
| "valid_audio_proj", |
| } |
|
|
| FP8_SKIP_PREFIXES = ( |
| "cond_module.", |
| "speaker_film_adapter.", |
| "speaker_reg_head.", |
| "reference_proj.", |
| ) |
|
|
|
|
| def module_filter_fn(mod, fqn): |
| if not isinstance(mod, nn.Linear): |
| return False |
| if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: |
| return False |
| if mod.in_features < 256 or mod.out_features < 256: |
| return False |
| if fqn in FP8_SKIP_EXACT or fqn.startswith(FP8_SKIP_PREFIXES): |
| return False |
| if ".mamba." in fqn: |
| return False |
| return True |
|
|
|
|
| def get_num_params(model): |
| n = sum(p.numel() for p in model.parameters()) |
| return f"{n / 1e9:.2f}B" if n >= 1e9 else f"{n / 1e6:.1f}M" |
|
|
|
|
| def format_eta(seconds): |
| h = int(seconds // 3600) |
| m = int((seconds % 3600) // 60) |
| return f"{h}h {m:02d}m" if h > 0 else f"{m}m" |
|
|
|
|
| def maybe_len(obj): |
| try: |
| return len(obj) |
| except TypeError: |
| return None |
|
|
|
|
| def unwrap_model(model): |
| while hasattr(model, "module"): |
| model = model.module |
| while hasattr(model, "_orig_mod"): |
| model = model._orig_mod |
| return model |
|
|
|
|
| def set_requires_grad(module, flag): |
| for p in module.parameters(): |
| p.requires_grad_(flag) |
|
|
|
|
| def sample_stratified_logit_normal(batch_size, device, dtype): |
| offset = torch.rand((), device=device) |
| shift = torch.randint(0, batch_size, (), device=device) |
| idx = (torch.arange(batch_size, device=device) + shift) % batch_size |
| q = (idx.float() + offset) / batch_size |
| z = math.sqrt(2.0) * torch.erfinv(2.0 * q - 1.0) |
| return torch.sigmoid(z).to(dtype) |
|
|
|
|
| def sample_masked_span( |
| latent_mask, |
| frac_min=0.7, |
| frac_max=1.0, |
| min_mask_len=1, |
| keep_at_least_one_unmasked=False, |
| ): |
| latent_mask = latent_mask.to(torch.bool) |
| B, _ = latent_mask.shape |
| device = latent_mask.device |
|
|
| span_mask = torch.zeros_like(latent_mask, dtype=torch.bool) |
| valid_lens = latent_mask.sum(dim=1).long() |
| frac_lengths = torch.empty(B, device=device).uniform_(frac_min, frac_max) |
|
|
| for i in range(B): |
| L = int(valid_lens[i].item()) |
| if L <= 0: |
| continue |
| max_mask_len = max(L - 1, 0) if keep_at_least_one_unmasked else L |
| if max_mask_len <= 0: |
| continue |
| mask_len = max(min(int(frac_lengths[i].item() * L), max_mask_len), min_mask_len) |
| start = torch.randint(0, L - mask_len + 1, (1,), device=device).item() |
| span_mask[i, start: start + mask_len] = True |
|
|
| return span_mask & latent_mask |
|
|
|
|
| def extract_disc_features( |
| model, |
| latents, |
| latent_lengths, |
| text_ids, |
| text_mask, |
| detach_generator, |
| speaker_emb=None, |
| disc_t=None, |
| disc_noise=None, |
| ): |
|
|
| max_len = int(latent_lengths.max().item()) |
| latents = latents[:, :max_len, :] |
| valid_mask = torch.arange(max_len, device=latents.device)[None, :] < latent_lengths[:, None] |
| B = latents.size(0) |
|
|
| if disc_t is not None and disc_noise is not None: |
| disc_noise = disc_noise[:, :max_len, :] |
| t_exp = disc_t[:, None, None] |
| noisy_latents = (1.0 - t_exp) * disc_noise + t_exp * latents |
| noisy_latents = noisy_latents * valid_mask.unsqueeze(-1).to(noisy_latents.dtype) |
| t_for_model = disc_t |
| else: |
| noisy_latents = latents |
| t_for_model = torch.ones(B, device=latents.device, dtype=latents.dtype) |
|
|
| kwargs = dict( |
| x=noisy_latents, |
| cond=noisy_latents, |
| t=t_for_model, |
| text_ids=text_ids, |
| text_mask=text_mask, |
| reference_latent=None, |
| reference_mask=None, |
| text_cond_drop=torch.zeros(B, device=latents.device, dtype=torch.bool), |
| speaker_cond_drop=None, |
| use_checkpoint=False, |
| latent_mask=valid_mask, |
| valid_audio_mask=valid_mask, |
| span_mask=torch.zeros_like(valid_mask), |
| return_hidden_states=True, |
| speaker_emb=speaker_emb, |
| ) |
|
|
| if detach_generator: |
| with torch.no_grad(): |
| _, gen_layers = model(**kwargs) |
| return [x.detach() for x in gen_layers] |
|
|
| _, gen_layers = model(**kwargs) |
| return gen_layers |
|
|
|
|
| def compute_discriminator_loss(model, discriminator, cache): |
| B = cache["fake_latents"].size(0) |
| device = cache["fake_latents"].device |
| dtype = cache["fake_latents"].dtype |
|
|
| disc_t = torch.rand(B, device=device, dtype=dtype) |
| disc_noise = torch.randn_like(cache["fake_latents"]) |
|
|
| shared = dict( |
| model=model, |
| latent_lengths=cache["latent_lengths"], |
| text_ids=cache["text_ids"], |
| text_mask=cache["text_mask"], |
| detach_generator=True, |
| speaker_emb=cache.get("speaker_embs"), |
| disc_t=disc_t, |
| disc_noise=disc_noise, |
| ) |
|
|
| fake_feats = extract_disc_features(latents=cache["fake_latents"], **shared) |
| real_feats = extract_disc_features(latents=cache["real_latents"], **shared) |
|
|
| fake_score = discriminator(fake_feats, cache["latent_lengths"]) |
| real_score = discriminator(real_feats, cache["latent_lengths"]) |
|
|
| return (fake_score ** 2).mean() + ((1.0 - real_score) ** 2).mean() |
|
|
|
|
| def compute_loss( |
| model, |
| batch, |
| use_checkpoint, |
| *, |
| stage2=False, |
| use_speaker=False, |
| use_adversarial=False, |
| discriminator=None, |
| adv_lambda=0.0, |
| speaker_aux_weight=0.1, |
| ): |
| latents = batch["latents"] |
| valid_audio_mask = batch["latent_mask"].to(torch.bool) |
| text_ids = batch["text_ids"] |
| text_mask = batch["text_mask"].to(torch.bool) |
|
|
| B = latents.size(0) |
| device = latents.device |
| dtype = latents.dtype |
|
|
| x1 = latents * valid_audio_mask.unsqueeze(-1).to(latents.dtype) |
| span_mask = sample_masked_span(valid_audio_mask, frac_min=0.7, frac_max=1.0) |
|
|
| x0 = torch.randn_like(x1) |
| t = sample_stratified_logit_normal(B, device=device, dtype=dtype) |
| t_exp = t[:, None, None] |
|
|
| x_t = (1.0 - t_exp) * x0 + t_exp * x1 |
| v_target = x1 - x0 |
|
|
| cond = torch.where(span_mask.unsqueeze(-1), torch.zeros_like(x1), x1) |
| cond = cond * valid_audio_mask.unsqueeze(-1).to(cond.dtype) |
| model_input = x_t * valid_audio_mask.unsqueeze(-1).to(x_t.dtype) |
|
|
|
|
| p_drop_both = 0.1 |
| p_drop_text = 0.1 |
| p_drop_speaker = 0.0 |
|
|
| r = torch.rand(B, device=device) |
| b0 = p_drop_both |
| b1 = b0 + p_drop_text |
| b2 = b1 + p_drop_speaker |
|
|
| drop_both_bucket = r < b0 |
| drop_text_bucket = (r >= b0) & (r < b1) |
| drop_speaker_bucket = (r >= b1) & (r < b2) |
|
|
| text_cond_drop = drop_both_bucket | drop_text_bucket |
|
|
| speaker_emb = batch.get("speaker_embs", None) |
| if speaker_emb is not None: |
| speaker_emb = speaker_emb.to(device) |
|
|
| speaker_cond_drop = None |
| if stage2 and use_speaker and speaker_emb is not None: |
| speaker_cond_drop = drop_both_bucket | drop_speaker_bucket |
|
|
| need_hidden = stage2 and use_speaker and speaker_emb is not None |
| out = model( |
| x=model_input, |
| cond=cond, |
| t=t, |
| text_ids=text_ids, |
| text_mask=text_mask, |
| reference_latent=None, |
| reference_mask=None, |
| text_cond_drop=text_cond_drop, |
| speaker_cond_drop=speaker_cond_drop, |
| use_checkpoint=use_checkpoint, |
| latent_mask=valid_audio_mask, |
| valid_audio_mask=valid_audio_mask, |
| span_mask=span_mask, |
| speaker_emb=speaker_emb if (stage2 and use_speaker) else None, |
| return_hidden_states=need_hidden, |
| ) |
| v_pred, hidden_states = out if need_hidden else (out, None) |
|
|
| per_pos_loss = F.mse_loss(v_pred.float(), v_target.float(), reduction="none").mean(dim=-1) |
| masked_loss = per_pos_loss[span_mask] |
| fm_loss = masked_loss.mean() if masked_loss.numel() > 0 else per_pos_loss.new_zeros(()) |
|
|
| zero = fm_loss.new_zeros(()) |
| if not stage2: |
| stats = {"fm_loss": fm_loss.detach(), "gen_adv_loss": zero, "speaker_aux_loss": zero} |
| return fm_loss, stats, None |
|
|
|
|
| raw = unwrap_model(model) |
| speaker_aux_loss = zero |
| if need_hidden and hasattr(raw, "predict_speaker"): |
| spk_pred = raw.predict_speaker(hidden_states[-1], valid_audio_mask) |
| if spk_pred is not None: |
| valid_spk = speaker_emb.abs().sum(dim=-1) > 0 |
| if speaker_cond_drop is not None: |
| valid_spk = valid_spk & ~speaker_cond_drop.to(torch.bool) |
| if valid_spk.any(): |
| speaker_aux_loss = F.mse_loss(spk_pred[valid_spk], speaker_emb[valid_spk]) |
|
|
| total_loss = fm_loss + speaker_aux_weight * speaker_aux_loss |
|
|
| gen_adv_loss = zero |
| disc_cache = None |
| if use_adversarial and discriminator is not None and adv_lambda > 0.0: |
| latent_lengths = valid_audio_mask.sum(dim=1).long() |
|
|
|
|
| x0_gen = torch.randn_like(x1) |
| t_gen = torch.rand(B, device=device, dtype=dtype) |
| t_gen_exp = t_gen[:, None, None] |
| x_t_gen = (1.0 - t_gen_exp) * x0_gen + t_gen_exp * x1 |
|
|
| v_gen = model( |
| x=x_t_gen * valid_audio_mask.unsqueeze(-1).to(x_t_gen.dtype), |
| cond=torch.zeros_like(x1), |
| t=t_gen, |
| text_ids=text_ids, |
| text_mask=text_mask, |
| reference_latent=None, |
| reference_mask=None, |
| text_cond_drop=torch.zeros(B, device=device, dtype=torch.bool), |
| use_checkpoint=use_checkpoint, |
| latent_mask=valid_audio_mask, |
| valid_audio_mask=valid_audio_mask, |
| span_mask=valid_audio_mask, |
| speaker_emb=speaker_emb if use_speaker else None, |
| ) |
| x1_hat = (x0_gen + v_gen) * valid_audio_mask.unsqueeze(-1).to(v_gen.dtype) |
|
|
| fake_feats = extract_disc_features( |
| model=model, |
| latents=x1_hat, |
| latent_lengths=latent_lengths, |
| text_ids=text_ids, |
| text_mask=text_mask, |
| detach_generator=False, |
| speaker_emb=speaker_emb if use_speaker else None, |
| disc_t=torch.rand(B, device=device, dtype=dtype), |
| disc_noise=torch.randn_like(x1), |
| ) |
| gen_adv_loss = ((1.0 - discriminator(fake_feats, latent_lengths)) ** 2).mean() |
|
|
| disc_cache = { |
| "fake_latents": x1_hat.detach(), |
| "real_latents": x1.detach(), |
| "latent_lengths": latent_lengths.detach(), |
| "text_ids": text_ids.detach(), |
| "text_mask": text_mask.detach(), |
| "speaker_embs": speaker_emb.detach() if speaker_emb is not None else None, |
| } |
|
|
| total_loss = total_loss + adv_lambda * gen_adv_loss |
|
|
| stats = { |
| "fm_loss": fm_loss.detach(), |
| "gen_adv_loss": gen_adv_loss.detach(), |
| "speaker_aux_loss": speaker_aux_loss.detach(), |
| } |
| return total_loss, stats, disc_cache |
|
|
|
|
| @torch.no_grad() |
| def update_ema(ema_model, model, decay): |
| ema_params = dict(ema_model.named_parameters()) |
| for name, param in model.named_parameters(): |
| if name in ema_params: |
| ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) |
|
|
| ema_bufs = dict(ema_model.named_buffers()) |
| for name, buf in model.named_buffers(): |
| if name in ema_bufs: |
| ema_bufs[name].copy_(buf) |
|
|
|
|
| def save_checkpoint(accelerator, model, ema, step, exp_dir, tag): |
| ckpt_dir = os.path.join(exp_dir, "checkpoints", tag) |
| os.makedirs(ckpt_dir, exist_ok=True) |
| accelerator.save_state(ckpt_dir) |
|
|
| if accelerator.is_main_process: |
| torch.save(unwrap_model(model).state_dict(), os.path.join(ckpt_dir, "model_state.pt")) |
| torch.save({"step": step, "ema": ema.state_dict()}, |
| os.path.join(ckpt_dir, "extra_state.pt")) |
|
|
| accelerator.wait_for_everyone() |
| logger.info(f"Saved checkpoint: {tag}") |
|
|
|
|
| def load_checkpoint(accelerator, model, ema, exp_dir, tag, drop_optimizer_state=False): |
| ckpt_dir = os.path.join(exp_dir, "checkpoints", tag) |
|
|
| if drop_optimizer_state: |
| logger.info("drop_optimizer_state=True -> skipping optimizer/scheduler/RNG restore") |
| else: |
| accelerator.load_state(ckpt_dir) |
|
|
| model_state_path = os.path.join(ckpt_dir, "model_state.pt") |
| if os.path.exists(model_state_path): |
| state = torch.load(model_state_path, map_location="cpu", weights_only=True) |
| unwrap_model(model).load_state_dict(state, strict=False) |
| logger.info("Loaded model state_dict from model_state.pt (strict=False)") |
|
|
| extra = torch.load(os.path.join(ckpt_dir, "extra_state.pt"), |
| map_location="cpu", weights_only=True) |
|
|
|
|
| ema.load_state_dict(extra["ema"], strict=False) |
| ema.to(accelerator.device) |
|
|
| if drop_optimizer_state: |
| logger.info(f"Loaded weights from {tag} | optimizer dropped | step reset to 0") |
| return 0 |
|
|
| logger.info(f"Resumed from {tag} (step {extra['step']})") |
| return extra["step"] |
|
|
|
|
| def find_latest_checkpoint(exp_dir): |
| dirs = sorted(glob(os.path.join(exp_dir, "checkpoints", "step_*"))) |
| return os.path.basename(dirs[-1]) if dirs else None |
|
|
|
|
| def cleanup_checkpoints(exp_dir, keep_last_k): |
| dirs = sorted(glob(os.path.join(exp_dir, "checkpoints", "step_*"))) |
| for d in dirs[:-keep_last_k]: |
| shutil.rmtree(d, ignore_errors=True) |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", type=str, required=True) |
| parser.add_argument("--dataset", type=str, required=True) |
| parser.add_argument("--streaming", action="store_true") |
| parser.add_argument("--dataset_split", type=str, default="train") |
| parser.add_argument("--tokenizer", type=str, default=None) |
| parser.add_argument("--exp_dir", type=str, required=True) |
| parser.add_argument("--resume", type=str, default=None) |
| parser.add_argument("--drop_optimizer_state", action="store_true") |
|
|
| parser.add_argument("--stage2", action="store_true") |
| parser.add_argument("--use_speaker", action="store_true") |
| parser.add_argument("--speaker_aux_weight", type=float, default=0.1) |
| parser.add_argument("--speaker_drop_prob", type=float, default=0.1) |
|
|
| parser.add_argument("--use_adversarial", action="store_true") |
| parser.add_argument("--adv_lambda", type=float, default=1.0) |
| parser.add_argument("--disc_lr", type=float, default=5e-5) |
| parser.add_argument("--disc_channels", type=int, default=512) |
| parser.add_argument("--disc_num_layers", type=int, default=3) |
| parser.add_argument("--disc_num_heads", type=int, default=8) |
| parser.add_argument("--disc_kernel", type=int, default=15) |
| return parser.parse_args() |
|
|
|
|
| def build_optimizer(args, model, train_cfg): |
| if not args.stage2: |
| return bnb.optim.AdamW8bit( |
| model.parameters(), |
| lr=train_cfg["lr"], |
| betas=tuple(train_cfg["betas"]), |
| weight_decay=train_cfg["weight_decay"], |
| ) |
|
|
| base_lr, speaker_lr = 2e-5, 1e-4 |
| raw = unwrap_model(model) |
| speaker_params, speaker_param_ids = [], set() |
|
|
| if args.use_speaker: |
| for module in (raw.speaker_film_adapter, raw.speaker_reg_head): |
| if module is None: |
| continue |
| for p in module.parameters(): |
| if p.requires_grad: |
| speaker_params.append(p) |
| speaker_param_ids.add(id(p)) |
|
|
| base_params = [p for p in model.parameters() |
| if p.requires_grad and id(p) not in speaker_param_ids] |
|
|
| param_groups = [] |
| if base_params: |
| param_groups.append({"params": base_params, "lr": base_lr, |
| "weight_decay": train_cfg["weight_decay"]}) |
| if speaker_params: |
| param_groups.append({"params": speaker_params, "lr": speaker_lr, |
| "weight_decay": train_cfg["weight_decay"]}) |
|
|
| logger.info(f"Optimizer: base params={len(base_params)} @ lr={base_lr:.0e}, " |
| f"speaker params={len(speaker_params)} @ lr={speaker_lr:.0e}") |
|
|
| return bnb.optim.AdamW8bit( |
| param_groups, |
| betas=tuple(train_cfg["betas"]), |
| weight_decay=train_cfg["weight_decay"], |
| ) |
|
|
|
|
| def build_data(args, cfg, need_audio): |
| data_cfg = cfg["data"] |
| raw_max = data_cfg["max_audio_seconds"] * data_cfg["codec_rate_hz"] |
|
|
| logger.info(f"Loading dataset from {args.dataset}") |
| return build_tts_dataloader( |
| dataset_path=args.dataset, |
| latent_dim=data_cfg["latent_dim"], |
| need_audio=need_audio, |
| max_text_len=data_cfg["max_text_length"], |
| max_latent_len=int(data_cfg.get("max_latent_length", int(raw_max))), |
| text_pad_id=data_cfg["text_pad_id"], |
| train_cfg=cfg["training"], |
| tokenizer=( |
| args.tokenizer |
| or data_cfg.get("tokenizer") |
| or data_cfg.get("tokenizer_name") |
| or data_cfg.get("text_tokenizer") |
| ), |
| streaming=args.streaming, |
| dataset_split=args.dataset_split, |
| ) |
|
|
|
|
| def main(): |
| args = parse_args() |
|
|
| cfg = load_config(args.config) |
| train_cfg = cfg["training"] |
|
|
| accelerator = Accelerator( |
| gradient_accumulation_steps=train_cfg["grad_accum_steps"], |
| log_with=None, |
| ) |
|
|
| os.makedirs(args.exp_dir, exist_ok=True) |
| logging.basicConfig( |
| level=logging.INFO, |
| format="[\033[34m%(asctime)s\033[0m] %(message)s", |
| datefmt="%Y-%m-%d %H:%M:%S", |
| handlers=[ |
| logging.StreamHandler(), |
| logging.FileHandler(os.path.join(args.exp_dir, "log.txt"), mode="a"), |
| ], |
| ) |
|
|
| if accelerator.is_main_process: |
| shutil.copy2(args.config, os.path.join(args.exp_dir, "config.json")) |
|
|
| set_seed(train_cfg["seed"]) |
|
|
| logger.info(f"Experiment directory: {args.exp_dir}") |
| logger.info(f"Stage: {'2' if args.stage2 else '1 (vanilla)'}" |
| f"{' + speaker conditioning' if args.use_speaker else ''}" |
| f"{' + adversarial' if args.use_adversarial else ''}") |
|
|
| model = model_from_config( |
| cfg, |
| use_speaker_conditioning=args.use_speaker, |
| speaker_emb_dim=192, |
| ) |
| logger.info(f"Model parameters: {get_num_params(model)}") |
|
|
| ema = deepcopy(model) |
| ema.requires_grad_(False) |
| ema.eval() |
|
|
| convert_to_float8_training( |
| model, |
| config=Float8LinearConfig.from_recipe_name("tensorwise"), |
| module_filter_fn=module_filter_fn, |
| ) |
|
|
| if train_cfg.get("compile", False): |
| logger.info("Compiling model with torch.compile...") |
| model = torch.compile(model) |
|
|
| optimizer = build_optimizer(args, model, train_cfg) |
| scheduler = get_cosine_schedule_with_warmup( |
| optimizer, |
| num_warmup_steps=train_cfg["warmup_steps"], |
| num_training_steps=train_cfg["max_steps"], |
| ) |
|
|
| titanet_model = None |
| if args.stage2 and args.use_speaker: |
| titanet_model = load_titanet(accelerator.device) |
| logger.info("Loaded frozen TitaNet-Large for on-the-fly speaker embeddings") |
|
|
| dataset, dataloader = build_data(args, cfg, need_audio=args.stage2 and args.use_speaker) |
| n_samples = maybe_len(dataset) |
| logger.info(f"Dataset size: {n_samples:,} samples" if n_samples else "Dataset size: streaming") |
|
|
| adversarial = args.stage2 and args.use_adversarial |
| discriminator = None |
| disc_optimizer = None |
|
|
| if adversarial: |
| base_model = unwrap_model(model) |
| discriminator = ConformerDiscirminator( |
| input_dim=base_model.num_layers * base_model.model_size, |
| channels=args.disc_channels, |
| num_layers=args.disc_num_layers, |
| num_heads=args.disc_num_heads, |
| depthwise_conv_kernel_size=args.disc_kernel, |
| use_group_norm=True, |
| ) |
| disc_optimizer = bnb.optim.AdamW8bit( |
| discriminator.parameters(), |
| lr=args.disc_lr, |
| betas=(0.9, 0.95), |
| weight_decay=train_cfg["weight_decay"], |
| ) |
|
|
| model, optimizer, dataloader, scheduler, discriminator, disc_optimizer = accelerator.prepare( |
| model, optimizer, dataloader, scheduler, discriminator, disc_optimizer |
| ) |
| else: |
| model, optimizer, dataloader, scheduler = accelerator.prepare( |
| model, optimizer, dataloader, scheduler |
| ) |
|
|
| ema = ema.to(accelerator.device) |
|
|
| global_step = 0 |
| if args.resume is None: |
| update_ema(ema, unwrap_model(model), decay=0) |
| else: |
| tag = find_latest_checkpoint(args.exp_dir) if args.resume == "latest" else args.resume |
| if tag is None: |
| logger.warning("No checkpoint found, training from scratch") |
| else: |
| global_step = load_checkpoint( |
| accelerator, model, ema, args.exp_dir, tag, |
| drop_optimizer_state=args.drop_optimizer_state, |
| ) |
|
|
| max_steps = train_cfg["max_steps"] |
| grad_accum = train_cfg["grad_accum_steps"] |
| log_every = train_cfg["log_every"] |
| save_every = train_cfg["save_every"] |
| keep_last_k = train_cfg["keep_last_k"] |
| grad_clip = train_cfg["grad_clip"] |
| ema_decay = train_cfg["ema_decay"] |
| use_ckpt = train_cfg["use_checkpoint"] |
|
|
| num_batches = maybe_len(dataloader) |
| steps_per_epoch = (num_batches // grad_accum) if num_batches else max(max_steps - global_step, 1) |
| steps_per_epoch = max(steps_per_epoch, 1) |
|
|
| logger.info(f"Training from step {global_step} to {max_steps}") |
| logger.info(f"Batch size per device: {train_cfg['batch_size']}") |
| logger.info(f"Gradient accumulation: {grad_accum}") |
| logger.info( |
| f"Effective batch size: " |
| f"{train_cfg['batch_size'] * grad_accum * accelerator.num_processes}" |
| ) |
| if args.stage2: |
| logger.info( |
| f"Stage 2 | use_speaker={args.use_speaker} | adversarial={args.use_adversarial} | " |
| f"adv_lambda={args.adv_lambda if args.use_adversarial else 0.0} | " |
| f"speaker_aux_weight={args.speaker_aux_weight} | " |
| f"speaker_drop_prob={args.speaker_drop_prob}" |
| ) |
|
|
| model.train() |
| if adversarial: |
| discriminator.train() |
|
|
| running = {"loss": 0.0, "fm": 0.0, "adv": 0.0, "spk_aux": 0.0, "disc": 0.0} |
| log_steps = 0 |
| start_time = time.time() |
| epoch = global_step // steps_per_epoch |
|
|
| while global_step < max_steps: |
| epoch_step = 0 |
| logger.info(f"Starting epoch {epoch} ({steps_per_epoch} steps)") |
|
|
| for batch in dataloader: |
| accumulate_ctx = ( |
| accelerator.accumulate(model, discriminator) |
| if adversarial |
| else accelerator.accumulate(model) |
| ) |
|
|
| with accumulate_ctx: |
| if titanet_model is not None: |
| batch["speaker_embs"] = extract_speaker_embs_from_batch( |
| titanet_model, batch, accelerator.device, |
| ) |
|
|
| if adversarial: |
| set_requires_grad(discriminator, False) |
|
|
| with accelerator.autocast(): |
| loss, stats, disc_cache = compute_loss( |
| model=model, |
| batch=batch, |
| use_checkpoint=use_ckpt, |
| stage2=args.stage2, |
| use_speaker=args.use_speaker, |
| use_adversarial=adversarial, |
| discriminator=discriminator, |
| adv_lambda=args.adv_lambda if adversarial else 0.0, |
| speaker_aux_weight=args.speaker_aux_weight, |
| ) |
|
|
| accelerator.backward(loss) |
|
|
| if adversarial: |
| set_requires_grad(discriminator, True) |
| disc_loss = compute_discriminator_loss( |
| model=model, discriminator=discriminator, cache=disc_cache, |
| ) |
| accelerator.backward(disc_loss) |
|
|
| if accelerator.sync_gradients: |
| accelerator.clip_grad_norm_(model.parameters(), grad_clip) |
| accelerator.clip_grad_norm_(discriminator.parameters(), grad_clip) |
|
|
| optimizer.step() |
| scheduler.step() |
| optimizer.zero_grad() |
| disc_optimizer.step() |
| disc_optimizer.zero_grad() |
| else: |
| if accelerator.sync_gradients: |
| accelerator.clip_grad_norm_(model.parameters(), grad_clip) |
| optimizer.step() |
| scheduler.step() |
| optimizer.zero_grad() |
| disc_loss = loss.new_zeros(()) |
|
|
| if not accelerator.sync_gradients: |
| continue |
|
|
| global_step += 1 |
| epoch_step += 1 |
| update_ema(ema, unwrap_model(model), ema_decay) |
|
|
| running["loss"] += loss.detach().item() |
| running["fm"] += stats["fm_loss"].item() |
| running["adv"] += stats["gen_adv_loss"].item() |
| running["spk_aux"] += stats["speaker_aux_loss"].item() |
| running["disc"] += disc_loss.detach().item() |
| log_steps += 1 |
|
|
| if global_step % log_every == 0: |
| n = log_steps |
| elapsed = time.time() - start_time |
| sps = n / elapsed |
| speed = f"{sps:.2f} it/s" if sps >= 1.0 else f"{elapsed / n:.2f} s/it" |
| head = ( |
| f"step={global_step:>7d} | " |
| f"epoch={epoch} [{epoch_step}/{steps_per_epoch}] | " |
| f"loss={running['loss'] / n:.4f}" |
| ) |
| if args.stage2: |
| head += ( |
| f" | fm={running['fm'] / n:.4f}" |
| f" | spk_aux={running['spk_aux'] / n:.4f}" |
| f" | g_adv={running['adv'] / n:.4f}" |
| f" | d={running['disc'] / n:.4f}" |
| ) |
| logger.info( |
| f"{head} | lr={scheduler.get_last_lr()[0]:.2e} | {speed} | " |
| f"mem={torch.cuda.max_memory_allocated() / 1e9:.1f}GB | " |
| f"eta={format_eta((max_steps - global_step) / sps)}" |
| ) |
|
|
| running = {k: 0.0 for k in running} |
| log_steps = 0 |
| start_time = time.time() |
|
|
| if global_step % save_every == 0: |
| save_checkpoint(accelerator, model, ema, global_step, |
| args.exp_dir, f"step_{global_step:07d}") |
| if accelerator.is_main_process: |
| cleanup_checkpoints(args.exp_dir, keep_last_k) |
|
|
| if global_step >= max_steps: |
| break |
|
|
| epoch += 1 |
|
|
| save_checkpoint(accelerator, model, ema, global_step, |
| args.exp_dir, f"step_{global_step:07d}") |
| logger.info("Training complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
