from contextlib import contextmanager, nullcontext import datetime import functools import json 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.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 data_utils_textualize import get_batch, _load_tokenizer from modules.base import ( apply_fsdp_checkpointing, configure_optimizers as base_configure_optimizers, RMSNorm, CausalSelfAttention, MLP, ) from modules.gpt_textualize import GPTConfig, GPT, 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}") # [MW] configuration data_dir = None out_dir = None train_filename = "data_tr.bin" train_metas_filename = "metas_tr.jsonl" train_info_filename = "info_tr.json" val_filename = "data_val.bin" val_metas_filename = "metas_val.jsonl" val_info_filename = "info_val.json" tokenizer_filename = "tokenizer_60k.json" use_private = False 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 add_fixed_loss = False suppress_text = False # vocab/time constants text_vocab_size = 60_032 # (multiple of 64) text_codebook_size = 60_001 text_pad_token = text_codebook_size text_infer_token = text_codebook_size + 1 semantic_vocab_size = 4032 # (multiple of 64) semantic_codebook_size = 4000 semantic_n_codebooks = 1 semantic_pad_token = semantic_codebook_size semantic_infer_token = semantic_codebook_size + 1 semantic_rate_hz = 25 semantic_shift_factor = 0 coarse_vocab_size = 2112 # (multiple of 64) coarse_codebook_size = 2048 coarse_n_codebooks = 12 coarse_pad_token = coarse_codebook_size coarse_infer_token = coarse_codebook_size + 1 coarse_rate_hz = 25 coarse_shift_factor = 0 t_text = 1152 t_audio = 3136 t_memmap = 3008 block_size = 4288 use_rotary_pos_emb = False last_codebook_weight = 0.5 # eval items custom_seed_offset = 0 eval_interval = 1 log_interval = 25 eval_iters = 250 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 = "suno-test" wandb_run_name = "test" # data gradient_accumulation_steps = 2 # used to simulate larger batch sizes batch_size = 8 # if gradient_accumulation_steps > 1, this is the micro-batch size # model n_layer = 24 n_head = 16 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 step_save_iters = 50_000 # at this checkpoint we save the model step_save_infer = False 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 = 5000 # how many steps to warm up for lr_decay_iters = None # should be ~= max_iters per Chinchilla min_lr = 0.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 # ----------------------------------------------------------------------------- assert t_text + t_audio == block_size assert dtype in ("bfloat16", "float32") if debug_val_only: train_filename = val_filename train_metas_filename = val_metas_filename train_info_filename = val_info_filename wandb_log = False 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, timeout=datetime.timedelta(seconds=2 * 60 * 60)) 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.concatenate( [ np.linspace(1, 1, 1), ], axis=0, ) loss_discount_facs = loss_discount_facs / loss_discount_facs.sum() if master_process: print("loss discounts for codebooks:", loss_discount_facs.round(3)) loss_discount_map = {} for n in range(1): loss_discount_map[f"text_{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 and not debug_val_only: os.makedirs(out_dir, exist_ok=True) print_with_time(f"logging checkpoint here: {out_dir}") # load data print_with_time("loading data...") if weights_multiplier is None: weights_multiplier_map = {} elif weights_multiplier == "base": weights_multiplier_map = { "youtube_music_lyrics": 2, "youtube_music_lyrics_foreign": 3, "genius_hq_lyrics": 2, "genius_hq_lyrics_foreign": 3, "deezer_lyrics": 2, "deezer_lyrics_foreign": 3, } elif weights_multiplier == "finetune": weights_multiplier_map = { "youtube_music_lyrics": 2, "youtube_music_lyrics_foreign": 3, "genius_hq_lyrics": 6, "genius_hq_lyrics_foreign": 8, "deezer_lyrics": 2, "deezer_lyrics_foreign": 3, } elif weights_multiplier == "pond5": weights_multiplier_map = { "pond5_music": 1, "genius_hq_lyrics_foreign": 0, "genius_hq_lyrics": 0, } else: weights_multiplier_map = { k.split(":")[0]: float(k.split(":")[1]) for k in weights_multiplier.split(";") if ":" in k } val_dataset_names = [] val_data_idx_lists = [] val_data_weights = [] # TODO: loading a memmap only once creates a memory 'leak', see here: # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122 val_data = np.memmap(os.path.join(data_dir, val_filename), dtype=np.uint16, mode="r") val_data = val_data.reshape(-1, t_memmap, semantic_n_codebooks + coarse_n_codebooks) assert val_data[:100, :, :semantic_n_codebooks].max() <= semantic_vocab_size assert val_data[:100, :, semantic_n_codebooks:].max() <= coarse_vocab_size with open(os.path.join(data_dir, val_info_filename)) as f: val_info = json.load(f) val_metas = [] with open(os.path.join(data_dir, val_metas_filename)) as f: for line in f: line = line.strip() if len(line) == 0: continue val_metas.append(json.loads(line)) assert len(val_data) == len(val_metas) idx_set = set() val_info = {"pond5_music": val_info["pond5_music"]} for dset_name, info in val_info.items(): val_dataset_names.append(dset_name) print(idx_set) idx_set |= set(info["idx_list"]) idx_list = [idx for idx in info["idx_list"]] random.shuffle(idx_list) val_data_idx_lists.append(idx_list) val_data_weights.append(len(info["idx_list"]) * weights_multiplier_map.get(dset_name, 1.0)) weights_norm = np.sum(val_data_weights) val_data_weights = [v / weights_norm for v in val_data_weights] # if not is_finetune: # assert len(idx_set) == len(val_data) for k in weights_multiplier_map.keys(): assert k in val_dataset_names del val_info, idx_set print_with_time(f"{len(val_data):,} lines of val loaded.") train_dataset_names = [] train_data_idx_lists = [] train_data_weights = [] train_data = np.memmap(os.path.join(data_dir, train_filename), dtype=np.uint16, mode="r") train_data = train_data.reshape(-1, t_memmap, semantic_n_codebooks + coarse_n_codebooks) assert train_data[:100, :, :semantic_n_codebooks].max() <= semantic_vocab_size assert train_data[:100, :, semantic_n_codebooks:].max() <= coarse_vocab_size with open(os.path.join(data_dir, train_info_filename)) as f: train_info = json.load(f) train_metas = [] with open(os.path.join(data_dir, train_metas_filename)) as f: for line in f: line = line.strip() if len(line) == 0: continue train_metas.append(json.loads(line)) assert len(train_data) == len(train_metas) idx_set = set() train_info = {"pond5_music": train_info["pond5_music"]} for dset_name, info in train_info.items(): train_dataset_names.append(dset_name) idx_set |= set(info["idx_list"]) idx_list = [idx for idx in info["idx_list"]] random.shuffle(idx_list) train_data_idx_lists.append(idx_list) train_data_weights.append(len(info["idx_list"]) * weights_multiplier_map.get(dset_name, 1.0)) weights_norm = np.sum(train_data_weights) train_data_weights = [v / weights_norm for v in train_data_weights] # if not is_finetune: # assert len(idx_set) == len(train_data) for k in weights_multiplier_map.keys(): assert k in train_dataset_names del train_info, idx_set print_with_time(f"{len(train_data):,} lines of train loaded.") # match val weights if match_val_weights: assert set(train_dataset_names) == set(val_dataset_names) weights_map = {k: v for k, v in zip(train_dataset_names, train_data_weights)} val_data_weights = [weights_map[k] for k in val_dataset_names] if master_process: weights_str = "train data weights:" for k, v in zip(train_dataset_names, train_data_weights): weights_str += f"\n {round(v*100, 1)}% {k}" print_with_time(weights_str) # make sure everything fits assert ( t_text + t_memmap + semantic_n_codebooks * semantic_shift_factor + (coarse_n_codebooks - 1) * coarse_shift_factor <= block_size ) # 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, block_size=block_size, bias=bias, dropout=dropout, text_vocab_size=text_vocab_size, text_codebook_size=text_codebook_size, text_pad_token=text_pad_token, text_infer_token=text_infer_token, semantic_vocab_size=semantic_vocab_size, semantic_codebook_size=semantic_codebook_size, semantic_n_codebooks=semantic_n_codebooks, semantic_pad_token=semantic_pad_token, semantic_infer_token=semantic_infer_token, semantic_rate_hz=semantic_rate_hz, semantic_shift_factor=semantic_shift_factor, 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_shift_factor=coarse_shift_factor, t_text=t_text, t_audio=t_audio, t_memmap=t_memmap, use_rotary_pos_emb=use_rotary_pos_emb, ) # init a new model from scratch print_with_time("Initializing a new model from scratch") gptconf = GPTConfig(**model_args) model = GPT(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 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) 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 min_lr + (learning_rate - min_lr) * 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) data_sampling_info = { "cfg": cfg, "batch_size": batch_size, "tokenizer_fp": os.path.join(data_dir, tokenizer_filename), "device": device, "device_type": device_type, "train": { "data": train_data, "metas": train_metas, "names": train_dataset_names, "weights": train_data_weights, "idx_lists": train_data_idx_lists, }, "val": { "data": val_data, "metas": val_metas, "names": val_dataset_names, "weights": val_data_weights, "idx_lists": val_data_idx_lists, }, } tokenizer = _load_tokenizer(data_sampling_info["tokenizer_fp"]) @torch.no_grad() def estimate_loss(random_idx=True): n_loss_modalities = len(train_dataset_names) + len(val_dataset_names) effective_eval_iters = int(round(eval_iters / n_loss_modalities)) 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() loss_prefix = "loss" if random_idx else "loss_fixed" n_loss_entries = n_loss_modalities * len(loss_discount_map) loss_tensor = torch.zeros(n_loss_entries, device=device) loss_tensor_keys = [] n_loss_entry = 0 for split in ["train", "val"]: n_datasets = len(train_dataset_names) if split == "train" else len(val_dataset_names) dataset_names = train_dataset_names if split == "train" else val_dataset_names for dataset_idx in range(n_datasets): losses = [] for k in range(effective_eval_iters): audio_offset, X, Y = get_batch( data_sampling_info, split, dataset_idx=dataset_idx, use_private=use_private, inference=True, dummy_data=dummy_data, suppress_text=suppress_text, n_offs=None if random_idx else k, ) print("[ground truth]") print(tokenizer.decode(Y[0, 0])) for iter_val in range(20): with ctx: loss_dict, logits = model(X, y=Y, audio_offset=audio_offset) X[:, 0, 3137 + iter_val] = logits.argmax(dim=-1)[:, iter_val] print("[prediction]") print(tokenizer.decode(X[0, 0, 3137:3157])) breakpoint() 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_prefix}_{dataset_names[dataset_idx]}_{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[f"train/{loss_prefix}"] = np.mean([v for k, v in tmp_out.items() if k.startswith("train/")]) out[f"val/{loss_prefix}"] = np.mean([v for k, v in tmp_out.items() if k.startswith("val/")]) out[f"train/{loss_prefix}_weighted"] = 0 out[f"val/{loss_prefix}_weighted"] = 0 weights_map = { "train": {k: v for k, v in zip(train_dataset_names, train_data_weights)}, "val": {k: v for k, v in zip(val_dataset_names, val_data_weights)}, } for split in ["train", "val"]: for dataset_idx in range(n_datasets): for n, loss_name in enumerate(loss_discount_map.keys()): dataset_name = dataset_names[dataset_idx] out[f"{split}/{loss_prefix}_weighted"] += ( weights_map[split][dataset_name] * tmp_out[f"{split}/{loss_prefix}_{dataset_name}_{loss_name}"] / len(loss_discount_map) / n_datasets ) model.train() for name, module in model.named_modules(): assert module.training # make sure we can undo everything return out ccnt = 0 # training loop print_with_time("training...") audio_offset, X, Y = get_batch( data_sampling_info, "train", use_private=use_private, dummy_data=dummy_data, suppress_text=suppress_text, ) # fetch the very first batch 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() if add_fixed_loss: fixed_losses = estimate_loss(random_idx=False) 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 if add_fixed_loss: for k, v in fixed_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 and not debug_val_only: # only write state dicts on rank 0 t_save = time.time() 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 % step_save_iters == 0: torch.save( checkpoint, os.path.join(out_dir, f"step_{iter_num/1000:.0f}k_ckpt.pt"), ) if step_save_infer: torch.save( {k: checkpoint[k] for k in ["model", "model_args", "best_val_loss"]}, os.path.join(out_dir, f"step_{iter_num/1000:.0f}k_infer.pt"), ) print_with_time(f"saving took {time.time() - t_save:.1f} seconds.") 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, logits = model(X, y=Y, audio_offset=audio_offset) # loss = sum( # v * loss_discount_map[k] # for k, v in loss_dict.items() # if k != "z_loss" # ) # if "z_loss" in loss_dict: # # use 1/10 of paper cause of multi-codebook # loss += 1e-5 * loss_dict["z_loss"] # if wandb_log and master_process: # d = { # "iter": iter_num, # "z_loss": loss_dict["z_loss"].item(), # } # wandb.log(d) # 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] # breakpoint() # # debug token prediction # # for i in range(len(Y)): # # print("[ground truth]") # # print(tokenizer.decode(Y[i, 0])) # # print("[prediction]") # # print(tokenizer.decode(logits.argmax(dim=-1)[i])) # # ccnt += 1 # # if ccnt > 50: # # breakpoint() # # immediately async prefetch next batch while model is doing the forward pass on the GPU # audio_offset, X, Y = get_batch( # data_sampling_info, # "train", # use_private=use_private, # dummy_data=dummy_data, # suppress_text=suppress_text, # ) # 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, # block_size, # batch_size * gradient_accumulation_steps, # dt, # ) # tokens_per_s = ( # world_size * batch_size * 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" # ) # if wandb_log and master_process: # wandb.log( # { # "iter": iter_num, # "avg_running_loss": avg_loss, # "curr_lr": lr, # } # ) # iter_num += 1 # local_iter_num += 1 # termination conditions if iter_num >= max_iters: print_with_time("done.") break if ddp: destroy_process_group()