import argparse from datetime import timedelta import functools import gc import importlib.metadata import json import logging import os import random import sys import tempfile import time import warnings import math import numpy as np import torch import torch.distributed as dist import torch.nn.functional as F import wandb from colorama import Fore, Style from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy from torch.distributed.fsdp.api import MixedPrecision from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.optim import AdamW from torch.utils.data import DataLoader from tqdm import tqdm from dataset import ( GeneralMemmapMapDataset, StemMemmapMapDataset, BundleDownloaderDataset, shard_data, ) from dynamic_dataset import collate_fn as dynamic_collate_fn from dynamic_dataset import DynamicDataset, CROP_PADDING_S from helpers import ( FSDP_EMA, save_checkpoint, dist_barrier, log_training_metrics, log_validation_metrics, maybe_enable_profiling, maybe_enable_memory_snapshot, print_with_time_master, load_checkpoint, ) from audio_metrics import calculate_stft_loss, calculate_mel_loss from generation import ( simple_generate, _load_dit_model, convert_to_precision, ratio_mask_semantic_codes, ) from prefix_model.model import Discriminator from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_codec_models, decode as codec_decode, encode_overlap as codec_encode, ) from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as semantic_encode, ) from suno_utils.tasks.ear import load_model as load_ear_model gc.disable() torch.set_float32_matmul_precision("high") # Set the TOKENIZERS_PARALLELISM environment variable to False to avoid warning os.environ["TOKENIZERS_PARALLELISM"] = "false" os.umask(0o003) # set umask to 0o003 to allow group write for created directories # base - 24 layers, 24 heads, 1536 dim # large - 32 layers, 32 heads, 2048 dim class TimingProfiler: """Simple profiling context manager for timing code sections.""" def __init__(self, enabled=False): self.enabled = enabled self.timings = {} self.counts = {} def __call__(self, name): return self._Timer(self, name) class _Timer: def __init__(self, profiler, name): self.profiler = profiler self.name = name self.start_time = None def __enter__(self): if self.profiler.enabled: self.start_time = time.perf_counter() return self def __exit__(self, *args): if self.profiler.enabled and self.start_time is not None: elapsed = time.perf_counter() - self.start_time if self.name not in self.profiler.timings: self.profiler.timings[self.name] = 0.0 self.profiler.counts[self.name] = 0 self.profiler.timings[self.name] += elapsed self.profiler.counts[self.name] += 1 def get_stats(self): """Get timing statistics.""" if not self.enabled: return {} stats = {} for name in self.timings: total = self.timings[name] count = self.counts[name] avg = total / count if count > 0 else 0 stats[name] = { "total": total, "count": count, "avg": avg, "percent": 0.0, # Will be calculated below } # Calculate percentages total_time = sum(self.timings.values()) if total_time > 0: for name in stats: stats[name]["percent"] = (stats[name]["total"] / total_time) * 100 return stats def reset(self): """Reset all timings.""" self.timings.clear() self.counts.clear() def print_stats(self, prefix="[Profiling]"): """Print timing statistics.""" if not self.enabled: return stats = self.get_stats() if not stats: return print(f"\n{prefix} Timing Statistics:") print(f"{'Section':<50} {'Total (s)':<12} {'Count':<8} {'Avg (ms)':<12} {'Percent':<8}") print("-" * 98) # Sort by total time descending sorted_stats = sorted(stats.items(), key=lambda x: x[1]["total"], reverse=True) for name, data in sorted_stats: print( f"{name:<50} {data['total']:<12.4f} {data['count']:<8} {data['avg'] * 1000:<12.2f} {data['percent']:<8.1f}%" ) print() DEFAULT_CONFIG = { "training": { "wandb_name": "48n_finetune", "wandb_project": "harmonai_train_vt", "model_type": "prefix", # default|prefix "tot_num_steps": 10_000_000, "ckpt_every": 50_000, "batch_size": 4, "learning_rate": 5e-5, "lr_warmup": 1_000, "lr_scheduler": "constant", # constant|cosine "diffusion_objective": "v", # v|rectified_flow|rf_denoiser "noise_schedule": "default", # default|logit_normal|early "use_fine_guidance": False, "weight_decay": 1e-3, "clip_grad_norm": 0.5, "betas": (0.9, 0.999), "seed_offset": 0, # increment when we continue from a checkpoint "compute_metrics": False, "use_ema": True, "activation_checkpointing": False, "compile": True, "preload_checkpoint": "/app/suno/checkpoints/2024-10-26_04-39-23_s121/step_1000000_ckpt.pt", "preload_optimizer": True, "check_model_args": True, "teacher_checkpoint": None, "use_critic": False, "acc_steps": 1, # Gradient accumulation steps "batch_store_size": 30, "batch_reuse_factor": 1, # Number of times to reuse each batch with different noise levels "profile_timing": False, # Enable detailed timing profiling "profile_print_every": 100, # Print profiling stats every N steps "num_workers": 1, "distill": { "finetune_mode": None, # "distill", "adversarial", or None for normal training "critic_lr": 5e-5, # Learning rate for critic optimizer "step_critic_N": 5, # Number of critic steps per generator step (N critic steps, 1 generator step) "critic_pretrain_steps": 0, # Number of steps to pretrain critic before alternating "use_discriminator": False, # Whether to use discriminator on intermediate features "residual": False, # Whether to use residual formulation in distillation "t_discretize": 1, # Discretization for timesteps in inference distribution "text_cfg_scale": 1.0, # Text guidance scale for teacher model "reward_model_checkpoint": None, # Path to reward model checkpoint "reward_model_pretrain_steps": 0, # Steps before using reward model }, }, "model": { "io_hz": 25, "io_channels": 128, "embed_dim": 2048, # depth * 64 "depth": 32, "n_heads": 32, # scale together with depth "qk_norm": True, "block_size": 750, "cond_semantic_n_vocab": 4001, "cond_semantic_len": 750, "cond_text_n_vocab": 60001, "cond_text_len": 1536, "ctx_len": 750, "infill_ctx_len": 0, "stem_ctx_len": 0, "shared_ctx": False, "use_actnorm": False, "actnorm_sync_init": False, "use_mmdit": False, "init_variant": "stability_v1", "use_rvq": False, "n_codebooks": 1, }, "data": { "mode": "pretraining", "shard_data": None, # None means auto (>=8 nodes) "shard_data_dir": "/mnt/localdisk/tmp", "allow_shard_reuse": False, "dataset_type": "general_memmap", "dataset_dir": "/app/suno/data/diffusion_mix/vae_25hz_30s", "train_metas_filepath": "/app/suno/data/diffusion_mix/vae_25hz_30s/metas_context_aligned_quality_tr.jsonl", "val_metas_filepath": "/app/suno/data/diffusion_mix/vae_25hz_30s/metas_context_aligned_quality_val.jsonl", "audio_chunk_s": 30.02, "audio_ctx_s": 30.02, "audio_vox_s": 30.02, "train_vae_memmap_filename": "data_vae_tr.bin", "train_semantic_memmap_filename": "data_semantic_tr.bin", "train_metas_filename": "metas_context_aligned_quality_tr.jsonl", "train_info_filename": None, "val_vae_memmap_filename": "data_vae_val.bin", "val_semantic_memmap_filename": "data_semantic_val.bin", "val_metas_filename": "metas_context_aligned_quality_val.jsonl", "val_info_filename": None, "codec_dir": "/app/suno/models/codecs", "vae_dim": 128, "semantic_rate_hz": 25, "vae_scale_factor": 2.5, "semantic_skip_factors": [], "patch_size": 1, "foreign_weight": 1.0, "text_aligned_weight": 1.0, "stem_weight": 1.0, "scale_vae_ctx": False, "prev_vae_ctx": True, "infill_vae_ctx": False, "always_pad_semantic": False, "noise_ctx": 0.2, "semantic_noise_level": 0.0, "always_skip_semantic": False, "aligned_text_prob": 0.5, "max_multi_instruments": 20, "respell_augment_prob": 0.0, "shared_ctx": False, "semantic_dropout": True, "use_cached_dataloader": False, "infill_prob": 0.1, "infill_min_ratio": 0.2, "infill_max_ratio": 0.5, "text_drop_prob": 0.1, "use_text_aligned_prob": 0.8, "target_loudness_db": -16.0, "codec_filepath": None, "semantic_model_filepath": None, "semantic_clusters_filepath": None, "semantic_mask_prob": 0.0, "ctx_mask_prob": 0.0, "use_stem_prob": 0.0, "use_vox_prob": 0.0, }, } RUN_START_TIME = time.strftime("%Y-%m-%d_%H-%M-%S") CHECKPOINT_DIR = f"/app2/suno/checkpoints/{RUN_START_TIME}_s{random.randint(0, 9999)}" torch_profiler = None memory_profiler = None # Utility functions def get_distill_config(run_config, key, default=None): """Get a distillation config value, with fallback to default.""" distill_config = run_config["training"].get("distill", {}) return distill_config.get(key, default) def update_nested_dict(original, updates): for key, value in updates.items(): if key not in original: raise KeyError( f"Invalid key '{key}' found in updates. It doesn't exist in the default configuration." ) if isinstance(value, dict): original[key] = update_nested_dict(original.get(key, {}), value) else: original[key] = value return original def load_and_update_config(default_config, config_path=None): if config_path is None: return default_config.copy() with open(config_path, "r") as f: config_overrides = json.load(f) return update_nested_dict(default_config.copy(), config_overrides) def _load_dit_model(dit_model_filepath, use_ema_if_exists, weights_precision, compile=False): state_dict, dit_config = load_checkpoint(dit_model_filepath, use_ema_if_exists=use_ema_if_exists) print(f"dit_config: {dit_config}") dit_model = DiffusionTransformer( io_hz=dit_config["io_hz"], io_channels=dit_config["io_channels"], embed_dim=dit_config["embed_dim"], depth=dit_config["depth"], n_heads=dit_config["n_heads"], qk_norm=dit_config["qk_norm"], block_size=dit_config["block_size"], cond_semantic_n_vocab=dit_config["cond_semantic_n_vocab"], cond_semantic_len=dit_config["cond_semantic_len"], cond_text_n_vocab=dit_config["cond_text_n_vocab"], cond_text_len=dit_config["cond_text_len"], ctx_len=dit_config.get("ctx_len", 0), infill_ctx_len=dit_config.get("infill_ctx_len", 0), shared_ctx=dit_config.get("shared_ctx", False), ) dit_model.eval() if dit_model.infill_ctx_len is None: dit_model.infill_ctx_len = 0 print("loading weights...") # TODO: gross hack, can be replaced with strict=True when we have new models model_sd_keys = set(dit_model.state_dict().keys()) checkpoint_sd_keys = set(state_dict.keys()) non_essential_keys = [ "semantic_conditioner.pos_embedding.inv_freq", "text_conditioner.pos_embedding.inv_freq", "ctx_conditioner.pos_embedding.inv_freq", ] for k in non_essential_keys: state_dict.pop(k, None) extra_keys = checkpoint_sd_keys - model_sd_keys missing_keys = model_sd_keys - checkpoint_sd_keys - set(non_essential_keys) assert len(extra_keys) == 0, f"extra keys in state dict: {extra_keys}" assert len(missing_keys) == 0, f"missing keys not in state dict: {missing_keys}" dit_model.load_state_dict(state_dict, strict=False) print(f"converting model to precision {weights_precision}...") convert_to_precision(dit_model, weights_precision=weights_precision) # Only move to CUDA and compile if compile=True # If compile=False (for critic model), keep on CPU for FSDP wrapping if compile: dit_model.to("cuda") dit_model = torch.compile(dit_model) return dit_model, dit_config def sample_timesteps_logsnr(batch_size, mean_logsnr=-1.2, std_logsnr=2.0): """ Sample timesteps for diffusion training by sampling logSNR values and converting to t. Args: batch_size (int): Number of timesteps to sample mean_logsnr (float): Mean of the logSNR Gaussian distribution std_logsnr (float): Standard deviation of the logSNR Gaussian distribution Returns: torch.Tensor: Tensor of shape (batch_size,) containing timestep values t in [0, 1] """ # Sample logSNR from Gaussian distribution logsnr = torch.randn(batch_size) * std_logsnr + mean_logsnr # Convert logSNR to timesteps using the logistic function # Since logSNR = ln((1-t)/t), we can solve for t: # t = 1 / (1 + exp(logsnr)) t = torch.sigmoid(-logsnr) # Clamp values to ensure numerical stability t = t.clamp(1e-4, 1 - 1e-4) return t def truncated_logistic_normal_rescaled(shape, left_trunc=0.075, right_trunc=1): """ shape: shape of the output tensor left_trunc: left truncation point, fraction of probability to be discarded right_trunc: right truncation boundary, should be 1 (never seen at test time) """ # Step 1: Sample from the logistic normal distribution (sigmoid of normal) logits = torch.randn(shape) # Step 2: Apply the CDF transformation of the normal distribution normal_dist = dist.Normal(0, 1) cdf_values = normal_dist.cdf(logits) # Step 3: Define the truncation bounds on the CDF lower_bound = normal_dist.cdf(torch.logit(torch.tensor(left_trunc))) upper_bound = normal_dist.cdf(torch.logit(torch.tensor(right_trunc))) # Step 4: Rescale linear CDF values into the truncated region (between lower_bound and upper_bound) truncated_cdf_values = lower_bound + (upper_bound - lower_bound) * cdf_values # Step 5: Map back to logistic-normal space using inverse CDF truncated_samples = torch.sigmoid(normal_dist.icdf(truncated_cdf_values)) # Step 6: Rescale values so that min is 0 and max is just below 1 rescaled_samples = (truncated_samples - left_trunc) / (right_trunc - left_trunc) return rescaled_samples def get_alphas_sigmas(t): """Returns the scaling factors for the clean image (alpha) and for the noise (sigma), given a timestep.""" return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2) class ConstantLRScheduler(torch.optim.lr_scheduler._LRScheduler): """Implements a constant learning rate schedule with a linear warmup. The learning rate increases linearly from 0 to the base learning rate during the warmup period. After warmup steps, the learning rate remains constant. Args: optimizer (Optimizer): Wrapped optimizer. warmup_steps (int): The number of steps for linear warmup. Default: 0. last_epoch (int): The index of last epoch. Default: -1. """ def __init__(self, optimizer, warmup_steps=0, last_epoch=-1): self.warmup_steps = warmup_steps super().__init__(optimizer, last_epoch) def get_lr(self): # If within warmup steps, apply linear warmup if self.last_epoch < self.warmup_steps: warmup_factor = (self.last_epoch + 1) / self.warmup_steps return [base_lr * warmup_factor for base_lr in self.base_lrs] # After warmup, return the constant learning rate return [base_lr for base_lr in self.base_lrs] class CosineLRScheduler(torch.optim.lr_scheduler._LRScheduler): """Implements a cosine learning rate schedule with a linear warmup. The learning rate increases linearly from 0 to the base learning rate during the warmup period. After warmup, it follows a cosine decay from 1 to 0 over the remaining steps. Args: optimizer (Optimizer): Wrapped optimizer. total_steps (int): The total number of steps in the schedule. warmup_steps (int): The number of steps for linear warmup. Default: 0. last_epoch (int): The index of last epoch. Default: -1. """ def __init__(self, optimizer, total_steps, warmup_steps=0, last_epoch=-1): self.warmup_steps = warmup_steps self.total_steps = total_steps super().__init__(optimizer, last_epoch) def get_lr(self): if self.last_epoch < self.warmup_steps: # Linear warmup warmup_factor = (self.last_epoch + 1) / self.warmup_steps return [base_lr * warmup_factor for base_lr in self.base_lrs] # Cosine decay progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps) cosine_factor = 0.5 * (1 + math.cos(math.pi * progress)) return [base_lr * cosine_factor for base_lr in self.base_lrs] def setup_distributed(master_addr, master_port): # Set environment variables based on parsed arguments os.environ["MASTER_ADDR"] = str(master_addr) os.environ["MASTER_PORT"] = str(master_port) if "SLURM_PROCID" in os.environ: # Running on SLURM if int(os.environ["SLURM_NTASKS_PER_NODE"]) != torch.cuda.device_count(): raise ValueError( f"SLURM_NTASKS_PER_NODE ({os.environ['SLURM_NTASKS_PER_NODE']}) does not match" f" the number of CUDA devices ({torch.cuda.device_count()}) on node {os.environ['HOSTNAME']}" ) rank = int(os.environ["SLURM_PROCID"]) local_rank = int(os.environ["SLURM_LOCALID"]) world_size = int(os.environ["SLURM_JOB_NUM_NODES"]) * int(os.environ["SLURM_NTASKS_PER_NODE"]) else: # Running locally rank = 0 local_rank = 0 world_size = 1 os.environ["RANK"] = str(rank) os.environ["LOCAL_RANK"] = str(local_rank) print(f"Initializing distributed process group on rank {local_rank}") torch.cuda.set_device(local_rank) try: dist.init_process_group( backend="nccl", timeout=timedelta(hours=6), rank=rank, world_size=world_size, device_id=torch.device(f"cuda:{local_rank}"), ) except Exception as e: print(f"Distributed error on rank {rank} with host {os.environ['HOSTNAME']}") raise e print(f"Done initializing on rank: {dist.get_rank()}") # barrier to check if nccl is working dist_barrier() print_with_time_master("distributed setup ready.") # Main training function def train( train_dataset, val_dataset, run_config, debug_mode=False, ): # Setup ddp_rank = int(os.environ["RANK"]) ddp_local_rank = int(os.environ["LOCAL_RANK"]) world_size = dist.get_world_size() group_size = min(world_size, 8) device = f"cuda:{ddp_local_rank}" torch.cuda.set_device(device) master_process = ddp_rank == 0 # Initialize timing profiler profiler = TimingProfiler(enabled=run_config["training"].get("profile_timing", False)) profile_print_every = run_config["training"].get("profile_print_every", 100) # Initialize wandb for logging (it will be disabled if in debug mode) if master_process: wandb.init( project=run_config["training"]["wandb_project"], name=run_config["training"]["wandb_name"], config={ "run_config": run_config, "world_size": world_size, "slurm_id": os.environ.get("SLURM_JOB_ID"), "slurm_name": os.environ.get("SLURM_JOB_NAME"), "slurm_script_path": os.environ.get("SLURM_SCRIPT_PATH"), "checkpoint_dir": CHECKPOINT_DIR, "pip_freeze": { dist.metadata["Name"]: dist.version for dist in importlib.metadata.distributions() }, "python_path": sys.executable, }, ) wandb.run.log_code(".") train_sampler = ( torch.utils.data.DistributedSampler( train_dataset, shuffle=True, rank=ddp_rank if not run_config["data"]["shard_data"] else ddp_local_rank, num_replicas=world_size if not run_config["data"]["shard_data"] else group_size, drop_last=True, seed=run_config["training"]["seed_offset"], ) if isinstance(train_dataset, GeneralMemmapMapDataset) else None ) val_sampler = ( torch.utils.data.DistributedSampler( val_dataset, shuffle=True, rank=ddp_rank, num_replicas=world_size, drop_last=True, seed=run_config["training"]["seed_offset"], ) if isinstance(val_dataset, GeneralMemmapMapDataset) else None ) # Prepare data loaders train_collate_fn = dynamic_collate_fn if isinstance(train_dataset, DynamicDataset) else None val_collate_fn = dynamic_collate_fn if isinstance(val_dataset, DynamicDataset) else None train_dataloader = DataLoader( train_dataset, batch_size=run_config["training"]["batch_size"], sampler=train_sampler, drop_last=True, num_workers=run_config["training"]["num_workers"], collate_fn=train_collate_fn, prefetch_factor=8, # Increased from 4 to prefetch more batches persistent_workers=True, worker_init_fn=lambda worker_id: np.random.seed( run_config["training"]["seed_offset"] + worker_id ), ) val_dataloader = DataLoader( val_dataset, batch_size=run_config["training"]["batch_size"], sampler=val_sampler, drop_last=True, num_workers=run_config["training"]["num_workers"], collate_fn=val_collate_fn, prefetch_factor=8, # Increased from 4 to prefetch more batches persistent_workers=True, ) def cached_dataloader_iter_fn(dataloader): # preload batch_store_size batches in memory. # #this way we make new batches every 10 steps, smoothing out the load batch_store = [] while True: if not batch_store: for _ in tqdm( range(run_config["training"]["batch_store_size"]), desc="preloading batches", disable=True, ): batch = next(dataloader) batch_store.append(batch) yield batch_store.pop(0) if run_config["data"]["use_cached_dataloader"]: train_dataloader = cached_dataloader_iter_fn(iter(train_dataloader)) val_dataloader = cached_dataloader_iter_fn(iter(val_dataloader)) # Initialize model print_with_time_master("setting up model...") model = DiffusionTransformer(**run_config["model"]) semantic_clusters_filepath = run_config["data"].get("semantic_clusters_filepath") model._initialize(semantic_clusters_filepath) global_step_offset = 0 if run_config["training"]["preload_checkpoint"]: print_with_time_master("preloading checkpoint...") ckpt = torch.load( run_config["training"]["preload_checkpoint"], weights_only=True, map_location="cpu" ) # verify model args for k in ckpt["model_args"].keys(): if k in run_config["model"] and run_config["training"]["check_model_args"]: assert ( ckpt["model_args"][k] == run_config["model"][k] ), f"model args mismatch for key {k} between config and checkpoint" # TODO: handling of ema model is a bit magical here if run_config["training"]["preload_optimizer"]: print_with_time_master("preloading optimizer...") optimizer_state = ckpt["optimizer"] state_dict = ckpt["model"] global_step_offset = ckpt["iter_num"] else: # Note: basic finetuning prefers ema model here, not obvious that's correct state_dict = ckpt.get("ema_model") or ckpt["model"] model.load_state_dict( state_dict, strict=True if run_config["training"]["check_model_args"] else False ) del state_dict num_params = sum(p.numel() for p in model.parameters()) num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print_with_time_master(f"Number of model parameters: {num_params:,}") print_with_time_master(f"Number of trainable parameters: {num_trainable_params:,}") # Initialize distillation-related models (only used if use_critic is True) teacher_model = None critic_model = None discriminator = None reward_model = None # load teacher model if run_config["training"].get("teacher_checkpoint"): print_with_time_master("loading teacher model...") teacher_model, teacher_config = _load_dit_model( run_config["training"]["teacher_checkpoint"], use_ema_if_exists=True, weights_precision=torch.bfloat16, compile=True, ) # teacher_model.eval() teacher_model.to(device) if run_config["training"]["use_critic"]: # print_with_time_master("creating fake score model from teacher model...") critic_model, critic_config = _load_dit_model( run_config["training"]["teacher_checkpoint"], use_ema_if_exists=True, weights_precision=torch.bfloat16, compile=False, # don't compile the critic because of intermediate features ) # Explicitly disable compilation on critic model to prevent compiled submodules critic_model.disable_compile() # create a small network to downsample the intermediate features print_with_time_master("creating discriminator...") if get_distill_config(run_config, "use_discriminator", False): discriminator = Discriminator(run_config["model"]["embed_dim"]) discriminator.to(device) # send to bfloat16 discriminator.to(torch.bfloat16) else: discriminator = None dist_barrier() print_with_time_master("wrapping in FSDP...") # Setup FSDP auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls={TransformerBlock} ) mixed_precision_policy = MixedPrecision( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, # Keep float32 for gradient precision buffer_dtype=torch.bfloat16, _module_classes_to_ignore=( ( ScaledSinusoidalEmbedding, RotaryEmbedding, ) if run_config["training"]["model_type"] == "default" else (RotaryEmbedding,) ), ) model = FSDP( model, auto_wrap_policy=auto_wrap_policy, mixed_precision=mixed_precision_policy, sharding_strategy=ShardingStrategy.HYBRID_SHARD, device_id=torch.cuda.current_device(), sync_module_states=True, use_orig_params=True, ) model.to(device) if run_config["training"]["use_critic"]: critic_model = FSDP( critic_model, auto_wrap_policy=auto_wrap_policy, mixed_precision=mixed_precision_policy, sharding_strategy=ShardingStrategy.HYBRID_SHARD, device_id=torch.cuda.current_device(), sync_module_states=True, use_orig_params=True, ) critic_model.to(device) dist_barrier() # set up ema ema_model = None if run_config["training"]["use_ema"]: print_with_time_master("initializing ema...") # 0.9995: 99.3% after 10k # 0.9999: 99.3% after 50k (64.2% after 10k) # 0.99998: 99.3% after 250k decays = [0.9999, 0.99998] update_every = [10, 10] warmup_steps = [50_000, 250_000] if run_config["training"]["preload_optimizer"]: warmup_steps = [1_000, 1_000] ema_model = FSDP_EMA( model, decays=decays, warmup_steps=warmup_steps, update_every=update_every, cpu_offload=False ) dist_barrier() # setup reward model reward_model_checkpoint = get_distill_config(run_config, "reward_model_checkpoint") if reward_model_checkpoint: print_with_time_master("loading reward model...") reward_model = load_ear_model(reward_model_checkpoint) reward_model.eval() else: reward_model = None dist_barrier() # add grad checkpointing and compile if needed if run_config["training"]["activation_checkpointing"]: print_with_time_master("applying grad checkpointing...") apply_fsdp_checkpointing(model) dist_barrier() # compile model model_copy_if_compiled = None should_compile = run_config["training"]["compile"] if should_compile and run_config["training"].get("use_critic", False): if master_process: print_with_time_master( "Skipping torch.compile for generator because distillation (use_critic=True) is enabled." ) should_compile = False if should_compile: print_with_time_master("compiling model...") model_copy_if_compiled = model # Increase cache limit to handle block_mask refresh with different batch sizes # This is needed because the block_mask is recreated when batch size changes, # causing different graph variants import torch._dynamo as dynamo dynamo.config.cache_size_limit = 64 # Increased from default 8 model = torch.compile(model, dynamic=False) if should_compile and run_config["training"]["preload_optimizer"]: # fix optimizer state if needed (add compile prefix) optimizer_state["state"] = { k if k.startswith("_orig_mod.") else "_orig_mod." + k: v for k, v in optimizer_state["state"].items() } for nn in range(len(optimizer_state["param_groups"])): optimizer_state["param_groups"][nn]["params"] = [ k if k.startswith("_orig_mod.") else "_orig_mod." + k for k in optimizer_state["param_groups"][nn]["params"] ] elif not should_compile and run_config["training"]["preload_optimizer"]: # remove compile prefix optimizer_state["state"] = { k.replace("_orig_mod.", ""): v for k, v in optimizer_state["state"].items() } for nn in range(len(optimizer_state["param_groups"])): optimizer_state["param_groups"][nn]["params"] = [ k.replace("_orig_mod.", "") for k in optimizer_state["param_groups"][nn]["params"] ] dist_barrier() # Setup optimizer and scheduler # Note: disabling weight decay for 1d layers showed similar accuracy; keeping things simple for now optimizer = AdamW( model.parameters(), lr=run_config["training"]["learning_rate"], betas=run_config["training"]["betas"], weight_decay=run_config["training"]["weight_decay"], fused=True, ) if run_config["training"]["use_critic"]: critic_optimizer = AdamW( list(critic_model.parameters()) + list(discriminator.parameters()) if get_distill_config(run_config, "use_discriminator", False) else critic_model.parameters(), lr=get_distill_config(run_config, "critic_lr", 5e-5), betas=run_config["training"]["betas"], weight_decay=run_config["training"]["weight_decay"], ) if run_config["training"]["preload_optimizer"]: print_with_time_master("preloading optimizer...") fsdp_optimizer_state = FSDP.optim_state_dict_to_load(model, optimizer, optimizer_state) optimizer.load_state_dict(fsdp_optimizer_state) del fsdp_optimizer_state, optimizer_state if run_config["training"]["lr_scheduler"] == "constant": lr_scheduler = ConstantLRScheduler(optimizer, warmup_steps=run_config["training"]["lr_warmup"]) elif run_config["training"]["lr_scheduler"] == "cosine": lr_scheduler = CosineLRScheduler( optimizer, run_config["training"]["tot_num_steps"], warmup_steps=run_config["training"]["lr_warmup"], ) else: raise ValueError(f"Unknown lr scheduler: {run_config['training']['lr_scheduler']}") torch.cuda.empty_cache() gc.collect() dist_barrier() print_with_time_master("finished model initialization.") # Training loop seed = ddp_rank + run_config["training"]["seed_offset"] rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=seed) rng_val = torch.quasirandom.SobolEngine(1, scramble=True, seed=seed) torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) global_step = global_step_offset epoch = 0 avg_val_loss = float("inf") best_val_loss = float("inf") if debug_mode: log_every = 10 val_every = 50 metrics_every = 20 else: finetune_mode = get_distill_config(run_config, "finetune_mode") if finetune_mode in ["distill", "adversarial"]: log_every = 1 val_every = 0 metrics_every = 10_000 else: log_every = 100 val_every = 1000 metrics_every = 10_000 train_start_time = time.time() distill_losses = {} while True: if train_sampler is not None: train_sampler.set_epoch(epoch) last_time = time.time() for batch in tqdm( train_dataloader, desc=f"{Fore.CYAN}Epoch {epoch + 1} - Training{Style.RESET_ALL}", disable=True, ): # Process the training batch once for both validation logging and training with profiler("01_total_batch_prep"): if run_config["data"]["dataset_type"] == "dynamic": processed_batch = prepare_batch_from_dynamic( batch, device, model, vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_mask_prob=run_config["data"].get("semantic_mask_prob", 0.0), cond_semantic_n_vocab=run_config["model"].get("cond_semantic_n_vocab", 4001), ctx_noise_level=run_config["data"]["noise_ctx"], ctx_mask_prob=run_config["data"].get("ctx_mask_prob", 0.0), infill_prob=run_config["data"]["infill_prob"], infill_min_ratio=run_config["data"]["infill_min_ratio"], infill_max_ratio=run_config["data"]["infill_max_ratio"], batch_reuse_factor=run_config["training"].get("batch_reuse_factor", 1), semantic_rate_hz=run_config["data"]["semantic_rate_hz"], profiler=profiler, ) else: processed_batch = prepare_batch_from_memmap( batch, device, model, run_config["training"]["use_fine_guidance"] ) if val_every > 0 and global_step % val_every == 0: avg_val_loss = validate( model_copy_if_compiled if model_copy_if_compiled is not None else model, val_dataloader, device, rng_val, run_config=run_config, ) tot_elapsed_time = time.time() - train_start_time log_validation_metrics( master_process, avg_val_loss, processed_batch, world_size, global_step, tot_elapsed_time, ) if run_config["training"]["use_critic"]: N = get_distill_config(run_config, "step_critic_N", 5) if global_step < get_distill_config(run_config, "critic_pretrain_steps", 0): step_critic = True else: step_critic = global_step % (N + 1) < N finetune_mode = get_distill_config(run_config, "finetune_mode") if finetune_mode == "distill": loss, distill_losses = train_step_distill( model, teacher_model, critic_model, processed_batch, device, rng, reward_model=reward_model, discriminator=discriminator, step_critic=step_critic, residual=get_distill_config(run_config, "residual", False), noise_schedule=run_config["training"]["noise_schedule"], diffusion_objective=run_config["training"]["diffusion_objective"], t_discretize=get_distill_config(run_config, "t_discretize", 1), text_cfg_scale=get_distill_config(run_config, "text_cfg_scale", 1.0), use_reward_model=global_step > get_distill_config(run_config, "reward_model_pretrain_steps", 0), ) else: raise ValueError(f"Unknown finetune mode: {finetune_mode}") std_data = distill_losses["x_std"] acc_steps = run_config["training"]["acc_steps"] grad_norm = 0 if step_critic: if global_step % acc_steps == 0: critic_optimizer.zero_grad() (loss / acc_steps).backward() if (global_step + 1) % acc_steps == 0: if get_distill_config(run_config, "use_discriminator", False): grad_norm = torch.nn.utils.clip_grad_norm_( list(critic_model.parameters()) + list(discriminator.parameters()), run_config["training"]["clip_grad_norm"], ) else: grad_norm = torch.nn.utils.clip_grad_norm_( list(critic_model.parameters()), run_config["training"]["clip_grad_norm"], ) distill_losses["critic_grad_norm"] = grad_norm critic_optimizer.step() lr_scheduler.step() else: if global_step % acc_steps == 0: optimizer.zero_grad() (loss / acc_steps).backward() if (global_step + 1) % acc_steps == 0: grad_norm = torch.nn.utils.clip_grad_norm_( list(model.parameters()), run_config["training"]["clip_grad_norm"] ) distill_losses["gen_grad_norm"] = grad_norm optimizer.step() lr_scheduler.step() else: with profiler("02_train_step_total"): loss, std_data = train_step( model, processed_batch, device, rng, noise_schedule=run_config["training"]["noise_schedule"], diffusion_objective=run_config["training"]["diffusion_objective"], profiler=profiler, ) with profiler("03_optimizer_step"): with profiler("03a_zero_grad"): optimizer.zero_grad() with profiler("03b_backward"): loss.backward() with profiler("03c_clip_grad"): grad_norm = model.clip_grad_norm_( max_norm=run_config["training"]["clip_grad_norm"] ) with profiler("03d_optimizer_step"): optimizer.step() with profiler("03e_lr_scheduler"): lr_scheduler.step() with profiler("04_ema_update"): if ema_model is not None: ema_model.update(step=global_step) # Only sync EMA every N steps to avoid overhead if global_step % 10 == 0: dist_barrier() # Logging and validation if global_step % log_every == 0: tot_elapsed_time = time.time() - train_start_time # Pass distill_losses if in distillation mode finetune_mode = get_distill_config(run_config, "finetune_mode", None) log_training_metrics( master_process, loss, lr_scheduler, epoch, tot_elapsed_time, (time.time() - last_time) / log_every, processed_batch, world_size, global_step, std_data, grad_norm, distill_losses=distill_losses if finetune_mode == "distill" else {}, ) last_time = time.time() # Print profiling stats periodically if profiler.enabled and master_process and global_step % profile_print_every == 0: profiler.print_stats(prefix=f"[Profiling Step {global_step}]") profiler.reset() if global_step % 1_000 == 0: gc.collect() if ( global_step - global_step_offset > 0 and global_step % run_config["training"]["ckpt_every"] == 0 ): best_val_loss = save_checkpoint( CHECKPOINT_DIR, model, ema_model, optimizer, best_val_loss, avg_val_loss, 250_000, run_config=run_config, model_args=run_config["model"], iter_num=global_step, ) if ( run_config["training"]["compute_metrics"] and global_step - global_step_offset > 0 and global_step % metrics_every == 0 ): print_with_time_master("Computing metrics...") metrics_start_time = time.time() try: compute_metrics( model_copy_if_compiled if model_copy_if_compiled is not None else model, val_dataloader, vae_scale_factor=run_config["data"]["vae_scale_factor"], device=device, num_batches=3, global_step=global_step, dataset_type=run_config["data"]["dataset_type"], ) except Exception as e: print(f"Error computing metrics: {e}") metrics_time = time.time() - metrics_start_time print_with_time_master(f"Time to compute metrics: {metrics_time:.2f} seconds") # signals the profiler that the next profiling step has started if torch_profiler: torch_profiler.step() if memory_profiler: memory_profiler.step() global_step += 1 if global_step >= run_config["training"]["tot_num_steps"]: break epoch += 1 if global_step >= run_config["training"]["tot_num_steps"]: break print_with_time_master(f"Done. Epoch {epoch + 1} completed. Total steps: {global_step}.") # Cleanup dist_barrier() dist.destroy_process_group() wandb.finish() # Batch preparation functions def prepare_batch_from_memmap(batch, device, model, use_fine_guidance=False): """ Extract data from old memmap format (diffusion_input, info dict) and move to device. Returns standardized dict with all training data. """ diffusion_input, info = batch result = { "diffusion_input": diffusion_input.to(device), "semantic_codes": None, "padding_mask": info["padding_mask"].to(device), "text_codes": info["text_codes"].to(device), "ctx_vae": None, "ctx_mask": None, "infill_ctx_vae": None, "infill_ctx_mask": None, "stem_ctx_vae": None, "stem_ctx_mask": None, "semantic_noise": None, } if "semantic_codes" in info: result["semantic_codes"] = info["semantic_codes"].to(device) if use_fine_guidance: result["w_scale"] = info["w_scale"].to(device) result["w_cond"] = info["w_cond"].to(device) result["uncond_text_codes"] = info["uncond_text_codes"].to(device) if (model.ctx_len != 0 or model.shared_ctx) and "ctx_vae" in info: result["ctx_vae"] = info["ctx_vae"].to(device) result["ctx_mask"] = info["ctx_mask"].to(device) if (model.infill_ctx_len != 0 or model.shared_ctx) and "infill_ctx_vae" in info: result["infill_ctx_vae"] = info["infill_ctx_vae"].to(device) result["infill_ctx_mask"] = info["infill_ctx_mask"].to(device) if model.stem_ctx_len != 0 and "stem_ctx_vae" in info: result["stem_ctx_vae"] = info["stem_ctx_vae"].to(device) result["stem_ctx_mask"] = info["stem_ctx_mask"].to(device) if "semantic_noise" in info: result["semantic_noise"] = info["semantic_noise"].to(device) return result def _codec_encode_audio_to_tensor(audio_list, device, scale_factor): """Helper to codec encode audio (VAE) and convert to scaled tensor on device.""" encoded = codec_encode(audio_list, normalize_volume=False) if isinstance(encoded, list): encoded = torch.stack([torch.from_numpy(x) if isinstance(x, np.ndarray) else x for x in encoded]) elif isinstance(encoded, np.ndarray): encoded = torch.from_numpy(encoded) return encoded.to(device) * scale_factor def _semantic_encode_audio_to_tensor(audio_list, device, n_codebooks=1): """Helper to semantic encode audio and convert to long tensor on device.""" semantic_codes_list = [] for audio in audio_list: sem_codes = semantic_encode(audio) if isinstance(sem_codes, np.ndarray): sem_codes = torch.from_numpy(sem_codes) sem_codes = sem_codes[:, :n_codebooks] if n_codebooks > 1 else sem_codes[..., 0] semantic_codes_list.append(sem_codes) # only use the n codebook levels return torch.stack(semantic_codes_list).to(device).long() # [batch, sequence, n_codebooks] def _apply_infill_augmentation( vae_target, padding_mask, infill_ctx_vae, infill_ctx_mask, infill_prob, infill_min_ratio, infill_max_ratio, ctx_indices=None, vox_indices=None, ctx_noise_level=0.0, ): """ Apply infill augmentation to a batch on a per-item basis. For each sample independently, with probability infill_prob, randomly selects a contiguous region (left/center/right) from vae_target, copies it to infill_ctx_vae, and masks it out from loss computation. Only applies infill to items that don't have ctx_vae or vox_vae. """ batch_size, n_vae_tokens = vae_target.shape[0], vae_target.shape[1] # Convert indices to sets for faster lookup ctx_indices_set = set(ctx_indices) if ctx_indices else set() vox_indices_set = set(vox_indices) if vox_indices else set() # add some optional noise to the infill region, like we do for ctx if ctx_noise_level > 0.0: noise = torch.randn_like(vae_target) * ctx_noise_level * random.random() vae_target = vae_target + noise for batch_idx in range(batch_size): # Skip this sample if it has ctx_vae or vox_vae if batch_idx in ctx_indices_set or batch_idx in vox_indices_set: continue # Skip this sample with probability (1 - infill_prob) if random.random() >= infill_prob: continue # Randomly choose infill region: 'left', 'center', or 'right' infill_position = random.choice(["left", "center", "right"]) # Determine infill region size infill_ratio = random.uniform(infill_min_ratio, infill_max_ratio) infill_length = int(n_vae_tokens * infill_ratio) # Calculate start and end indices based on position if infill_position == "left": start_idx, end_idx = 0, infill_length elif infill_position == "center": start_idx = (n_vae_tokens - infill_length) // 2 end_idx = start_idx + infill_length else: # 'right' start_idx = n_vae_tokens - infill_length end_idx = n_vae_tokens # Extract infill region from vae_target infill_ctx_vae[batch_idx, start_idx:end_idx] = vae_target[batch_idx, start_idx:end_idx] infill_ctx_mask[batch_idx, :] = True # Update padding_mask to mask out the infill region (we don't compute loss there) padding_mask[batch_idx, start_idx:end_idx] = False def prepare_batch_from_dynamic( batch, device, model, vae_scale_factor=2.5, semantic_mask_prob=0.0, cond_semantic_n_vocab=4001, ctx_noise_level=0.0, ctx_mask_prob=0.0, infill_prob=0.1, infill_min_ratio=0.2, infill_max_ratio=0.5, batch_reuse_factor=1, semantic_rate_hz=25, profiler=None, ): """ Encode Audio objects from DynamicDataset format and move to device. Returns standardized dict with all training data. The function expects audio_target to be 2 seconds longer than the target chunk size (1 second padding on each side). After encoding, it crops semantic_rate_hz tokens from the start and end of both codec and semantic encoded vectors to remove the padding. Args: batch_reuse_factor: Number of times to reuse each batch with different noise levels infill_prob: Probability of doing infill when no ctx is present infill_min_ratio: Minimum ratio of sequence length for infill region infill_max_ratio: Maximum ratio of sequence length for infill region ctx_mask_prob: Probability of masking the context semantic_mask_prob: Probability of masking the semantic codes cond_semantic_n_vocab: Number of semantic tokens ctx_noise_level: Noise level for the context semantic_rate_hz: Semantic token rate in Hz (used to determine crop size) """ # Use a dummy profiler if none provided if profiler is None: profiler = TimingProfiler(enabled=False) with profiler("01a_extract_batch"): ( audio_target_list, audio_ctx_list, audio_vox_list, text_codes_list, raw_text_list, audio_target_24k_list, ) = batch # Encode audio targets (codec VAE encoding) with profiler("01b_encode_vae_target"): vae_target = _codec_encode_audio_to_tensor(audio_target_list, device, vae_scale_factor) # Encode semantic targets (24kHz mono audio) with profiler("01c_encode_semantic"): semantic_codes = _semantic_encode_audio_to_tensor( audio_target_24k_list, device, n_codebooks=model.semantic_conditioner.n_codebooks, ) # Crop tokens corresponding to CROP_PADDING_S seconds from start and end with profiler("01d_crop_tokens"): # The dataloader ensures proper padding, so we can safely crop crop_tokens = int(CROP_PADDING_S * semantic_rate_hz) assert vae_target.shape[1] >= 2 * crop_tokens, ( f"vae_target has {vae_target.shape[1]} tokens, need at least {2 * crop_tokens} " f"(CROP_PADDING_S={CROP_PADDING_S}s * semantic_rate_hz={semantic_rate_hz} * 2)" ) vae_target = vae_target[:, crop_tokens:-crop_tokens, :] assert semantic_codes.shape[1] >= 2 * crop_tokens, ( f"semantic_codes has {semantic_codes.shape[1]} tokens, need at least {2 * crop_tokens} " f"(CROP_PADDING_S={CROP_PADDING_S}s * semantic_rate_hz={semantic_rate_hz} * 2)" ) if semantic_codes.ndim == 2: semantic_codes = semantic_codes[:, crop_tokens:-crop_tokens] else: # ndim == 3 semantic_codes = semantic_codes[:, crop_tokens:-crop_tokens, :] # Mask semantic codes with profiler("01e_mask_semantic"): if np.random.random() < semantic_mask_prob: if np.random.random() < 0.5: # 50% chance to apply ratio mask # uniformly sample a ratio between 0 and 1 semantic_mask_ratio = random.uniform(0, 1) semantic_codes = ratio_mask_semantic_codes(semantic_codes, semantic_mask_ratio) else: # 50% chance to replace tokens on the right with PAD token # randomy select the mask size, max is 500 tokens mask_size = random.randint(0, 500) start_idx = semantic_codes.shape[1] - mask_size semantic_codes[:, start_idx:] = cond_semantic_n_vocab - 1 # Stack text codes text_codes = torch.stack(text_codes_list).to(device) # Get batch dimensions batch_size, n_vae_tokens, vae_dim = vae_target.shape # Encode context audio (may be None for some samples) ctx_vae = torch.zeros((batch_size, n_vae_tokens, vae_dim), dtype=vae_target.dtype, device=device) ctx_mask = torch.zeros((batch_size, n_vae_tokens), device=device).bool() ctx_indices = [idx for idx, a in enumerate(audio_ctx_list) if a is not None] if ctx_indices: ctx_batch = [audio_ctx_list[idx] for idx in ctx_indices] encoded_ctx = _codec_encode_audio_to_tensor(ctx_batch, device, vae_scale_factor) # Crop tokens corresponding to CROP_PADDING_S seconds from start and end (same as target) assert encoded_ctx.shape[1] >= 2 * crop_tokens, ( f"encoded_ctx has {encoded_ctx.shape[1]} tokens, need at least {2 * crop_tokens} " f"(CROP_PADDING_S={CROP_PADDING_S}s * semantic_rate_hz={semantic_rate_hz} * 2)" ) encoded_ctx = encoded_ctx[:, crop_tokens:-crop_tokens, :] # Scatter ctx results into full batch tensors for k, global_idx in enumerate(ctx_indices): # randomize the noise level from 0 to ctx_noise_level noise_scale = ctx_noise_level * random.random() noise = torch.randn_like(encoded_ctx[k]) * noise_scale encoded_ctx[k] = encoded_ctx[k] + noise # Set mask to 1.0 (use all context by default) ctx_mask[global_idx] = 1.0 # Apply context masking augmentation with probability ctx_mask_prob if np.random.random() < ctx_mask_prob: ctx_length = encoded_ctx[k].shape[1] mask_size = random.randint(0, min(500, ctx_length)) # Mask up to 500 tokens encoded_ctx[k][:, :mask_size] = 0.0 ctx_vae[global_idx] = encoded_ctx[k] # Create padding mask - per-token for dynamic dataset (batch_size, n_tokens) padding_mask = torch.ones(batch_size, n_vae_tokens, device=device).bool() # Encode vox audio (may be None for some samples) vox_vae = torch.zeros((batch_size, n_vae_tokens, vae_dim), dtype=vae_target.dtype, device=device) vox_mask = torch.zeros((batch_size, n_vae_tokens), device=device).bool() vox_indices = [idx for idx, a in enumerate(audio_vox_list) if a is not None] if vox_indices: vox_batch = [audio_vox_list[idx] for idx in vox_indices] encoded_vox = _codec_encode_audio_to_tensor(vox_batch, device, vae_scale_factor) # Scatter vox results into full batch tensors for k, global_idx in enumerate(vox_indices): vox_vae[global_idx] = encoded_vox[k] vox_mask[global_idx] = 1.0 # Set full vox mask # Apply infill augmentation: X% of samples (independently) when they have no ctx or vox # Initialize infill ctx vae and mask infill_ctx_vae = torch.zeros_like(ctx_vae) infill_ctx_mask = torch.zeros(batch_size, n_vae_tokens, device=device).bool() # Apply infill augmentation before reshaping (when tensors are in correct shape) # Apply infill on a per-item basis for items that don't have ctx_vae or vox_vae with profiler("01g_apply_augmentation"): if infill_prob > 0: _apply_infill_augmentation( vae_target.clone(), padding_mask, infill_ctx_vae, infill_ctx_mask, infill_prob, infill_min_ratio, infill_max_ratio, ctx_indices=ctx_indices, vox_indices=vox_indices, ctx_noise_level=ctx_noise_level, ) # do some reshaping vae_target = vae_target.permute(0, 2, 1) ctx_vae = ctx_vae.permute(0, 2, 1) infill_ctx_vae = infill_ctx_vae.permute(0, 2, 1) vox_vae = vox_vae.permute(0, 2, 1) # Repeat tensors for batch reuse if batch_reuse_factor > 1: vae_target = vae_target.repeat(batch_reuse_factor, 1, 1) if semantic_codes.ndim == 2: semantic_codes = semantic_codes.repeat(batch_reuse_factor, 1) elif semantic_codes.ndim == 3: semantic_codes = semantic_codes.repeat(batch_reuse_factor, 1, 1) else: raise ValueError(f"Invalid semantic codes dimension: {semantic_codes.ndim}") padding_mask = padding_mask.repeat(batch_reuse_factor, 1) text_codes = text_codes.repeat(batch_reuse_factor, 1) if ctx_vae is not None: ctx_vae = ctx_vae.repeat(batch_reuse_factor, 1, 1) if ctx_mask is not None: ctx_mask = ctx_mask.repeat(batch_reuse_factor, 1) if infill_ctx_vae is not None: infill_ctx_vae = infill_ctx_vae.repeat(batch_reuse_factor, 1, 1) if infill_ctx_mask is not None: infill_ctx_mask = infill_ctx_mask.repeat(batch_reuse_factor, 1) if vox_vae is not None: vox_vae = vox_vae.repeat(batch_reuse_factor, 1, 1) if vox_mask is not None: vox_mask = vox_mask.repeat(batch_reuse_factor, 1) # stem_ctx_vae and stem_ctx_mask are always None in dynamic dataset # Create empty text codes for unconditional generation (used in distillation) # Use pad_idx from text conditioner or default to cond_text_n_vocab - 1 pad_idx = ( getattr(model.text_conditioner, "pad_idx", None) if hasattr(model, "text_conditioner") else None ) if pad_idx is None: # Default: use the last token ID as pad (typically cond_text_n_vocab - 1) pad_idx = model.cond_text_n_vocab - 1 empty_text_codes = torch.full_like(text_codes, pad_idx) return { "diffusion_input": vae_target, "semantic_codes": semantic_codes, "padding_mask": padding_mask, "text_codes": text_codes, "empty_text_codes": empty_text_codes, "ctx_vae": ctx_vae, # Always tensor (zeros when no ctx, mask controls usage) "ctx_mask": ctx_mask, # Always tensor (zeros when no ctx) "infill_ctx_vae": infill_ctx_vae, "infill_ctx_mask": infill_ctx_mask, "stem_ctx_vae": None, "stem_ctx_mask": None, "vox_vae": vox_vae, "vox_mask": vox_mask, } # Training and validation steps def train_step( model, prepared_batch, device, rng, noise_schedule="default", diffusion_objective="v", profiler=None, ): # prepared_batch is already processed and contains all necessary data # Use a dummy profiler if none provided if profiler is None: profiler = TimingProfiler(enabled=False) # Extract prepared data diffusion_input = prepared_batch["diffusion_input"] semantic_codes = prepared_batch["semantic_codes"] padding_mask = prepared_batch["padding_mask"] text_codes = prepared_batch["text_codes"] ctx_vae = prepared_batch["ctx_vae"] ctx_mask = prepared_batch["ctx_mask"] infill_ctx_vae = prepared_batch["infill_ctx_vae"] infill_ctx_mask = prepared_batch["infill_ctx_mask"] stem_ctx_vae = prepared_batch["stem_ctx_vae"] stem_ctx_mask = prepared_batch["stem_ctx_mask"] vox_vae = prepared_batch["vox_vae"] vox_mask = prepared_batch["vox_mask"] batch_size = diffusion_input.shape[0] if noise_schedule == "default" or noise_schedule == "early": # Draw uniformly distributed continuous timesteps t = rng.draw(batch_size)[:, 0].to(device).to(torch.bfloat16) elif noise_schedule == "logit_normal": # Draw from a logit-normal distribution t = torch.sigmoid(torch.randn(batch_size, device=device).to(torch.bfloat16)) elif noise_schedule == "trunc_logit_normal": t = truncated_logistic_normal_rescaled(batch_size).to(device).to(torch.bfloat16) # Flip the distribution t = 1 - t elif noise_schedule == "log_snr": t = sample_timesteps_logsnr(batch_size).to(device).to(torch.bfloat16) else: raise ValueError(f"Invalid noise schedule: {noise_schedule}") # Replace 1% of t with ones to ensure training on terminal SNR with profiler("02b_prepare_noise_schedule"): pct_ones = 0.2 if noise_schedule == "early" else 0.01 t = torch.where(torch.rand_like(t) < pct_ones, torch.ones_like(t), t) # Calculate the noise schedule parameters for those timesteps if diffusion_objective in ["v"]: alphas, sigmas = get_alphas_sigmas(t) elif diffusion_objective in ["rectified_flow", "rf_denoiser"]: alphas, sigmas = 1 - t, t else: raise ValueError(f"Invalid diffusion objective: {diffusion_objective}") diffusion_input = diffusion_input.to(device) diffusion_input = diffusion_input.to(t.dtype) # Combine the ground truth data and the noise alphas = alphas[:, None, None] sigmas = sigmas[:, None, None] noise = torch.randn_like(diffusion_input) noised_inputs = diffusion_input * alphas + noise * sigmas if diffusion_objective == "v": targets = noise * alphas - diffusion_input * sigmas elif diffusion_objective in ["rectified_flow", "rf_denoiser"]: targets = noise - diffusion_input else: raise ValueError(f"Invalid diffusion objective: {diffusion_objective}") with profiler("02c_model_forward"): v = model.forward( noised_inputs, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) # Calculate L2 loss between combined prediction and target with profiler("02d_compute_loss"): loss = F.mse_loss(v, targets, reduction="none").mean(dim=1) loss = loss[padding_mask].mean() return loss, diffusion_input.std() def diffuse(x, dist, noise_schedule, rng, batch_size, device, diffusion_objective, t_discretize=1): if dist == "training": if noise_schedule == "default" or noise_schedule == "early": # Draw uniformly distributed continuous timesteps t = rng.draw(batch_size)[:, 0].to(device).to(torch.bfloat16) elif noise_schedule == "logit_normal": # Draw from a logit-normal distribution t = torch.sigmoid(torch.randn(batch_size, device=device).to(torch.bfloat16)) elif noise_schedule == "trunc_logit_normal": t = truncated_logistic_normal_rescaled(batch_size).to(device).to(torch.bfloat16) # Flip the distribution t = 1 - t elif noise_schedule == "log_snr": t = sample_timesteps_logsnr(batch_size).to(device).to(torch.bfloat16) else: raise ValueError(f"Invalid noise schedule: {noise_schedule}") pct_ones = 0.2 if noise_schedule == "early" else 0.01 elif dist == "inference": # use early noise schedule for inference # For inference, use linearly spaced timesteps if False: if t_discretize == 1: t = torch.ones(batch_size, device=device).to(torch.bfloat16) else: t = torch.linspace(0, 1, t_discretize + 1, dtype=torch.bfloat16, device=device) # Select random timesteps from the discretized schedule indices = torch.randint(0, t_discretize + 1, (batch_size,), device=device) t = t[indices] else: t = rng.draw(batch_size)[:, 0].to(device).to(torch.bfloat16) pct_ones = 0.2 # schedule = [1.0, 0.5] ## randomly select the schedule values the number of times in the batch else: raise ValueError(f"Invalid diffusion distribution: {dist}") # Replace % of t with ones to ensure training on terminal SNR t = torch.where(torch.rand_like(t) < pct_ones, torch.ones_like(t), t) # Calculate the noise schedule parameters for those timesteps if diffusion_objective in ["v"]: alphas, sigmas = get_alphas_sigmas(t) elif diffusion_objective in ["rectified_flow", "rf_denoiser"]: alphas, sigmas = 1 - t, t else: raise ValueError(f"Invalid diffusion objective: {diffusion_objective}") x = x.to(device) x = x.to(t.dtype) # Combine the ground truth data and the noise alphas = alphas[:, None, None] sigmas = sigmas[:, None, None] noise = torch.randn_like(x) noised_inputs = x * alphas + noise * sigmas if diffusion_objective == "v": targets = noise * alphas - x * sigmas elif diffusion_objective in ["rectified_flow", "rf_denoiser"]: targets = noise - x else: raise ValueError(f"Invalid diffusion objective: {diffusion_objective}") return noised_inputs, targets, t, sigmas, noise def train_step_distill( model, teacher_model, critic_model, prepared_batch, device, rng, reward_model: torch.nn.Module = None, discriminator: torch.nn.Module = None, noise_schedule="default", diffusion_objective="v", text_cfg_scale=1.0, residual: bool = False, step_critic: bool = True, nu_1: float = 0.1, # default 0.01 nu_2: float = 0.1, # was 0.005 nu_3: float = 0.001, t_discretize: int = 1, use_reward_model: bool = False, ): # Extract prepared data - matching train_step exactly diffusion_input = prepared_batch["diffusion_input"] batch_size = diffusion_input.shape[0] semantic_codes = prepared_batch["semantic_codes"] padding_mask = prepared_batch["padding_mask"] text_codes = prepared_batch["text_codes"] ctx_vae = prepared_batch["ctx_vae"] ctx_mask = prepared_batch["ctx_mask"] infill_ctx_vae = prepared_batch["infill_ctx_vae"] infill_ctx_mask = prepared_batch["infill_ctx_mask"] stem_ctx_vae = prepared_batch["stem_ctx_vae"] stem_ctx_mask = prepared_batch["stem_ctx_mask"] vox_vae = prepared_batch["vox_vae"] vox_mask = prepared_batch["vox_mask"] empty_text_codes = prepared_batch.get("empty_text_codes") text_cfg_scale = prepared_batch.get("text_cfg_scale", 1.0) if not step_critic: # train the generator noised_inputs, targets, t, sigmas, noise = diffuse( diffusion_input, "inference", noise_schedule, rng, batch_size, device, diffusion_objective, t_discretize=t_discretize, ) # one-step generator inference x_gen = model.forward( noised_inputs, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) if residual: x_gen = noised_inputs - t[:, None, None] * x_gen latents = x_gen x_noise, targets, t, sigma, noise = diffuse( x_gen, "training", noise_schedule, rng, batch_size, device, diffusion_objective, ) pred_fake_noise = critic_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) pred_fake_latents = x_noise - pred_fake_noise with torch.no_grad(): pred_cond = teacher_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) # if text_cfg_scale > 1.0: pred_text_uncond = teacher_model( x_noise, t, text_codes=empty_text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) # print(f"text_cfg_scale: {text_cfg_scale}, ({text_cfg_scale.shape})") pred_real_noise = pred_cond + (pred_cond - pred_text_uncond) * (text_cfg_scale - 1.0) # else: # pred_real_noise = pred_cond pred_real_latents = x_noise - pred_real_noise # Compute distributional error p_real = latents - pred_real_latents p_fake = latents - pred_fake_latents denom = torch.abs(p_real).mean(dim=[1, 2], keepdim=True) # per batch item normalization grad = (p_real - p_fake) / denom grad = torch.nan_to_num(grad) dmd_loss = 0.5 * torch.nn.functional.mse_loss( latents.float(), (latents - grad).detach().float(), reduction="none" ).mean(dim=1) dmd_loss = dmd_loss[padding_mask].mean() # dmd_loss = F.mse_loss(fake_denoised, real_denoised, reduction="none").mean(dim=1) # dmd_loss = dmd_loss[padding_mask].mean() # compute adversarial loss if discriminator is not None: # if False: x_noise, targets, t, sigma, noise = diffuse( x_gen, "training", noise_schedule, rng, batch_size, device, diffusion_objective, ) d_in = critic_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, intermediate_layer_idx=16, ) d_out = discriminator(d_in) g_loss = F.mse_loss(d_out, torch.ones_like(d_out)) loss = dmd_loss + g_loss * nu_1 else: loss = dmd_loss g_loss = 0 if reward_model is not None: reward = reward_model(x_gen.float().permute(0, 2, 1)).mean() # higher reward is better if use_reward_model: loss = loss - reward * nu_3 else: reward = 0 return loss, { "gen_loss": loss, "dmd_loss": dmd_loss, "g_loss": g_loss, "x_gen_std": x_gen.std(), "x_std": diffusion_input.std(), "reward": reward, } else: # train the critic and discriminator with torch.no_grad(): # generate noised inputs using inference distribution noised_inputs, targets, t, sigmas, noise = diffuse( diffusion_input, "inference", # "inference", temporary noise_schedule, rng, batch_size, device, diffusion_objective, t_discretize=t_discretize, ) # one-step generator inference x_gen = model.forward( noised_inputs, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) if residual: x_gen = noised_inputs - t[:, None, None] * x_gen # compute adversarial loss if discriminator is not None: x_noise, targets, t, sigma, noise = diffuse( diffusion_input, "training", noise_schedule, rng, batch_size, device, diffusion_objective, ) d_in = critic_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, intermediate_layer_idx=16, ) d_out = discriminator(d_in) real_loss = F.mse_loss(d_out, torch.ones_like(d_out)) x_noise, targets, t, sigma, noise = diffuse( x_gen.detach(), "training", noise_schedule, rng, batch_size, device, diffusion_objective, ) d_in = critic_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, intermediate_layer_idx=16, ) d_out = discriminator(d_in) fake_loss = F.mse_loss(d_out, torch.zeros_like(d_out)) d_loss = real_loss + fake_loss # critic (dsm loss) # if False: x_noise, targets, t, sigma, noise = diffuse( x_gen.detach(), "training", noise_schedule, rng, batch_size, device, diffusion_objective, ) x_denoised = critic_model( x_noise, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, ) dsm_loss = F.mse_loss(x_denoised, targets, reduction="none").mean(dim=1) dsm_loss = dsm_loss[padding_mask].mean() # dsm_loss = dsm_loss.mean() if discriminator is not None: loss = dsm_loss + nu_2 * d_loss else: loss = dsm_loss d_loss = 0 real_loss = 0 fake_loss = 0 return loss, { "critic_loss": loss, "dsm_loss": dsm_loss, "d_loss": d_loss, "d_real_loss": real_loss, "d_fake_loss": fake_loss, "x_gen_std": x_gen.std(), "x_gen_mean": x_gen.mean(), "x_std": diffusion_input.std(), "x_mean": diffusion_input.mean(), } def all_reduce_metric(metric, device): """Helper function to all-reduce a metric across all GPUs.""" metric_tensor = torch.tensor(metric).to(device) dist.all_reduce(metric_tensor, op=dist.ReduceOp.SUM) return metric_tensor.item() @torch.no_grad() # this causes an error with compiled models def validate(model, val_dataloader, device, rng, run_config=None): total_val_loss = 0.0 num_batches = 0 for batch in val_dataloader: if num_batches >= 50: break # Process validation batch if run_config["data"]["dataset_type"] == "dynamic": processed_batch = prepare_batch_from_dynamic( batch, device, model, vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_mask_prob=0.0, # Disable semantic masking during validation cond_semantic_n_vocab=run_config["model"].get("cond_semantic_n_vocab", 4001), ctx_noise_level=run_config["data"]["noise_ctx"], ctx_mask_prob=run_config["data"]["ctx_mask_prob"], infill_prob=run_config["data"]["infill_prob"], infill_min_ratio=run_config["data"]["infill_min_ratio"], infill_max_ratio=run_config["data"]["infill_max_ratio"], batch_reuse_factor=run_config["training"].get("batch_reuse_factor", 1), semantic_rate_hz=run_config["data"]["semantic_rate_hz"], ) else: processed_batch = prepare_batch_from_memmap( batch, device, model, run_config["training"]["use_fine_guidance"] ) loss, std_data = train_step( model, processed_batch, device, rng, noise_schedule=run_config["training"]["noise_schedule"], diffusion_objective=run_config["training"]["diffusion_objective"], ) total_val_loss += loss.item() num_batches += 1 # All-reduce the total loss and number of batches across all GPUs total_val_loss = all_reduce_metric(total_val_loss, device) num_batches = all_reduce_metric(num_batches, device) if num_batches == 0: print("No batches to validate on?") return 0 avg_val_loss = total_val_loss / num_batches return avg_val_loss @torch.no_grad() def compute_metrics( model, val_dataloader, vae_scale_factor=2.5, device="cuda", num_batches=1, global_step=0, dataset_type="general_memmap", ): """Compute metrics for the validation set.""" is_master = dist.get_rank() == 0 total_stft_loss = 0.0 total_mel_loss = 0.0 total_samples = 0 total_upload = 3 n_uploaded = 0 for i, batch in enumerate(val_dataloader): if i >= num_batches: break # Prepare data based on dataset type from config if dataset_type == "dynamic": # New dynamic dataset format - need to encode audio_target_list, audio_ctx_list, text_codes_list, audio_target_24k_list = batch # Encode audio targets (VAE encoding) vae_target = codec_encode(audio_target_list, normalize_volume=False) if isinstance(vae_target, list): vae_target = torch.stack( [torch.from_numpy(x) if isinstance(x, np.ndarray) else x for x in vae_target] ) elif isinstance(vae_target, np.ndarray): vae_target = torch.from_numpy(vae_target) diffusion_input = vae_target.to(device) * vae_scale_factor # Encode semantic targets semantic_codes_list = [] for audio_24k in audio_target_24k_list: sem_codes = semantic_encode(audio_24k) if isinstance(sem_codes, np.ndarray): sem_codes = torch.from_numpy(sem_codes) semantic_codes_list.append(sem_codes) semantic_codes = torch.stack(semantic_codes_list).to(device).long() # Stack text codes text_codes = torch.stack(text_codes_list).to(device) # Encode context audio batch_size = diffusion_input.shape[0] n_vae_tokens = diffusion_input.shape[1] vae_dim = diffusion_input.shape[2] ctx_vae = torch.zeros( (batch_size, n_vae_tokens, vae_dim), dtype=diffusion_input.dtype, device=device ) ctx_mask = torch.zeros( (batch_size, n_vae_tokens, 1), dtype=diffusion_input.dtype, device=device ) ctx_indices = [idx for idx, a in enumerate(audio_ctx_list) if a is not None] if ctx_indices: ctx_batch = [audio_ctx_list[idx] for idx in ctx_indices] encoded_ctx = codec_encode(ctx_batch, normalize_volume=False) if isinstance(encoded_ctx, list): encoded_ctx = torch.stack( [torch.from_numpy(x) if isinstance(x, np.ndarray) else x for x in encoded_ctx] ) elif isinstance(encoded_ctx, np.ndarray): encoded_ctx = torch.from_numpy(encoded_ctx) encoded_ctx = encoded_ctx.to(device) * vae_scale_factor for k, global_idx in enumerate(ctx_indices): ctx_vae[global_idx] = encoded_ctx[k] ctx_mask[global_idx] = 1.0 else: # Old memmap format diffusion_input, info = batch semantic_codes = info["semantic_codes"].to(device) text_codes = info["text_codes"].to(device) ctx_vae = info["ctx_vae"].to(device) ctx_mask = info["ctx_mask"].to(device) diffusion_input = diffusion_input.to(device) out_pred_z = simple_generate( model, semantic_codes, text_codes, ctx_vae, ctx_mask, n_steps=8, ) try: # decode latents input_audio_len = int( round(semantic_codes.shape[1] / model.cond_semantic_len * model.block_size) ) scaled_pred_z = out_pred_z[..., :input_audio_len] / vae_scale_factor pred_audio = [codec_decode(latent.T) for latent in scaled_pred_z] clean_audio = [ codec_decode(diffusion_input[i].T / vae_scale_factor) for i in range(len(pred_audio)) ] # Calculate STFT loss for each audio pair stft_losses = [ calculate_stft_loss(clean, pred) for clean, pred in zip(clean_audio, pred_audio) ] mel_losses = [ calculate_mel_loss(clean, pred) for clean, pred in zip(clean_audio, pred_audio) ] total_stft_loss += sum(stft_losses) total_mel_loss += sum(mel_losses) total_samples += len(pred_audio) # save audio as mp3 if is_master: with tempfile.TemporaryDirectory() as tmp_dir: for i, audio in enumerate(pred_audio): if n_uploaded >= total_upload: break mp3_filename = f"generated_{i}.mp3" mp3_path = os.path.join(tmp_dir, mp3_filename) # save audio in a tmp dir for upload (using tmpdir import) audio.write_hq_mp3(mp3_path) audio_dict = { f"generated_audio/{i}": wandb.Audio( mp3_path, caption=f"Generated {global_step}" ) } audio_dict["generated_audio/global_step"] = global_step wandb.log(audio_dict, commit=True) n_uploaded += 1 # TODO: hacky fix for wandb upload to make sure it's finished time.sleep(5) except Exception as e: # if latents are random can lead to fail print(f"metrics failed with error: {e}.") dist_barrier() # All-reduce the total loss and number of samples across all GPUs total_stft_loss = all_reduce_metric(total_stft_loss, device) total_mel_loss = all_reduce_metric(total_mel_loss, device) total_samples = all_reduce_metric(total_samples, device) # Calculate average STFT loss across all batches avg_stft_loss = 0 avg_mel_loss = 0 if total_samples > 0: avg_stft_loss = total_stft_loss / total_samples avg_mel_loss = total_mel_loss / total_samples if is_master: wandb.log( { "avg_stft_loss": avg_stft_loss, "avg_mel_loss": avg_mel_loss, }, step=global_step, ) print_with_time_master(f"Average STFT loss across all batches: {avg_stft_loss}") print_with_time_master(f"Average mel loss across all batches: {avg_mel_loss}") dist_barrier() return avg_stft_loss def parse_args(): parser = argparse.ArgumentParser(description="Train the diffusion model") parser.add_argument("--master_addr", type=str, default="localhost", help="Master node address") parser.add_argument("--master_port", type=str, default="12355", help="Master node port") parser.add_argument("--debug", action="store_true", help="Enable debug mode") # Profiling arguments parser.add_argument("--enable_profiling", action="store_true", help="Enable profiling") parser.add_argument( "--dump_folder", type=str, default="/app/suno/diffusion_profiling/", help="Folder to dump profiling data", ) parser.add_argument( "--save_traces_folder", type=str, default="traces", help="Folder to save profiling traces", ) parser.add_argument("--profile_freq", type=int, default=20, help="Profiling frequency") parser.add_argument("--enable_memory_snapshot", action="store_true", help="Enable memory snapshot") parser.add_argument( "--save_memory_snapshot_folder", type=str, default="memory_snapshots", help="Folder to save memory snapshots", ) parser.add_argument("--config_path", type=str, default=None, help="Path to config file to override") return parser.parse_args() # Main execution if __name__ == "__main__": args = parse_args() setup_distributed(args.master_addr, args.master_port) master_process = int(os.environ["RANK"]) == 0 # load config and run some sanity checks run_config = load_and_update_config(DEFAULT_CONFIG, config_path=args.config_path) assert run_config["model"]["block_size"] % run_config["model"]["io_hz"] == 0 assert run_config["model"]["cond_semantic_len"] % run_config["data"]["semantic_rate_hz"] == 0 # assert (run_config["model"]["block_size"] / run_config["model"]["io_hz"]) == round( # run_config["model"]["cond_semantic_len"] / run_config["data"]["semantic_rate_hz"] # ) # some overrides if run_config["data"]["shard_data"] is None and not run_config["data"]["dataset_type"] == "dynamic": # shard if >= 8 nodes run_config["data"]["shard_data"] = dist.get_world_size() >= 8 * 8 if run_config["training"]["model_type"] == "default": from model import DiffusionTransformer from base import ( ScaledSinusoidalEmbedding, TransformerBlock, RotaryEmbedding, apply_fsdp_checkpointing, ) elif run_config["training"]["model_type"] == "prefix": from prefix_model.model import DiffusionTransformer from prefix_model.base import TransformerBlock, RotaryEmbedding, apply_fsdp_checkpointing else: raise ValueError(f"Unknown model type: {run_config['model']['model_type']}") # Disable wandb if in debug mode if args.debug: os.environ["WANDB_MODE"] = "disabled" # Preload codec and semantic models for dynamic dataset (needed for on-the-fly encoding) if run_config["data"]["dataset_type"] == "dynamic": print_with_time_master("Preloading codec and semantic models for DynamicDataset...") # Verify required config fields are present required_fields = [ "codec_filepath", "semantic_model_filepath", "semantic_clusters_filepath", "train_metas_filepath", "val_metas_filepath", ] for field in required_fields: if field not in run_config["data"]: raise ValueError( f"Missing required config field 'data.{field}' for dataset_type='dynamic'" ) codec_filepath = run_config["data"]["codec_filepath"] semantic_model_filepath = run_config["data"]["semantic_model_filepath"] semantic_clusters_filepath = run_config["data"]["semantic_clusters_filepath"] preload_codec_models(checkpoint_filepath=codec_filepath, device="cuda") preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath, device="cuda") print_with_time_master("Models preloaded successfully") # Also preload for compute_metrics if needed if run_config["training"]["compute_metrics"]: if run_config["model"]["io_hz"] == 100 and run_config["model"]["io_channels"] == 128: codec_filepath = os.path.join(run_config["data"]["codec_dir"], "100hz_vae_peaq_kl_0.005.pth") elif run_config["model"]["io_hz"] == 25 and run_config["model"]["io_channels"] == 128: codec_filepath = os.path.join(run_config["data"]["codec_dir"], "25hz_vae_peaq_kl_0.005.pth") elif run_config["model"]["io_hz"] == 25 and run_config["model"]["io_channels"] == 64: codec_filepath = os.path.join( run_config["data"]["codec_dir"], "25hz_vae_peaq_64_kl_0.005.pth" ) else: raise NotImplementedError("codec not supported") preload_codec_models(checkpoint_filepath=codec_filepath, device="cuda") dist_barrier() # tone down some logging unless in debug mode: if not args.debug: # Suppress specific module logs logging.getLogger("torch.distributed.fsdp._wrap_utils").setLevel(logging.ERROR) logging.getLogger("torch.fx.experimental.symbolic_shapes").setLevel(logging.ERROR) # Filter specific warnings warnings.filterwarnings( "ignore", message="Graph break due to unsupported builtin flash_attn_2_cuda.PyCapsule.fwd", category=UserWarning, module="torch._dynamo.variables.functions", ) warnings.filterwarnings( "ignore", message="Both mixed precision and an auto_wrap_policy were specified to FSDP", category=UserWarning, module="torch.distributed.fsdp._wrap_utils", ) warnings.filterwarnings( "ignore", message="Profiler function will be ignored", category=UserWarning, module="torch._logging._internal", ) if run_config["data"]["shard_data"]: print_with_time_master("sharding data...") # shard data memmap_fnames = [] for s in [ "val_semantic_memmap_filename", "val_metas_filename", "val_vae_memmap_filename", "val_info_filename", "train_vae_memmap_filename", "train_semantic_memmap_filename", "train_metas_filename", "train_info_filename", ]: fn = run_config["data"].get(s, None) if fn is not None: memmap_fnames.append(fn) shard_data( run_config["data"]["dataset_dir"], run_config["data"]["shard_data_dir"], memmap_fnames, vae_n_tokens_memmap=run_config["model"]["block_size"] * run_config["data"]["patch_size"], vae_dim=run_config["data"]["vae_dim"], semantic_n_tokens_memmap=run_config["model"]["cond_semantic_len"], allow_shard_reuse=run_config["data"]["allow_shard_reuse"], ) dist_barrier() # Initialize datasets print_with_time_master("loading datasets...") if run_config["data"]["dataset_type"] == "general_memmap": dataset_val = GeneralMemmapMapDataset( mode=run_config["data"]["mode"], dataset_dir=( run_config["data"]["shard_data_dir"] if run_config["data"]["shard_data"] else run_config["data"]["dataset_dir"] ), vae_memmap_filename=run_config["data"]["val_vae_memmap_filename"], semantic_memmap_filename=run_config["data"]["val_semantic_memmap_filename"], metas_filename=run_config["data"]["val_metas_filename"], info_filename=run_config["data"]["val_info_filename"], vae_dim=run_config["data"]["vae_dim"], vae_n_tokens=run_config["model"]["block_size"] * run_config["data"]["patch_size"], semantic_n_tokens=run_config["model"]["cond_semantic_len"], cond_text_len=run_config["model"]["cond_text_len"], vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_pad_token=run_config["model"]["cond_semantic_n_vocab"] - 1, ctx_len=run_config["model"]["ctx_len"], aligned_text_prob=run_config["data"]["aligned_text_prob"], patch_size=run_config["data"]["patch_size"], foreign_weight=run_config["data"]["foreign_weight"], is_training=False, scale_vae_ctx=run_config["data"]["scale_vae_ctx"], always_pad_semantic=run_config["data"]["always_pad_semantic"], semantic_rate_hz=run_config["data"]["semantic_rate_hz"], prev_vae_ctx=run_config["data"]["prev_vae_ctx"], infill_vae_ctx=run_config["data"]["infill_vae_ctx"], noise_ctx=run_config["data"]["noise_ctx"], always_skip_semantic=run_config["data"]["always_skip_semantic"], semantic_skip_factors=run_config["data"]["semantic_skip_factors"], respell_augment_prob=run_config["data"]["respell_augment_prob"], shared_ctx=run_config["model"]["shared_ctx"], semantic_dropout=run_config["data"]["semantic_dropout"], ) elif run_config["data"]["dataset_type"] == "dynamic": # DynamicDataset handles splitting across DDP ranks and workers internally dataset_val = DynamicDataset( metas=run_config["data"]["val_metas_filepath"], audio_chunk_s=run_config["data"]["audio_chunk_s"], audio_ctx_s=run_config["data"]["audio_ctx_s"], audio_vox_s=run_config["data"]["audio_vox_s"], cond_text_len=run_config["model"]["cond_text_len"], text_drop_prob=run_config["data"]["text_drop_prob"], target_loudness_db=run_config["data"]["target_loudness_db"], use_stem_prob=run_config["data"]["use_stem_prob"], use_vox_prob=run_config["data"]["use_vox_prob"], use_text_aligned_prob=run_config["data"]["use_text_aligned_prob"], foreign_weight=run_config["data"]["foreign_weight"], text_aligned_weight=run_config["data"]["text_aligned_weight"], stem_weight=run_config["data"]["stem_weight"], is_training=False, ) elif run_config["data"]["dataset_type"] == "stems": n_vae_tokens = (run_config["model"]["block_size"] + run_config["model"]["ctx_len"]) * run_config[ "data" ]["patch_size"] audio_dataset_val = BundleDownloaderDataset( dataset_dir=run_config["data"]["dataset_dir"], metas_filename=run_config["data"]["val_metas_filename"], duration_s=(run_config["model"]["block_size"] + run_config["model"]["ctx_len"]) / run_config["model"]["io_hz"], max_multi_instruments=run_config["data"]["max_multi_instruments"], ) audio_dataset_val_dl = DataLoader( audio_dataset_val, batch_size=None, num_workers=4, prefetch_factor=40 ) dataset_val = StemMemmapMapDataset( dataset_dir=( run_config["data"]["shard_data_dir"] if run_config["data"]["shard_data"] else run_config["data"]["dataset_dir"] ), bundle_iter=iter(audio_dataset_val_dl), metas_filename=run_config["data"]["val_metas_filename"], vae_dim=run_config["data"]["vae_dim"], vae_n_tokens=n_vae_tokens, semantic_n_tokens=run_config["model"]["cond_semantic_len"], cond_text_len=run_config["model"]["cond_text_len"], vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_pad_token=run_config["model"]["cond_semantic_n_vocab"] - 1, ctx_len=run_config["model"]["ctx_len"], patch_size=run_config["data"]["patch_size"], is_training=False, scale_vae_ctx=run_config["data"]["scale_vae_ctx"], max_multi_instruments=run_config["data"]["max_multi_instruments"], ) print_with_time_master(f"Loaded validation dataset on rank {os.environ['RANK']}.") # Use validation set as training set if in debug mode if args.debug: dataset_train = dataset_val elif run_config["data"]["dataset_type"] == "general_memmap": dataset_train = GeneralMemmapMapDataset( mode=run_config["data"]["mode"], dataset_dir=( run_config["data"]["shard_data_dir"] if run_config["data"]["shard_data"] else run_config["data"]["dataset_dir"] ), vae_memmap_filename=run_config["data"]["train_vae_memmap_filename"], semantic_memmap_filename=run_config["data"]["train_semantic_memmap_filename"], metas_filename=run_config["data"]["train_metas_filename"], info_filename=run_config["data"]["train_info_filename"], vae_dim=run_config["data"]["vae_dim"], vae_n_tokens=run_config["model"]["block_size"] * run_config["data"]["patch_size"], semantic_n_tokens=run_config["model"]["cond_semantic_len"], cond_text_len=run_config["model"]["cond_text_len"], vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_pad_token=run_config["model"]["cond_semantic_n_vocab"] - 1, ctx_len=run_config["model"]["ctx_len"], aligned_text_prob=run_config["data"]["aligned_text_prob"], patch_size=run_config["data"]["patch_size"], foreign_weight=run_config["data"]["foreign_weight"], is_training=True, scale_vae_ctx=run_config["data"]["scale_vae_ctx"], always_pad_semantic=run_config["data"]["always_pad_semantic"], noise_ctx=run_config["data"]["noise_ctx"], semantic_rate_hz=run_config["data"]["semantic_rate_hz"], semantic_skip_factors=run_config["data"]["semantic_skip_factors"], always_skip_semantic=run_config["data"]["always_skip_semantic"], prev_vae_ctx=run_config["data"]["prev_vae_ctx"], infill_vae_ctx=run_config["data"]["infill_vae_ctx"], semantic_noise_level=run_config["data"]["semantic_noise_level"], shared_ctx=run_config["model"]["shared_ctx"], semantic_dropout=run_config["data"]["semantic_dropout"], ) elif run_config["data"]["dataset_type"] == "dynamic": # DynamicDataset handles splitting across DDP ranks and workers internally dataset_train = DynamicDataset( metas=run_config["data"]["train_metas_filepath"], audio_chunk_s=run_config["data"]["audio_chunk_s"], audio_ctx_s=run_config["data"]["audio_ctx_s"], audio_vox_s=run_config["data"]["audio_vox_s"], cond_text_len=run_config["model"]["cond_text_len"], text_drop_prob=run_config["data"]["text_drop_prob"], target_loudness_db=run_config["data"]["target_loudness_db"], use_stem_prob=run_config["data"]["use_stem_prob"], use_vox_prob=run_config["data"]["use_vox_prob"], use_text_aligned_prob=run_config["data"]["use_text_aligned_prob"], foreign_weight=run_config["data"]["foreign_weight"], text_aligned_weight=run_config["data"]["text_aligned_weight"], stem_weight=run_config["data"]["stem_weight"], is_training=True, ) elif run_config["data"]["dataset_type"] == "stems": n_vae_tokens = (run_config["model"]["block_size"] + run_config["model"]["ctx_len"]) * run_config[ "data" ]["patch_size"] audio_dataset_train = BundleDownloaderDataset( dataset_dir=run_config["data"]["dataset_dir"], metas_filename=run_config["data"]["train_metas_filename"], duration_s=(run_config["model"]["block_size"] + run_config["model"]["ctx_len"]) / run_config["model"]["io_hz"], max_multi_instruments=run_config["data"]["max_multi_instruments"], ) audio_dataset_train_dl = DataLoader( audio_dataset_train, batch_size=None, num_workers=4, prefetch_factor=40 ) dataset_train = StemMemmapMapDataset( dataset_dir=( run_config["data"]["shard_data_dir"] if run_config["data"]["shard_data"] else run_config["data"]["dataset_dir"] ), bundle_iter=iter(audio_dataset_train_dl), metas_filename=run_config["data"]["train_metas_filename"], vae_dim=run_config["data"]["vae_dim"], vae_n_tokens=n_vae_tokens, semantic_n_tokens=run_config["model"]["cond_semantic_len"], cond_text_len=run_config["model"]["cond_text_len"], vae_scale_factor=run_config["data"]["vae_scale_factor"], semantic_pad_token=run_config["model"]["cond_semantic_n_vocab"] - 1, ctx_len=run_config["model"]["ctx_len"], patch_size=run_config["data"]["patch_size"], is_training=True, scale_vae_ctx=run_config["data"]["scale_vae_ctx"], max_multi_instruments=run_config["data"]["max_multi_instruments"], ) print_with_time_master(f"Loaded training dataset on rank {os.environ['RANK']}.") dist_barrier() print_with_time_master("finished data loading.") # Start training with ( maybe_enable_profiling( args.enable_profiling, args.dump_folder, args.save_traces_folder, args.profile_freq, global_step=0, # Assuming iter_num is not defined, set to 0 ) as torch_profiler, maybe_enable_memory_snapshot( args.enable_memory_snapshot, args.dump_folder, args.save_memory_snapshot_folder, args.profile_freq, global_step=0, # Assuming iter_num is not defined, set to 0 ) as memory_profiler, ): train( dataset_train, dataset_val, run_config, debug_mode=args.debug, )