from prefigure.prefigure import get_all_args, push_wandb_config import json import os import time import torch import pytorch_lightning as pl import random import time import shutil import numpy as np import funcy import gc import datetime from tqdm import tqdm from suno_utils.utils.text import read_jsonl, write_jsonl from typing import List from pytorch_lightning.callbacks import TQDMProgressBar from stable_audio_tools.data.dataset import create_dataloader_from_config from stable_audio_tools.models import create_model_from_config from stable_audio_tools.models.utils import ( load_ckpt_state_dict, remove_weight_norm_from_model, ) from stable_audio_tools.models.transformer import TransformerBlock from stable_audio_tools.training.diffusion import DiffusionCondTrainingWrapper from stable_audio_tools.training import ( create_training_wrapper_from_config, create_demo_callback_from_config, ) from stable_audio_tools.training.utils import copy_state_dict class ExceptionCallback(pl.Callback): def on_exception(self, trainer, module, err): print(f"{type(err).__name__}: {err}") class ModelConfigEmbedderCallback(pl.Callback): def __init__(self, model_config): self.model_config = model_config def on_save_checkpoint(self, trainer, pl_module, checkpoint): checkpoint["model_config"] = self.model_config class PerformanceMetricsCallback(pl.Callback): def __init__(self, log_interval=50): self.start_time = None self.iter_count = 0 self.token_count = 0 self.log_interval = log_interval self.world_size = self.get_world_size() def get_world_size(self): # Get the world size from Slurm environment variables if "SLURM_NTASKS" in os.environ: return int(os.environ["SLURM_NTASKS"]) elif "WORLD_SIZE" in os.environ: return int(os.environ["WORLD_SIZE"]) return 1 # Default to 1 if not in a distributed environment def on_train_start(self, trainer, pl_module): self.start_time = time.time() def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): self.iter_count += 1 # Calculate the number of tokens in this batch # Adjust this based on your actual batch structure batch_size = batch[0].size(0) # Assuming the first element is the input tensor seq_length = batch[0].size( -1 ) # Assuming the last dimension is the sequence length # Account for effective batch size across all nodes effective_batch_size = batch_size * self.world_size self.token_count += effective_batch_size * seq_length if self.iter_count % self.log_interval == 0: current_time = time.time() elapsed_time = current_time - self.start_time ips = self.iter_count / elapsed_time tps = self.token_count / elapsed_time # Log metrics to wandb trainer.logger.log_metrics( { "train/iterations_per_second": ips, "train/tokens_per_second": tps, "train/training": pl_module.training, }, step=trainer.global_step, ) # Reset counters and start time for the next interval self.start_time = current_time self.iter_count = 0 self.token_count = 0 def on_train_epoch_start(self, trainer, pl_module): # Refresh world_size at the start of each epoch in case it changes self.world_size = self.get_world_size() def shard_data( input_dir: str, output_dir: str, memmap_filenames: List[str], vae_n_tokens_memmap: int = 3000, vae_dim: int = 128, vae_use_float16: bool = False, semantic_n_tokens_memmap: int = 750, semantic_n_codebooks: int = 1, codec_n_tokens_memmap: int = 750, codec_n_codebooks: int = 12, allow_shard_reuse: bool = False, ): """ Args: input_dir (str): directory where the memmap files are stored output_dir (str): directory where the sharded memmap files will be stored (on local disk) memmap_filenames (List[str]): list of filenames to shard vae_n_tokens_memmap (int): number of tokens in the VAE memmap vae_dim (int): dimension of the VAE memmap vae_use_float16 (bool): whether to use float16 for VAE memmap semantic_n_tokens_memmap (int): number of tokens in the semantic memmap semantic_n_codebooks (int): number of codebooks in the semantic memmap codec_n_tokens_memmap (int): number of tokens in the codec memmap codec_n_codebooks (int): number of codebooks in the codec memmap allow_shard_reuse (bool): whether to allow reusing shards if they already exist on local disk """ ddp_rank = int(os.environ["SLURM_PROCID"]) ddp_local_rank = int(os.environ["SLURM_LOCALID"]) world_size = int(os.environ["SLURM_NTASKS"]) n_gpus_per_node = torch.cuda.device_count() master_process = ddp_rank == 0 done_fp = os.path.join(output_dir, "done.txt") print(f"start sharding ...") while True: try: if os.path.exists(done_fp): # might fail cause nodes might do this concurrently os.remove(done_fp) else: break except: time.sleep(10) if ddp_local_rank == 0: print(f"sharding data on {ddp_rank}...") # check for existing shards if reusing if allow_shard_reuse and os.path.exists(done_fp): # check each file exists for fn in memmap_filenames: if not os.path.exists(os.path.join(output_dir, fn)): break # if all files exist, we're done return # remove existing shards shutil.rmtree(output_dir, ignore_errors=True) os.makedirs(output_dir) # load from data_dir and shard based on fraction that node should receive from_frac = ddp_rank / world_size to_frac = (ddp_rank + n_gpus_per_node) / world_size assert 0 <= from_frac <= 1 assert 0 <= to_frac <= 1 for fn in memmap_filenames: if "val" in fn: shutil.copyfile( os.path.join(input_dir, fn), os.path.join(output_dir, fn), ) continue if "vae" in fn: if vae_use_float16: data = np.memmap( os.path.join(input_dir, fn), dtype=np.float16, mode="r" ) out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.float16, mode="w+", shape=(1, vae_n_tokens_memmap, vae_dim), ) else: data = np.memmap( os.path.join(input_dir, fn), dtype=np.float32, mode="r" ) out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.float32, mode="w+", shape=(1, vae_n_tokens_memmap, vae_dim), ) data = data.reshape(-1, vae_n_tokens_memmap, vae_dim) print(f"sharding {fn} with shape {data.shape}") elif "semantic" in fn: data = np.memmap(os.path.join(input_dir, fn), dtype=np.uint16, mode="r") data = data.reshape(-1, semantic_n_tokens_memmap) out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="w+", shape=(1, semantic_n_tokens_memmap), ) print(f"sharding {fn} with shape {data.shape}") elif "codec" in fn: data = np.memmap(os.path.join(input_dir, fn), dtype=np.uint16, mode="r") data = data.reshape(-1, codec_n_tokens_memmap, codec_n_codebooks) out_data = np.memmap( os.path.join(output_dir, codec_n_tokens_memmap), dtype=np.uint16, mode="w+", shape=(1, codec_n_tokens_memmap, codec_n_codebooks), ) print(f"sharding {fn} with shape {data.shape}") elif "metas" in fn: # open jsonl file data = read_jsonl(os.path.join(input_dir, fn)) print(f"sharding {fn} with {len(data)} samples") else: raise ValueError(f"{fn} unknown format") from_idx = int(round(from_frac * len(data))) to_idx = int(round(to_frac * len(data))) idx_chunks = list(funcy.chunks(100_000, list(range(from_idx, to_idx)))) n_offs = 0 # idx_chunk is a list of indices for n_chunk, idx_chunk in enumerate(tqdm(idx_chunks)): if master_process: print(f"writing shard {n_chunk+1}/{len(idx_chunks)} for {fn}") if "vae" in fn: if vae_use_float16: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.float16, mode="r+", shape=( n_offs + len(idx_chunk), vae_n_tokens_memmap, vae_dim, ), ) else: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.float32, mode="r+", shape=( n_offs + len(idx_chunk), vae_n_tokens_memmap, vae_dim, ), ) elif "semantic" in fn: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="r+", shape=(n_offs + len(idx_chunk), semantic_n_tokens_memmap), ) elif "codec" in fn: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="r+", shape=( n_offs + len(idx_chunk), codec_n_tokens_memmap, codec_n_codebooks, ), ) elif "metas" in fn: shard_metas = [] else: raise ValueError(f"{fn} unknown format") # iterate over indices and store in out_data for n, idx in enumerate(idx_chunk): if "metas" in fn: shard_metas.append(data[idx]) else: out_data[n_offs + n] = data[idx] # write metas to disk if "metas" in fn: write_jsonl( shard_metas, os.path.join(output_dir, fn), do_append=True ) del shard_metas else: out_data.flush() n_offs += len(idx_chunk) del out_data del data gc.collect() with open(done_fp, "w") as f: f.write("") # loop until we're done while not os.path.exists(done_fp): time.sleep(10) print(f"done sharding on rank {ddp_rank}") def main(): args = get_all_args() seed = args.seed # Set a different seed for each process if using SLURM if os.environ.get("SLURM_PROCID") is not None: seed += int(os.environ.get("SLURM_PROCID")) random.seed(seed) torch.manual_seed(seed) torch.set_float32_matmul_precision("high") # Get JSON config from args.model_config with open(args.model_config) as f: model_config = json.load(f) with open(args.dataset_config) as f: dataset_config = json.load(f) if args.local_data_shard_dir is not None and len(args.local_data_shard_dir) > 0: shard = True # 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 = dataset_config.get(s, None) if fn is not None: memmap_fnames.append(fn) shard_data( dataset_config["dataset_dir"], args.local_data_shard_dir, memmap_fnames, vae_n_tokens_memmap=dataset_config["vae_n_tokens_memmap"], vae_dim=dataset_config["vae_dim"], vae_use_float16=dataset_config["vae_use_float16"], semantic_n_tokens_memmap=dataset_config["semantic_n_tokens_memmap"], semantic_n_codebooks=dataset_config["semantic_n_codebooks"], # codec_n_tokens_memmap=dataset_config["codec_n_tokens_memmap"], # codec_n_codebooks=dataset_config["codec_n_codebooks"], allow_shard_reuse=args.allow_shard_reuse, ) dataset_config["dataset_dir"] = args.local_data_shard_dir else: shard = False train_dl, val_dl = create_dataloader_from_config( dataset_config, batch_size=args.batch_size, num_workers=args.num_workers, sample_rate=model_config["sample_rate"], sample_size=model_config["sample_size"], audio_channels=model_config.get("audio_channels", 2), ) model = create_model_from_config(model_config) num_params = sum(p.numel() for p in model.parameters()) print(f"Number of model parameters: {num_params/1e9:.1f} B") if args.pretrained_ckpt_path: print(f"Loading pretrained model from {args.pretrained_ckpt_path}") copy_state_dict(model, load_ckpt_state_dict(args.pretrained_ckpt_path)) if args.remove_pretransform_weight_norm == "pre_load": remove_weight_norm_from_model(model.pretransform) if args.pretransform_ckpt_path: model.pretransform.load_state_dict( load_ckpt_state_dict(args.pretransform_ckpt_path) ) # Remove weight_norm from the pretransform if specified if args.remove_pretransform_weight_norm == "post_load": remove_weight_norm_from_model(model.pretransform) # model = torch.compile(model) training_wrapper = create_training_wrapper_from_config(model_config, model) wandb_logger = pl.loggers.WandbLogger(project=args.name, name=args.run_name) wandb_logger.watch(training_wrapper) exc_callback = ExceptionCallback() if args.save_dir and isinstance(wandb_logger.experiment.id, str): checkpoint_dir = os.path.join( args.save_dir, wandb_logger.experiment.project, wandb_logger.experiment.id, "checkpoints", ) else: checkpoint_dir = None ckpt_callback = pl.callbacks.ModelCheckpoint( every_n_train_steps=args.checkpoint_every, dirpath=checkpoint_dir, save_top_k=1, save_last=True, monitor="train/mse_loss", ) save_model_config_callback = ModelConfigEmbedderCallback(model_config) # demo_callback = create_demo_callback_from_config(model_config, demo_dl=train_dl) # Combine args and config dicts args_dict = vars(args) args_dict.update({"model_config": model_config}) args_dict.update({"dataset_config": dataset_config}) push_wandb_config(wandb_logger, args_dict) # Set multi-GPU strategy if specified if args.strategy: if args.strategy == "deepspeed": from pytorch_lightning.strategies import DeepSpeedStrategy strategy = DeepSpeedStrategy( stage=2, contiguous_gradients=True, overlap_comm=True, reduce_scatter=True, reduce_bucket_size=5e8, allgather_bucket_size=5e8, load_full_weights=True, ) else: strategy = args.strategy else: strategy = "ddp_find_unused_parameters_true" if args.num_gpus > 1 else "auto" # overwrite for now cause we need timeout control if args.use_fsdp: policy = {TransformerBlock} strategy = pl.strategies.FSDPStrategy( timeout=datetime.timedelta(seconds=2 * 60 * 60), sharding_strategy="FULL_SHARD", # activation_checkpointing_policy=policy, mixed_precision=torch.distributed.fsdp.MixedPrecision( param_dtype=torch.float16, reduce_dtype=torch.float16, buffer_dtype=torch.float16, ), ) else: strategy = pl.strategies.DDPStrategy( timeout=datetime.timedelta(seconds=2 * 60 * 60), find_unused_parameters=True, ) trainer = pl.Trainer( devices=args.num_gpus, accelerator="gpu", num_nodes=args.num_nodes, strategy=strategy, precision=args.precision, accumulate_grad_batches=args.accum_batches, callbacks=[ ckpt_callback, exc_callback, save_model_config_callback, TQDMProgressBar(refresh_rate=args.refresh_rate or 1000), PerformanceMetricsCallback(log_interval=50), ], logger=wandb_logger, log_every_n_steps=100, max_epochs=-1, check_val_every_n_epoch=None, val_check_interval=None, # args.checkpoint_every, default_root_dir=args.save_dir, gradient_clip_val=args.gradient_clip_val, reload_dataloaders_every_n_epochs=0, use_distributed_sampler=False if shard else True, # don't split data if shard ) trainer.fit( training_wrapper, train_dl, # val_dl, ckpt_path=args.ckpt_path if args.ckpt_path else None, ) if __name__ == "__main__": main()