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 boto3 import subprocess from collections import defaultdict 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 DPOGeneralMemmapMapDataset, shard_data 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, ) from audio_metrics import calculate_stft_loss, calculate_mel_loss from generation import simple_generate from suno_utils.tasks.dac_vae_100hz_peaq import ( # NOTE: works for 25hz as well preload_models as preload_codec_models, decode as codec_decode, ) gc.disable() torch.set_float32_matmul_precision("high") # Set the TOKENIZERS_PARALLELISM environment variable to False to avoid warning os.environ["TOKENIZERS_PARALLELISM"] = "false" # base - 24 layers, 24 heads, 1536 dim # large - 32 layers, 32 heads, 2048 dim DEFAULT_CONFIG = { "training": { "wandb_name": "8n_prefix", "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": "cosine", "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, "ipo_beta": 1.0, "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, "check_model_args": True, "preload_checkpoint": None, "preload_optimizer": False, "revert_label": False, "semantic_chunk_frac": 0.2, # 20% of the time, we mask tail semantic codes "t_discretize": 12, # number of descrete steps of t "ave_latent_frac": 0.0, # percentage of times average latent as input "ctx_vae_noise": 1.0, # percentage of times add noise to ctx_vae "ctx_vae_noise_frac": 1.0, # fraction of noise augmentation "ctx_vae_mask_frac": 0.5, # fraction of ctx_vae masked fully "freeze_embeddings": False, "grad_accum_steps": 1, }, "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, }, "data": { "mode": "dpo", "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_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, "scale_vae_ctx": False, "prev_vae_ctx": True, "infill_vae_ctx": False, "always_pad_semantic": False, "noise_ctx": 0.0, "semantic_noise_level": 0.0, "always_skip_semantic": False, "aligned_text_prob": 0.5, "max_multi_instruments": 1, "respell_augment_prob": 0.0, "shared_ctx": False, }, } 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 t_steps_lookup = { 4: [0.98828, 0.94141, 0.74219, 0.29492], 6: [0.98828, 0.96875, 0.92188, 0.80469, 0.57031, 0.29492], 8: [0.98828, 0.97656, 0.95312, 0.91016, 0.82422, 0.6875, 0.48828, 0.29492], 12: [ 0.98828, 0.98047, 0.96875, 0.95312, 0.92969, 0.89453, 0.84375, 0.76953, 0.67188, 0.54688, 0.41602, 0.29492, ], } def preference_loss( policy_chosen_logps: torch.Tensor, policy_rejected_logps: torch.Tensor, reference_chosen_logps: torch.Tensor, reference_rejected_logps: torch.Tensor, beta: float, label_smoothing: float = 0.0, ipo: bool = True, reference_free: bool = False, kto: bool = False, ): """Compute the DPO loss for a batch of policy and reference model log probabilities. Args: policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,) reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,) beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0. label_smoothing: conservativeness for DPO loss, which assumes that preferences are noisy (flipped with probability label_smoothing) ipo: If True, use the IPO loss instead of the DPO loss. reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses. Returns: A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the DPO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively. """ if not kto: pi_logratios = policy_chosen_logps - policy_rejected_logps # win - lose ref_logratios = reference_chosen_logps - reference_rejected_logps # win - lose if reference_free: ref_logratios = 0 # for diffusion, this is regression loss, unlike negative log likelcihood logits = -(pi_logratios - ref_logratios) # also known as h_{\pi_\theta}^{y_w,y_l} if ipo: losses = (logits - 1 / (2 * beta)) ** 2 # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf else: # Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf) losses = ( -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing ) elif kto: chosen_logratios = policy_chosen_logps - reference_chosen_logps rejected_logratios = policy_rejected_logps - reference_rejected_logps chosen_KL = chosen_logratios.mean().clamp(min=0) rejected_KL = rejected_logratios.mean().clamp(min=0) # simplified from the original KTO loss losses = -F.sigmoid(beta * (chosen_logratios - rejected_KL)) - F.sigmoid( beta * (chosen_KL - rejected_logratios) ) # these are the same chosen_rewards = -beta * (policy_chosen_logps - reference_chosen_logps).detach() rejected_rewards = -beta * (policy_rejected_logps - reference_rejected_logps).detach() return losses, chosen_rewards, rejected_rewards # Utility functions def update_nested_dict(original, updates): for key, value in updates.items(): if key not in original: print( f"Invalid key '{key}' found in updates. It doesn't exist in the default configuration." ) continue 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 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) 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 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 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"), "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(".") # Prepare data loaders 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"], ) 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"], ) train_dataloader = DataLoader( train_dataset, batch_size=run_config["training"]["batch_size"], sampler=train_sampler, drop_last=True, num_workers=4, ) val_dataloader = DataLoader( val_dataset, batch_size=run_config["training"]["batch_size"], sampler=val_sampler, drop_last=True, ) # Initialize model print_with_time_master("setting up model...") model = DiffusionTransformer(**run_config["model"]) model._initialize(semantic_centroid_path="/app/suno/models/semantic/mert_25_2x4k.npy") ref_model = DiffusionTransformer(**run_config["model"]) ref_model._initialize(semantic_centroid_path="/app/suno/models/semantic/mert_25_2x4k.npy") 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"]: if ckpt["model_args"][k] != run_config["model"][k]: print( f"model args mismatch for key {k} between config and checkpoint, config: {run_config['model'][k]}, checkpoint: {ckpt['model_args'][k]}" ) # TODO: handling of ema model is a bit magical here if run_config["training"]["preload_optimizer"]: if master_process: print("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=False) ref_model.load_state_dict(state_dict, strict=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:,}") 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, 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"]["freeze_embeddings"]: embed_names = [ "vae_pad_embed", "vae_infill_pad_embed", "vae_stem_pad_embed", "vae_default_embed", ] for param_name, param in model.named_parameters(): if any(name in param_name for name in embed_names): print(f"Freezing {param_name}") param.requires_grad = False ref_model = FSDP( ref_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, ) ref_model.to(device) ref_model.requires_grad_(False) ref_model.eval() 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() # 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 if run_config["training"]["compile"]: print_with_time_master("compiling model...") model_copy_if_compiled = model model = torch.compile(model, dynamic=False) ref_model = torch.compile(ref_model, dynamic=False) if run_config["training"]["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 run_config["training"]["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"]["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 rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=ddp_rank) rng_val = torch.quasirandom.SobolEngine(1, scramble=True, seed=ddp_rank) torch.manual_seed(ddp_rank) random.seed(ddp_rank) np.random.seed(ddp_rank) global_step = global_step_offset epoch = 0 avg_val_loss = float("inf") best_val_loss = float("inf") log_every = 10 if debug_mode else 100 val_every = 50 if debug_mode else 1000 metrics_every = 20 if debug_mode else 10_000 train_start_time = time.time() while True: 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, ): # do the last val log if global_step % val_every == 0: avg_val_loss, dpo_losses = validate( model_copy_if_compiled if model_copy_if_compiled is not None else model, val_dataloader, device, rng_val, ref_model, run_config, ) tot_elapsed_time = time.time() - train_start_time log_validation_metrics( master_process, avg_val_loss, batch, world_size, global_step, tot_elapsed_time, dpo_losses, ) loss, std_data, dpo_losses = train_step(model, batch, device, rng, ref_model, run_config) # Scale loss for gradient accumulation loss = loss / run_config["training"]["grad_accum_steps"] # Backpropagate loss.backward() # Only clip gradients and update model after accumulation steps if global_step % run_config["training"]["grad_accum_steps"] == 0: grad_norm = model.clip_grad_norm_(max_norm=run_config["training"]["clip_grad_norm"]) optimizer.step() optimizer.zero_grad() lr_scheduler.step() if ema_model is not None: ema_model.update(step=global_step) dist_barrier() # Logging and validation if global_step % log_every == 0: tot_elapsed_time = time.time() - train_start_time log_training_metrics( master_process, loss, lr_scheduler, epoch, tot_elapsed_time, (time.time() - last_time) / log_every, batch, world_size, global_step, std_data, grad_norm, dpo_losses, ) last_time = time.time() if global_step % 1_000 == 0: gc.collect() can_save_checkpoint = ( global_step > 0 and (global_step + 1) % run_config["training"]["ckpt_every"] == 0 ) end_of_checkpoint = (global_step + 1) == run_config["training"]["tot_num_steps"] if can_save_checkpoint or end_of_checkpoint: best_val_loss = save_checkpoint( CHECKPOINT_DIR, model, ema_model, optimizer, best_val_loss, avg_val_loss, step_save_iters=run_config["training"]["ckpt_every"], run_config=run_config, model_args=run_config["model"], iter_num=(global_step + 1), save_best_ckpt=False, save_last_ckpt=True, ) if end_of_checkpoint: avg_val_loss, dpo_losses = validate( model_copy_if_compiled if model_copy_if_compiled is not None else model, val_dataloader, device, rng_val, ref_model, run_config, ) tot_elapsed_time = time.time() - train_start_time log_validation_metrics( master_process, avg_val_loss, batch, world_size, global_step, tot_elapsed_time, dpo_losses, ) 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=1, global_step=global_step, ) 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"] * run_config["training"]["grad_accum_steps"] ): break epoch += 1 if ( global_step >= run_config["training"]["tot_num_steps"] * run_config["training"]["grad_accum_steps"] ): break print_with_time_master( f"Done. Epoch {epoch + 1} completed. Total steps: {global_step}. Optimizer steps: {global_step // run_config['training']['grad_accum_steps']}." ) # Cleanup dist_barrier() dist.destroy_process_group() if master_process: # launch modal eval run_name = wandb.run.name launch_eval(run_name, objective=run_config["training"]["objective"]) wandb.finish() def launch_eval(run_name: str, objective: str): # Upload checkpoint local_ckpt_filepath = os.path.join(CHECKPOINT_DIR, "last_ckpt_infer.pt") if not os.path.exists(local_ckpt_filepath): raise FileNotFoundError(f"Last checkpoint file not found at {local_ckpt_filepath}") s3_client = boto3.client("s3") s3_client.upload_file( local_ckpt_filepath, "suno-data", "christian/checkpoints/diffusion/modal_eval_ckpt.pt" ) print("Uploaded checkpoint to s3.") # Launch modal eval print(f"Launching modal eval for run {run_name} with objective {objective}") modal_path = "/home/christian/miniconda3/envs/suno_diff/bin/modal" script_path = "/home/christian/code/neon/sunoDiff/modal_eval.py" try: result = subprocess.run( [ modal_path, "run", script_path, "--run-name", run_name, "--objective", objective, ], capture_output=True, text=True, check=True, ) print("Modal eval succeeded.") print("STDOUT:", result.stdout) except subprocess.CalledProcessError as e: print("Modal eval failed.") print("STDERR:", e.stderr) # copy results to local # now download the output from s3 to local and compute the statistics try: result = subprocess.run( [ "aws", "s3", "sync", "s3://suno-data/christian/outputs/auk-clips-up-u-2/", "/app2/suno/data/eval_outputs/auk-clips-up-u-2/", ], capture_output=True, text=True, check=True, # This raises CalledProcessError if the script fails ) print("S3 sync succeeded.") except subprocess.CalledProcessError as e: print("S3 sync failed.") print("STDERR:", e.stderr) # Handle failure as needed (log, raise, etc.) # Training and validation steps def train_step(model, batch, device, rng, ref_model, run_config): diffusion_input, info = batch noise_schedule = run_config["training"]["noise_schedule"] diffusion_objective = run_config["training"]["objective"] # print_with_time_master(f" diff_input: {diffusion_input.shape}") # TODO: this is a hack to get the batch size right, stupid diffusion_input_emb_dim = diffusion_input.shape[1] // 2 diffusion_neg_input = diffusion_input[:, :diffusion_input_emb_dim, :] diffusion_pos_input = diffusion_input[:, diffusion_input_emb_dim:, :] diffusion_ave_both_input = (diffusion_neg_input + diffusion_pos_input) / 2 diffusion_input = torch.repeat_interleave(diffusion_neg_input, repeats=2, dim=0) diffusion_ave_input = torch.repeat_interleave(diffusion_ave_both_input, repeats=2, dim=0) diffusion_input[1::2] = diffusion_pos_input semantic_codes = info["semantic_codes"].to(device) padding_mask = info["padding_mask"].to(device) # b, t text_codes = info["text_codes"].to(device) # b, t # print_with_time_master(f"info: {info['idx']}") # for dpo we have to load all the even indices assert torch.all(info["idx"] % 2 == 0), "All indices should be even" # print_with_time_master( # f" padding_mask: {padding_mask.shape}, semantic_codes: {semantic_codes.shape}, text_codes: {text_codes.shape}" # ) # print_with_time_master(f"text_codes: {text_codes.shape}") # TODO: add text code asserts and semantic code asserts # if master_process: # # print_with_time_master(f"diffusion_input: {diffusion_input.shape}") # np.save( # f"/home/tony/Data/test_npz/diffusion_input_{info['idx']}.npy", diffusion_input.cpu().numpy() # ) batch_size = diffusion_input.shape[0] ctx_vae = None # b, d, t ctx_mask = None # b, t if model.ctx_len is not None or model.shared_ctx: original_ctx_vae = info["ctx_vae"].to(device) original_ctx_mask = info["ctx_mask"].to(device) ctx_neg_vae = original_ctx_vae[:, :diffusion_input_emb_dim, :] ctx_pos_vae = original_ctx_vae[:, diffusion_input_emb_dim:, :] if (random.random() < run_config["training"]["ctx_vae_noise_frac"]) and ( ctx_vae_noise := run_config["training"]["ctx_vae_noise"] > 0 ): # augment with additional ctx_vae noise -- but make pos and neg has the same additional noise ctx_vae_additional_noise = torch.randn_like(ctx_neg_vae) * random.random() * ctx_vae_noise # find the zero value batch, and don't add noise to it # zero_mask: (batch,) True if all elements in 2nd and 3rd dims are zero zero_mask = torch.all(ctx_neg_vae == 0, dim=(1, 2)) ctx_vae_additional_noise[zero_mask] = 0 ctx_neg_vae += ctx_vae_additional_noise ctx_pos_vae += ctx_vae_additional_noise ctx_vae = torch.repeat_interleave(ctx_neg_vae, repeats=2, dim=0) ctx_vae[1::2] = ctx_pos_vae ctx_mask = torch.repeat_interleave(original_ctx_mask, repeats=2, dim=0) # print(f"ctx_mask: {ctx_mask.shape}") # print(f"ctx_vae: {ctx_vae.shape}") # we also can augment the ctx_mask if run_config["training"]["semantic_chunk_frac"] > 0: for ctx_batch_idx in range(batch_size // 2): dice_roll = random.random() if dice_roll < run_config["training"]["semantic_chunk_frac"]: n = random.randint(25 * 5, run_config["model"]["cond_semantic_len"] - 1) ctx_vae[2 * ctx_batch_idx : 2 * (ctx_batch_idx + 1), :, :n] = 0 ctx_mask[2 * ctx_batch_idx : 2 * (ctx_batch_idx + 1), :n] = 0 # set everything to empty, so treat later chunks as the first one if random.random() < run_config["training"]["ctx_vae_mask_frac"]: ctx_vae = torch.zeros_like(ctx_vae) ctx_mask = torch.zeros(ctx_mask.shape).bool() if run_config["training"]["semantic_chunk_frac"] > 0: # augment semantic codes with mask for semantic_batch_idx in range(batch_size // 2): dice_roll = random.random() if dice_roll < run_config["training"]["semantic_chunk_frac"]: # augment with mask -- so we don't forget how to infer n = random.randint(5, run_config["model"]["cond_semantic_len"] - 1) semantic_codes[semantic_batch_idx, -n:] = ( run_config["model"]["cond_semantic_n_vocab"] - 1 ) semantic_codes = torch.repeat_interleave(semantic_codes, repeats=2, dim=0) text_codes = torch.repeat_interleave(text_codes, repeats=2, dim=0) if noise_schedule == "default" or noise_schedule == "early": # Draw uniformly distributed continuous timesteps t = rng.draw(batch_size // 2)[:, 0].to(device).to(torch.bfloat16) elif noise_schedule == "logit_normal": # Draw from a logit-normal distribution t = torch.sigmoid(torch.randn(batch_size // 2, device=device).to(torch.bfloat16)) elif noise_schedule == "trunc_logit_normal": t = truncated_logistic_normal_rescaled(batch_size // 2).to(device).to(torch.bfloat16) # Flip the distribution t = 1 - t elif noise_schedule == "log_snr": t = sample_timesteps_logsnr(batch_size // 2).to(device).to(torch.bfloat16) else: raise ValueError(f"Invalid noise schedule: {noise_schedule}") # descretize to fixed steps if run_config["training"]["t_discretize"] > 0: # if run_config["training"]["t_discretize"] in t_steps_lookup: # t = torch.tensor( # random.choices(t_steps_lookup[run_config["training"]["t_discretize"]], k=batch_size // 2) # ) # else: # # double sample # random_steps = random.choice(list(t_steps_lookup.keys())) # t = torch.tensor(random.choices(t_steps_lookup[random_steps], k=batch_size // 2)) t = ( torch.round(t * run_config["training"]["t_discretize"]) / run_config["training"]["t_discretize"] ) t = torch.repeat_interleave(t, repeats=2, dim=0).to(device).to(torch.bfloat16) # print_with_time_master( # f"batch_size: {batch_size}, diff_input: {diffusion_input.shape}, semantic_codes: {semantic_codes.shape}" # ) # Replace 1% of t with ones to ensure training on terminal SNR # 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] # TODO: should these be different??? noise = torch.randn_like(diffusion_input) noise = noise.chunk(2)[0] noise = torch.repeat_interleave(noise, repeats=2, dim=0) assert torch.all(noise[::2] == noise[1::2]) noised_inputs = diffusion_input * alphas + noise * sigmas if ( run_config["training"]["ave_latent_frac"] > 0 and random.random() < run_config["training"]["ave_latent_frac"] ): diffusion_ave_input = diffusion_ave_input.to(device) diffusion_ave_input = diffusion_ave_input.to(t.dtype) noised_inputs = diffusion_ave_input * alphas + noise * sigmas targets = noise * alphas - diffusion_input * sigmas # assert torch.all(noised_inputs[::2] == noised_inputs[1::2]) assert torch.all(t[::2] == t[1::2]) assert torch.all(semantic_codes[::2] == semantic_codes[1::2]) assert torch.all(text_codes[::2] == text_codes[1::2]) # if ctx_vae is not None: # assert torch.all(ctx_vae[::2] == ctx_vae[1::2]) if ctx_mask is not None: assert torch.all(ctx_mask[::2] == ctx_mask[1::2]) # TODO: check this works with real infill ctx vae if run_config["model"]["infill_ctx_len"] > 0 or run_config["model"]["shared_ctx"]: infill_ctx_vae = info["infill_ctx_vae"].to(device) infill_ctx_mask = info["infill_ctx_mask"].to(device) infill_ctx_vae = torch.repeat_interleave(infill_ctx_vae, repeats=2, dim=0) infill_ctx_mask = torch.repeat_interleave(infill_ctx_mask, repeats=2, dim=0) assert torch.all(infill_ctx_vae[::2] == infill_ctx_vae[1::2]) assert torch.all(infill_ctx_mask[::2] == infill_ctx_mask[1::2]) else: infill_ctx_vae = None infill_ctx_mask = None 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}") 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, ) with torch.no_grad(): ref_v = ref_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, ) pi_rmse_loss = F.mse_loss(v, targets, reduction="none") # b, h, t # print_with_time_master(f"pi_rmse_loss: {pi_rmse_loss.shape}, padding_mask: {padding_mask.shape}") if not run_config["training"]["revert_label"]: pi_neg_rmse_loss = pi_rmse_loss[0::2, :, :].mean(dim=1).mean(dim=1) # neg even pi_pos_rmse_loss = pi_rmse_loss[1::2, :, :].mean(dim=1).mean(dim=1) # pos odd else: # labels are reverted pi_neg_rmse_loss = pi_rmse_loss[1::2, :, :].mean(dim=1).mean(dim=1) # neg odd pi_pos_rmse_loss = pi_rmse_loss[0::2, :, :].mean(dim=1).mean(dim=1) # pos even ref_rmse_loss = F.mse_loss(ref_v, targets, reduction="none") if not run_config["training"]["revert_label"]: ref_neg_rmse_loss = ref_rmse_loss[0::2, :, :].mean(dim=1).mean(dim=1) # neg even ref_pos_rmse_loss = ref_rmse_loss[1::2, :, :].mean(dim=1).mean(dim=1) # pos odd else: ref_neg_rmse_loss = ref_rmse_loss[1::2, :, :].mean(dim=1).mean(dim=1) # neg odd ref_pos_rmse_loss = ref_rmse_loss[0::2, :, :].mean(dim=1).mean(dim=1) # pos even # half of bt size is the shape # print_with_time_master( # f"pi_neg_rmse_loss: {pi_neg_rmse_loss.shape}, {pi_neg_rmse_loss}, pi_pos_rmse_loss: {pi_pos_rmse_loss.shape}, {pi_pos_rmse_loss}" # ) # print_with_time_master( # f"ref_neg_rmse_loss: {ref_neg_rmse_loss.shape}, {ref_neg_rmse_loss}, ref_pos_rmse_loss: {ref_pos_rmse_loss.shape}, {ref_pos_rmse_loss}" # ) loss, chosen_rewards, rejected_rewards = preference_loss( pi_pos_rmse_loss, pi_neg_rmse_loss, ref_pos_rmse_loss, ref_neg_rmse_loss, beta=run_config["training"]["ipo_beta"], ipo=True, ) # print(f"loss: {loss}, chosen_rewards: {chosen_rewards}, rejected_rewards: {rejected_rewards}") loss = loss.mean() reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() # print_with_time_master( # f"loss: {loss.shape}, reward acc: {reward_accuracies}, local rank: {dist.get_rank()}" # ) # loss = F.mse_loss(v, targets, reduction="none").mean(dim=1) # loss = loss[padding_mask].mean() dpo_losses = {} dpo_losses["reward_accuracies"] = reward_accuracies.item() dpo_losses["pos_loss"] = pi_pos_rmse_loss.mean().item() dpo_losses["neg_loss"] = pi_neg_rmse_loss.mean().item() dpo_losses["chosen_rewards"] = chosen_rewards.mean().item() dpo_losses["rejected_rewards"] = rejected_rewards.mean().item() return loss, diffusion_input.std(), dpo_losses 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, ref_model, run_config): total_val_loss = 0.0 num_batches = 0 dpo_ave_losses = defaultdict(float) for batch in val_dataloader: if num_batches >= 50: break loss, std_data, dpo_losses = train_step(model, batch, device, rng, ref_model, run_config) total_val_loss += loss.item() for dpo_loss_name, dpo_loss_value in dpo_losses.items(): dpo_ave_losses[dpo_loss_name] += dpo_loss_value 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 for dpo_loss_name, dpo_loss_value in dpo_ave_losses.items(): dpo_ave_losses[dpo_loss_name] = all_reduce_metric(dpo_loss_value, device) dpo_ave_losses[dpo_loss_name] /= num_batches return avg_val_loss, dpo_ave_losses @torch.no_grad() def compute_metrics( model, val_dataloader, vae_scale_factor=2.5, device="cuda", num_batches=1, global_step=0, ): """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 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) 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) wandb.log( { f"generated_audio/{i}": wandb.Audio( mp3_path, caption=f"Generated {global_step}" ) }, step=global_step, 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: # 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" 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", "train_semantic_memmap_filename", "train_metas_filename", "train_vae_memmap_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"], 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...") dataset_val = DPOGeneralMemmapMapDataset( 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=True if run_config["model"]["ctx_len"] > 0 else False, infill_vae_ctx=False, noise_ctx=run_config["data"]["noise_ctx"], shared_ctx=run_config["model"]["shared_ctx"], ) print_with_time_master(f"Loaded {len(dataset_val)} val samples on rank {os.environ['RANK']}.") # Use validation set as training set if in debug mode if args.debug: dataset_train = dataset_val else: dataset_train = DPOGeneralMemmapMapDataset( 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=False, 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"], prev_vae_ctx=True if run_config["model"]["ctx_len"] > 0 else False, infill_vae_ctx=False, shared_ctx=run_config["model"]["shared_ctx"], ) print_with_time_master(f"Loaded {len(dataset_train)} train samples 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, )