from contextlib import contextmanager, nullcontext import datetime import logging import math import os import random import time import numpy as np import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group import torchaudio from moo.model import Moo, MooConfig, MooFeatureExtractor @contextmanager def suppress_logging(highest_level=logging.CRITICAL): previous_level = logging.root.manager.disable logging.disable(highest_level) try: yield finally: logging.disable(previous_level) CUR_PATH = os.path.dirname(os.path.abspath(__file__)) data_dir = None out_dir = None train_filename = "moo_tr.bin" val_filename = "moo_val.bin" min_duration_s = 3 max_duration_s = 5 debug_val_only = False preload_checkpoint = None preload_optimizer = False preload_strict = True suppress_compile_warnings = True # eval items custom_seed_offset = 0 eval_interval = 2000 log_interval = 25 eval_iters = 250 eval_only = False # if True, script exits right after the first eval debug_gradients = False always_save_checkpoint = True # if True, always save a checkpoint after each eval init_from = "scratch" # "scratch" or "resume" or "gpt2*" # wandb logging wandb_log = False wandb_project = "suno-test" wandb_run_name = "test" # data gradient_accumulation_steps = 1 # used to simulate larger batch sizes batch_size = 12 # if gradient_accumulation_steps > 1, this is the micro-batch size # model be_awesome = True # adamw optimizer learning_rate = 3e-4 # max learning rate max_iters = 100000 # total number of training iterations weight_decay = 1e-1 beta1 = 0.9 beta2 = 0.95 grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0 # learning rate decay settings decay_lr = True # whether to decay the learning rate warmup_iters = 1000 # how many steps to warm up for lr_decay_iters = None # should be ~= max_iters per Chinchilla min_lr = 1e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla # DDP settings backend = "nccl" # "nccl", "gloo", etc. # system device = "cuda" # examples: "cpu", "cuda", "cuda:0", "cuda:1" etc., or try "mps" on macbooks dtype = "bfloat16" # "float32", "bfloat16", or "float16" (implements a GradScaler) compile = False # use PyTorch 2.0 to compile the model to be faster # ----------------------------------------------------------------------------- config_keys = [ k for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str)) ] exec(open(os.path.join(CUR_PATH, "configurator.py")).read()) # overrides from command line or config file config = {k: globals()[k] for k in config_keys} # will be useful for logging # ----------------------------------------------------------------------------- assert dtype in ("bfloat16", "float32") if debug_val_only: train_filename = val_filename eval_iters = int(eval_iters * gradient_accumulation_steps) if lr_decay_iters is None: lr_decay_iters = max_iters # various inits, derived attributes, I/O setup ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run? if ddp: init_process_group(backend=backend) ddp_rank = int(os.environ["RANK"]) ddp_local_rank = int(os.environ["LOCAL_RANK"]) device = f"cuda:{ddp_local_rank}" torch.cuda.set_device(device) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. seed_offset = ddp_rank # each process gets a different seed else: # if not ddp, we are running on a single gpu, and one process master_process = True seed_offset = 0 seed_offset += custom_seed_offset torch.manual_seed(1337 + seed_offset) random.seed(6006 + seed_offset) torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn device_type = "cuda" if "cuda" in device else "cpu" # for later use in torch.autocast # note: float16 data type will automatically use a GradScaler ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype] ctx = ( nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type=device_type, dtype=ptdtype) ) # logging if wandb_log and master_process: import wandb wandb.init(project=wandb_project, name=wandb_run_name, config=config) date_time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") out_dir = os.path.join(out_dir, date_time_str) if master_process: os.makedirs(out_dir, exist_ok=True) print(f"logging checkpoint here: {out_dir}") # load data print(f"NOT loading data...") # val_data = np.memmap(os.path.join(data_dir, val_filename), dtype=np.int16, mode="r") # train_data = np.memmap(os.path.join(data_dir, train_filename), dtype=np.int16, mode="r") # model init print("Initializing a new model from scratch") model_args = dict( be_awesome=be_awesome, ) conf = MooConfig(**model_args) model = Moo(conf) model.to(device) feature_extractor = MooFeatureExtractor.load() # init these up here, can override if init_from="resume" (i.e. from a checkpoint) iter_num = 0 best_val_loss = 1e9 # initialize a GradScaler. If enabled=False scaler is a no-op scaler = torch.cuda.amp.GradScaler(enabled=(dtype == "float16")) # optimizer optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type) # load checkpoint if preload_checkpoint is not None: print("preloading checkpoint") cur_state_dict = model.state_dict() checkpoint = torch.load(preload_checkpoint, map_location=device) state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint # fix the keys of the state dictionary :( # honestly no idea how checkpoints sometimes get this prefix, have to debug more unwanted_prefix = "_orig_mod." for k, v in list(state_dict.items()): if k.startswith(unwanted_prefix): state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) model.load_state_dict(state_dict, strict=preload_strict) if preload_optimizer: print("preloading optimizer") optimizer.load_state_dict(checkpoint["optimizer"]) iter_num = checkpoint["iter_num"] best_val_loss = checkpoint["best_val_loss"] del cur_state_dict, state_dict, checkpoint # compile the model if compile: print("compiling the model... (takes a ~minute)") compile_ctx = suppress_logging if suppress_compile_warnings else nullcontext with compile_ctx(): model = torch.compile(model) # requires PyTorch 2.0 # wrap model into DDP container if ddp: model = DDP(model, device_ids=[ddp_local_rank]) raw_model = model.module if ddp else model # unwrap DDP container if needed def get_sample(split): # if split == "train": # idx = random.randint(0, len(train_data)-max_duration_s*16_000) # x = train_data[idx:idx+max_duration_s*16_000] # else: # idx = random.randint(0, len(val_data)-max_duration_s*16_000) # x = val_data[idx:idx+max_duration_s*16_000] x = np.zeros(max_duration_s*16_000, dtype=torch.float32) return x def get_batch(split): x_list = [] for _ in range(batch_size): x = get_sample(split) x_list.append(x) x = torch.stack(x_list) del x_list x = x.to(device) x_in = { 'input_values': x, 'attention_mask': None, 'mask_time_indices': None, 'sampled_negative_indices': None, } return x_in # helps estimate an arbitrarily accurate loss over either split using many batches @torch.no_grad() def estimate_loss(): n_loss_modalities = 2 effective_eval_iters = int(round(eval_iters / n_loss_modalities)) model.eval() loss_tensor = torch.zeros(n_loss_modalities, device=device) tensor_keys = [] n_loss_entry = 0 for split in ["train", "val"]: losses = [] for k in range(effective_eval_iters): X_text, X_semantic, Y = get_batch(split, inference=True) with ctx: loss = model(X_text, X_semantic, Y) losses.append(loss.item()) loss_tensor[n_loss_entry] = np.mean(losses) tensor_keys.append(f"{split}/loss") n_loss_entry += 1 if ddp: torch.distributed.all_reduce(loss_tensor, op=torch.distributed.ReduceOp.AVG) out = {k: loss_tensor[n].item() for n, k in enumerate(tensor_keys)} model.train() return out # learning rate decay scheduler (cosine with warmup) def get_lr(it): # 1) linear warmup for warmup_iters steps if it < warmup_iters: return learning_rate * it / warmup_iters # 2) if it > lr_decay_iters, return min learning rate if it > lr_decay_iters: return min_lr # 3) in between, use cosine decay down to min learning rate decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters) assert 0 <= decay_ratio <= 1 coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1 return min_lr + coeff * (learning_rate - min_lr) # training loop if master_process: print("Training...") t0 = time.time() t00 = time.time() local_iter_num = 0 # number of iterations in the lifetime of this process running_loss = [] running_contrastive_loss = [] running_diversity_loss = [] while True: # determine and set the learning rate for this iteration lr = get_lr(iter_num) if decay_lr else learning_rate for param_group in optimizer.param_groups: param_group["lr"] = lr # evaluate the loss on train/val sets and write checkpoints if iter_num % eval_interval: time_since_last_loss = time.time() - t00 t00 = time.time() losses = estimate_loss() estimation_time = time.time() - t00 eval_time_pct = np.clip(estimation_time / time_since_last_loss * 100, 0, 100) if master_process: print( f"loss estimation took {estimation_time:.1f} seconds." f" ({eval_time_pct:.1f}% of loop)") print( f"step {iter_num}: train loss {losses['train/loss']:.4f}," f" val loss {losses['val/loss']:.4f}" ) if wandb_log: log_dict = { "iter": iter_num, "lr": lr, } for k, v in losses.items(): log_dict[k] = v wandb.log(log_dict) if losses["val/loss"] < best_val_loss or always_save_checkpoint: if iter_num > 0: checkpoint = { "model": raw_model.state_dict(), "optimizer": optimizer.state_dict(), "model_args": model_args, "iter_num": iter_num, "best_val_loss": losses["val/loss"], "config": config, } print(f"saving checkpoint to {out_dir}") if losses["val/loss"] < best_val_loss: torch.save(checkpoint, os.path.join(out_dir, "best_ckpt.pt")) if always_save_checkpoint: torch.save(checkpoint, os.path.join(out_dir, "last_ckpt.pt")) if losses["val/loss"] < best_val_loss: best_val_loss = losses["val/loss"] # end if eval test only if iter_num == 0 and eval_only: print("eval test done.") break # forward backward update, with optional gradient accumulation to simulate larger batch size # and using the GradScaler if data type is float16 loss = None contrastive_loss = None diversity_loss = None for micro_step in range(gradient_accumulation_steps): if ddp: # in DDP training we only need to sync gradients at the last micro step. # the official way to do this is with model.no_sync() context manager, but # I really dislike that this bloats the code and forces us to repeat code # looking at the source of that context manager, it just toggles this variable model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1) X = get_batch("train") with ctx: outputs = model(X) loss = loss + outputs.loss if loss is not None else outputs.loss contrastive_loss = ( contrastive_loss + outputs.contrastive_loss if loss is not None else outputs.contrastive_loss ) diversity_loss = ( diversity_loss + outputs.diversity_loss if loss is not None else outputs.diversity_loss ) loss = loss / gradient_accumulation_steps contrastive_loss = contrastive_loss / gradient_accumulation_steps diversity_loss = diversity_loss / gradient_accumulation_steps loss_val = loss.item() # loss as float. note: this is a CPU-GPU sync point contrastive_loss_val = contrastive_loss.item() diversity_loss_val = diversity_loss.item() if debug_gradients and wandb_log and master_process: d = { "iter": iter_num, "debug_loss": loss_val, "debug_contrastive_loss": contrastive_loss_val, "debug_diversity_loss": diversity_loss_val, } grads = [] for name, param in model.named_parameters(): if param.grad is not None: grads.append(param.grad.norm().item()) if len(grads) > 0: d["debug_grads"] = np.mean(grads) wandb.log(d) # TODO: other losses get ignored?? scaler.scale(loss).backward() # scaler.scale(loss+contrastive_loss+diversity_loss).backward() running_loss.append(loss_val) running_contrastive_loss.append(contrastive_loss_val) running_diversity_loss.append(diversity_loss_val) # clip the gradient if grad_clip != 0.0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) # step the optimizer and scaler if training in fp16 scaler.step(optimizer) scaler.update() # flush the gradients as soon as we can, no need for this memory anymore optimizer.zero_grad(set_to_none=True) # timing and logging t1 = time.time() dt = t1 - t0 t0 = t1 if iter_num % log_interval == 0 and master_process: avg_loss = np.mean(running_loss) avg_contrastive_loss = np.mean(running_contrastive_loss) avg_diversity_loss = np.mean(running_diversity_loss) running_loss = [] running_contrastive_loss = [] running_diversity_loss = [] print( f"iter {iter_num}:" f" avg_loss {avg_loss:.3f}," f" avg_contrastive_loss {avg_contrastive_loss:.3f}," f" avg_diversity_loss {avg_diversity_loss:.3f}," f" step_time {dt*1000:.1f}ms" ) iter_num += 1 local_iter_num += 1 # termination conditions if iter_num > max_iters: break if ddp: destroy_process_group()