from contextlib import contextmanager, nullcontext import datetime import functools import logging import math import os import random import time from pathlib import Path import numpy as np import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, ShardingStrategy, ) from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed import init_process_group, destroy_process_group from torch.utils.data import DataLoader # from data_utils_fine import get_batch from data_utils_fine_dataset import CodesDir, FineDataset from modules.base import ( apply_fsdp_checkpointing, configure_optimizers as base_configure_optimizers, estimate_mfu_no_model, RMSNorm, CausalSelfAttention, MLP, ) from modules.fine import FineConfig, Fine, Block from utils.fsdp_policies import bfSixteen from utils.model_io import ( get_model_state_dict_on_rank_0, get_optimizer_state_dict_on_rank_0, ) @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) def print_with_time(content): """Print the content with the current time.""" print(f"[{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}]: {content}") out_dir = None coarse_dir = "/app/suno/victor/music_sample/concat/dac_2c_25_12" fine_dir = "/app/suno/victor/music_sample/concat/dac_2c_100x16" debug_val_only = False dummy_data = False preload_checkpoint = None preload_optimizer = False preload_strict = True preload_remove_embeddings = False preload_remove_wpe = False suppress_compile_warnings = True grad_checkpointing = False match_val_weights = True weights_multiplier = None # eg "genius_lyrics:2;genius_hq:0.5" is_finetune = False # vocab/time constants coarse_vocab_size = 4160 # (multiple of 64) coarse_codebook_size = 4096 coarse_n_codebooks = 12 coarse_pad_token = coarse_codebook_size coarse_infer_token = coarse_codebook_size + 1 coarse_rate_hz = 25 coarse_samples = 50 coarse_mask_period: int = 1 coarse_masked_samples: int = 0 fine_clone_coarse = False fine_vocab_size: int = 1152 fine_codebook_size: int = 1024 fine_n_codebooks: int = 16 fine_pad_token: int = 1024 fine_infer_token: int = 1025 fine_rate_hz: int = 100 fine_shift_factor: int = 5 fine_samples: int = 200 t_memmap = 3008 # eval items custom_seed_offset = 0 eval_interval = 2000 log_interval = 25 eval_iters = 50 eval_only = False # if True, script exits right after the first eval model_as_bfloat16 = False # only really for eval debug_gradients = False init_from = "scratch" # "scratch" or "resume" or "gpt2*" # wandb logging wandb_log = False wandb_project = "fine" wandb_run_name = "test" # data gradient_accumulation_steps = 1 # used to simulate larger batch sizes batch_size = 8 # if gradient_accumulation_steps > 1, this is the micro-batch size remove_last_coarse = 0 # remove the last N coarse codebooks from the input noise_level = 0.0 # add noise to the input data # model n_layer = 12 n_head = 12 n_kv_head = None d_head = 64 dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+ bias = False # do we use bias inside LayerNorm and Linear layers? # adamw optimizer learning_rate = 8e-4 # max learning rate max_iters = 50_000 # 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 = 500 # how many steps to warm up for lr_decay_iters = None # should be ~= max_iters per Chinchilla min_lr = 0 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla # DDP pr FSDP 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" compile = False # use PyTorch 2.0 to compile the model to be faster fsdp = False # fully sharded data parallel sharding_strategy = "no_shard" # ----------------------------------------------------------------------------- config_keys = [ k for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str)) ] exec(open("configurator.py").read()) # overrides from command line or config file config = {k: globals()[k] for k in config_keys} # will be useful for logging # ----------------------------------------------------------------------------- if fine_clone_coarse: print_with_time("cloning coarse config for fine") fine_dir = coarse_dir fine_vocab_size = coarse_vocab_size fine_codebook_size = coarse_codebook_size fine_n_codebooks = coarse_n_codebooks fine_pad_token = coarse_pad_token fine_infer_token = coarse_infer_token fine_rate_hz = coarse_rate_hz fine_samples = coarse_samples assert dtype in ("bfloat16", "float32") 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 fsdp: assert ddp, "found fsdp = True but ddp is False" if ddp: init_process_group(backend=backend) ddp_rank = int(os.environ["RANK"]) ddp_local_rank = int(os.environ["LOCAL_RANK"]) world_size = torch.distributed.get_world_size() 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 + 1 # each process gets a different seed print_with_time(f"ddp init, rank {ddp_rank}, local_rank {ddp_local_rank}") else: ddp_rank = 0 world_size = 1 # if not ddp, we are running on a single gpu, and one process master_process = True seed_offset = 1 seed_offset *= custom_seed_offset + 1 # multiply to not just shift torch.manual_seed(6006 + seed_offset) random.seed(6006 + seed_offset) np.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 ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype] ctx = ( nullcontext() if device_type == "cpu" or fsdp else torch.amp.autocast(device_type=device_type, dtype=ptdtype) ) # fairly arbitrary loss averaging here, eg: # [0.40] + [0.10, 0.09, 0.09, 0.08, 0.07, 0.06, 0.06, 0.05] loss_discount_facs = np.linspace(1, 0.5, fine_n_codebooks) loss_discount_facs = loss_discount_facs / loss_discount_facs.sum() loss_discount_map = {} for n in range(fine_n_codebooks): loss_discount_map[f"fine_{n}"] = loss_discount_facs[n] # 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_with_time(f"logging checkpoint here: {out_dir}") # init these up here, can override if init_from="resume" (i.e. from a checkpoint) iter_num = 0 best_val_loss = 1e9 # model init model_args = dict( n_layer=n_layer, n_head=n_head, n_kv_head=n_kv_head, d_head=d_head, bias=bias, dropout=dropout, coarse_vocab_size=coarse_vocab_size, coarse_codebook_size=coarse_codebook_size, coarse_n_codebooks=coarse_n_codebooks, coarse_pad_token=coarse_pad_token, coarse_infer_token=coarse_infer_token, coarse_rate_hz=coarse_rate_hz, coarse_samples=coarse_samples, coarse_mask_period=coarse_mask_period, coarse_masked_samples=coarse_masked_samples, fine_vocab_size=fine_vocab_size, fine_codebook_size=fine_codebook_size, fine_n_codebooks=fine_n_codebooks, fine_pad_token=fine_pad_token, fine_infer_token=fine_infer_token, fine_rate_hz=fine_rate_hz, fine_shift_factor=fine_shift_factor, fine_samples=fine_samples, t_memmap=t_memmap, ) # init a new model from scratch print_with_time("Initializing a new model from scratch") gptconf = FineConfig(**model_args) model = Fine(gptconf) if model_as_bfloat16: model.to(torch.bfloat16) if not fsdp: model.to(device) cfg = model.config # this is needed to calculate MFU later # it will get messed up by FSDP, so calculate now raw_model_n_params = model.get_num_params() # load data print_with_time("loading data...") coarse_dir = CodesDir(Path(coarse_dir)) fine_dir = CodesDir(Path(fine_dir)) train_dataset = FineDataset( gptconf, coarse_dir, fine_dir, split="train", ) val_dataset = FineDataset( gptconf, coarse_dir, fine_dir, split="val", ) train_dl = DataLoader( train_dataset, batch_size=batch_size, num_workers=4, pin_memory=True, ) val_dl = DataLoader( val_dataset, batch_size=batch_size, num_workers=4, pin_memory=True, ) train_iter = iter(train_dl) val_iter = iter(val_dl) def get_batch(split="train"): if split == "train": it = train_iter elif split == "val": it = val_iter else: raise ValueError(f"split {split} not recognized") coarse_offset, X, Y = next(it) X = X.to(device) Y = Y.to(device) return coarse_offset[0].item(), X, Y print_with_time(f"{len(coarse_dir):,} lines of train loaded.") # load checkpoint # loading the checkpoint into the model needs to happen before wrapping # loading checkpoint into optimizer into FSDP needs to happen AFTER wrapping # to shard the optimizer to all the ranks model_state_dict, optimizer_state_dict = None, None if preload_checkpoint is not None: print_with_time("preloading checkpoint") cur_state_dict = model.state_dict() checkpoint = torch.load(preload_checkpoint, map_location="cpu") # fix checkpoint based on gqa change if "n_embd" in checkpoint.get("model_args", {}): n_emb = checkpoint["model_args"]["n_embd"] new_state_dict = {} for k, v in checkpoint["model"].items(): if "c_attn" in k: new_state_dict[k.replace("c_attn", "c_attn_q")] = v[: n_emb * 1] new_state_dict[k.replace("c_attn", "c_attn_k")] = v[n_emb * 1 : n_emb * 2] new_state_dict[k.replace("c_attn", "c_attn_v")] = v[n_emb * 2 :] else: new_state_dict[k] = v checkpoint["model_args"]["n_kv_head"] = None checkpoint["model_args"]["d_head"] = n_emb // checkpoint["model_args"]["n_head"] del checkpoint["model_args"]["n_embd"] checkpoint["model"] = new_state_dict model_state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint # fix the keys of the state dict # depending on how model was saved, it may have a prefix in the keys unwanted_prefix = "_orig_mod." loaded_has_prefix = any(k.startswith(unwanted_prefix) for k in model_state_dict) if loaded_has_prefix and not compile: for k, v in list(model_state_dict.items()): if k.startswith(unwanted_prefix): model_state_dict[k[len(unwanted_prefix) :]] = model_state_dict.pop(k) if preload_remove_embeddings: print_with_time("redoing preloaded embeddings") model_state_dict.pop("transformer.wte.weight") model_state_dict.pop("lm_head.weight") if preload_remove_wpe: model_state_dict.pop("transformer.wpe.weight") if preload_optimizer: optimizer_state_dict = checkpoint["optimizer"] iter_num = checkpoint["iter_num"] best_val_loss = checkpoint["best_val_loss"] del cur_state_dict, checkpoint # import torch._dynamo # torch._dynamo.config.cache_size_limit = 512 #64 # compile the model if compile: print_with_time("compiling the model... (takes a ~minute)") compile_ctx = suppress_logging if suppress_compile_warnings else nullcontext with compile_ctx(): # model = torch.compile(model, fullgraph=True, mode="max-autotune") model = torch.compile(model, fullgraph=True, mode="reduce-overhead") else: print_with_time("not compiling model.") # order matters: # FSDP: load model ckpt, wrap model, make optim (sharded), shard ckpt into optimizer # DDP: load model ckpt, make optimizer, load model and optimizer checkpoints if fsdp: if master_process: print_with_time("wrapping model in FSDP ....") # loading the state dict only happens on rank 0 # this is different from DDP if model_state_dict is not None: print_with_time("loading state dict into model ... on rank 0") model.load_state_dict(model_state_dict, strict=preload_strict) else: # using FSDP.shard_full_optim_state_dict instead of scatter... # needs state dict on all ranks. This is more CPU memory costs # and lower communication costs. This is also more robust to different # sharding strategies, so change with care. pass model_state_dict = None # free this memory auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls={Block}, ) model = FSDP( model, auto_wrap_policy=auto_wrap_policy, mixed_precision=bfSixteen, sharding_strategy=getattr(ShardingStrategy, sharding_strategy.upper()), device_id=torch.cuda.current_device(), sync_module_states=True, use_orig_params=True, # cpu_offload=torch.distributed.fsdp.CPUOffload(offload_params=True), ) if grad_checkpointing: apply_fsdp_checkpointing(model) optimizer = base_configure_optimizers( model, weight_decay, learning_rate, (beta1, beta2), device_type, use_fused=False, is_fsdp=True, ) if preload_optimizer: if master_process: assert optimizer_state_dict is not None, "no optimizer state dict found" print_with_time("sharding optimizer state dict") # needs to be called on all ranks sharded_osd = FSDP.shard_full_optim_state_dict(optimizer_state_dict, model, optim=optimizer) optimizer.load_state_dict(sharded_osd) else: # both DDP and single-worker if model_state_dict is not None: if master_process: print_with_time("loading model state dict") model.load_state_dict(model_state_dict, strict=preload_strict) # optimizer optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type) if preload_optimizer: assert optimizer_state_dict is not None, "no optimizer state dict found" if master_process: print_with_time("loading optimizer state dict") optimizer.load_state_dict(optimizer_state_dict) optimizer_state_dict = None if ddp: if master_process: print_with_time("wrapping model in DDP") model = DDP(model, device_ids=[ddp_local_rank]) del model_state_dict del optimizer_state_dict torch.cuda.empty_cache() # 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) @torch.no_grad() def estimate_loss(): if fsdp: modules_for_eval = ( torch.nn.Linear, torch.nn.Dropout, torch.nn.Embedding, torch.nn.SiLU, RMSNorm, MLP, CausalSelfAttention, ) for name, module in model.named_modules(): if isinstance(module, modules_for_eval): module.train(False) else: pass # might want a print_with_time here for debugging else: model.eval() n_loss_entries = len(loss_discount_map) * 2 loss_tensor = torch.zeros(n_loss_entries, device=device) loss_tensor_keys = [] n_loss_entry = 0 for split in ["train", "val"]: losses = [] for k in range(eval_iters): coarse_offset, X, Y = get_batch(split) with ctx: loss_dict = model(X, y=Y, coarse_offset=coarse_offset) losses.append([loss_dict[k].item() for k in loss_discount_map.keys()]) for n, loss_name in enumerate(loss_discount_map.keys()): loss_tensor[n_loss_entry] = np.mean([e[n] for e in losses]) loss_tensor_keys.append(f"{split}/loss_{loss_name}") n_loss_entry += 1 if ddp: torch.distributed.all_reduce(loss_tensor, op=torch.distributed.ReduceOp.AVG) tmp_out = {k: loss_tensor[n].item() for n, k in enumerate(loss_tensor_keys)} # add extra loss items out = {k: v for k, v in tmp_out.items()} out["train/loss"] = np.mean([v for k, v in tmp_out.items() if k.startswith("train/")]) out["val/loss"] = np.mean([v for k, v in tmp_out.items() if k.startswith("val/")]) model.train() for name, module in model.named_modules(): assert module.training # make sure we can undo everything return out # training loop print_with_time("training...") coarse_offset, X, Y = get_batch("train") t0 = time.time() t00 = time.time() t_start = time.time() # absolute time since starting to train local_iter_num = 0 # number of iterations in the lifetime of this process # raw_model = model.module if ddp else model # unwrap DDP container if needed mfu = 0 tokens_per_s = 0 running_loss = [] # number samples trained on since last time this number is synced # this will get reset gathered and then reset to 0 every log_interval samples_fetched_since_last_gather = torch.zeros(1, device=device) # total number of samples trained on total_samples_fetched = 0 # with torch.autograd.set_detect_anomaly(True): 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 == 0 or iter_num == max_iters - 1: 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_with_time( f"loss estimation took {estimation_time:.1f} seconds." f" ({eval_time_pct:.1f}% of loop)" ) print_with_time( 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, "mfu": mfu * 100, # convert to percentage "tok/s": tokens_per_s, "samples_trained": total_samples_fetched, } for k, v in losses.items(): log_dict[k] = v wandb.log(log_dict) if iter_num > 0: # state dicts need to be collected on all ranks model_state = get_model_state_dict_on_rank_0(model, ddp_rank) optim_state = get_optimizer_state_dict_on_rank_0(model, optimizer, ddp_rank) if master_process: # only write state dicts on rank 0 assert model_state is not None, "some sort of distributed bug" assert optim_state is not None, "some sort of distributed bug" checkpoint = { "model": model_state, "optimizer": optim_state, "model_args": model_args, "iter_num": iter_num, "best_val_loss": losses["val/loss"], "config": config, } print_with_time(f"saving checkpoint to {out_dir}") if losses["val/loss"] < best_val_loss: torch.save(checkpoint, os.path.join(out_dir, "best_ckpt.pt")) torch.save(checkpoint, os.path.join(out_dir, "last_ckpt.pt")) torch.save( {k: checkpoint[k] for k in ["model", "model_args", "best_val_loss"]}, os.path.join(out_dir, "last_ckpt_infer.pt"), ) if iter_num % 50_000 == 0: torch.save( checkpoint, os.path.join(out_dir, f"step_{iter_num/1000:.0f}k_ckpt.pt"), ) del model_state, optim_state torch.cuda.empty_cache() 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_with_time("eval test done.") break # forward backward update, with optional gradient accumulation to simulate larger batch size for micro_step in range(gradient_accumulation_steps): if ddp and micro_step < gradient_accumulation_steps - 1: grad_sync_context = model.no_sync else: grad_sync_context = nullcontext with grad_sync_context(): with ctx: loss_dict = model(X, y=Y, coarse_offset=coarse_offset) loss = sum(v * loss_discount_map[k] for k, v in loss_dict.items()) loss_val = loss.item() # loss as float. this is a CPU-GPU sync loss = loss / gradient_accumulation_steps samples_fetched_since_last_gather[0] += X.shape[0] # immediately async prefetch next batch while model is doing the forward pass on the GPU coarse_offset, X, Y = get_batch("train") if debug_gradients and wandb_log and master_process: d = { "iter": iter_num, "debug_loss": 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) # backward pass loss.backward() running_loss.append(loss_val) # clip the gradient if grad_clip != 0.0: if fsdp: grad_norm = model.clip_grad_norm_(grad_clip) if torch.isnan(grad_norm): raise RuntimeError("Found NaN infinite grad") else: grad_norm = torch.nn.utils.clip_grad_norm_( model.parameters(), grad_clip, error_if_nonfinite=True ) grad_norm = grad_norm.item() optimizer.step() # flush the gradients as soon as we can, no need for this memory anymore optimizer.zero_grad(set_to_none=True) if master_process and wandb_log and grad_norm is not None: wandb.log( { "iter": iter_num, "misc/grad_norm": grad_norm, } ) # timing and logging t1 = time.time() dt = t1 - t0 t0 = t1 if iter_num % log_interval == 0: if ddp: torch.distributed.all_reduce( samples_fetched_since_last_gather, op=torch.distributed.ReduceOp.SUM ) total_samples_fetched += samples_fetched_since_last_gather.item() samples_fetched_since_last_gather[0] = 0 if master_process and (iter_num % log_interval == 0 or iter_num == max_iters - 1): if local_iter_num >= 5: # let the training loop settle a bit mfu = estimate_mfu_no_model( raw_model_n_params, n_layer, n_head, n_head * d_head, cfg.block_size, batch_size * gradient_accumulation_steps, dt, ) tokens_per_s = world_size * batch_size * cfg.block_size * gradient_accumulation_steps / dt avg_loss = np.mean(running_loss) running_loss = [] print_with_time( f"iter {iter_num}:" f" avg_loss {avg_loss:.3f}," f" step_time {dt*1000:.1f}ms," f" mfu {mfu*100:.1f}%," f" throughput {tokens_per_s/1e3:,.0f}k tok/s," f" total time {t1 - t_start:.0f}s" ) iter_num += 1 local_iter_num += 1 # termination conditions if iter_num >= max_iters: print_with_time("done.") break if ddp: destroy_process_group()