import shutil from contextlib import nullcontext import datetime import funcy import functools import math import gc import json import logging import math import os import random import time from collections import defaultdict 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 destroy_process_group, init_process_group from torch.nn import functional as F from data_utils import get_batch, read_jsonl, write_jsonl from modules.gpt import GPTConfig, GPTTrainConfig, GPT, Block from utils.fsdp_policies import bfSixteen from utils.helpers import ( dist_barrier, load_checkpoint, load_old_state_dict, print_with_time, print_with_time_master, suppress_logging, verify_preload_model_args, ) # turn down some annoying fsdp logging logging.getLogger("torch.distributed.fsdp._debug_utils").setLevel(logging.ERROR) logging.getLogger("torch.distributed.fsdp._optim_utils").setLevel(logging.ERROR) logging.getLogger("torch.distributed.checkpoint._dedup_tensors").setLevel(logging.ERROR) data_dir = None local_data_shard_dir = None allow_data_shard_reuse = False pre_processed_idx_path = "" # in case prev job crashed, filter indices pre_fixed_idx_path = "" # in case just run on a list of indices out_dir = None out_sub_dir = None 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 # use restricted data if available dummy_data = False # fake batches to test get_batch overhead preload_checkpoint = None local_cache_dir = None # checkpoint will get copied here, 1 per node to allow for faster loading preload_strict = True # enforce keys in dict on load suppress_compile_warnings = True # vocab/time constants text_vocab_size = 60_032 # (multiple of 64) text_codebook_size = 60_001 text_pad_token = text_codebook_size 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 = 50 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 = 5 t_text = 1152 t_audio = 3136 t_memmap = 3008 block_size = 4288 use_rotary_pos_emb = True rope_theta = 500_000 use_qk_norm = True activation_f = "silu" embed_scale_factor = 1.0 semantic_codebook_weight = 4.0 first_codebook_weight = 1.0 last_codebook_weight = 0.5 mask_padding = False layer_init = True # eval items log_interval = 25 write_interval = 100 model_as_bfloat16 = True # only really for eval # data batch_size = 1 # careful, higher creates issues if model was trained with pack # model n_layer = 24 n_head = 16 n_kv_head = 4 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? attention_type = "torch" attention_sliding_window_size = -1 global_every_n_layers = 1 # system device = "cuda" 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" # wandb logging wandb_log = True wandb_project = "eval-loss" wandb_run_name = "test" wandb_dir = None # ----------------------------------------------------------------------------- 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 # ----------------------------------------------------------------------------- batch_size_tokens = block_size * batch_size assert t_text + t_audio == block_size assert dtype in ("bfloat16", "float32") # 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="nccl", timeout=datetime.timedelta(seconds=24 * 60 * 60)) ddp_rank = int(os.environ["RANK"]) # global gpu rank ddp_local_rank = int(os.environ["LOCAL_RANK"]) # gpu rank within node world_size = torch.distributed.get_world_size() # total number of gpus device = f"cuda:{ddp_local_rank}" torch.cuda.set_device(device) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. print_with_time(f"ddp init, rank {ddp_rank}, local_rank {ddp_local_rank}, world_size {world_size}") else: ddp_rank = 0 ddp_local_rank = 0 world_size = 1 # if not ddp, we are running on a single gpu, and one process master_process = True n_gpus_per_node = torch.cuda.device_count() dist_barrier() # logging if wandb_log and master_process: import wandb wandb.init(project=wandb_project, name=wandb_run_name, config=config, dir=wandb_dir) wandb.run.log_code(".") torch.manual_seed(6006) random.seed(6006) np.random.seed(6006) 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] print_with_time_master(f"using {dtype}, {ptdtype}") 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, 0.5, semantic_n_codebooks) * semantic_codebook_weight, np.linspace(first_codebook_weight, last_codebook_weight, coarse_n_codebooks), ], axis=0, ) loss_discount_facs = loss_discount_facs / loss_discount_facs.sum() print_with_time_master(f"loss discounts for codebooks: {loss_discount_facs.round(3)}") loss_discount_map = {} for n in range(semantic_n_codebooks): loss_discount_map[f"semantic_{n}"] = loss_discount_facs[n] for n in range(coarse_n_codebooks): n2 = n + semantic_n_codebooks loss_discount_map[f"coarse_{n}"] = loss_discount_facs[n2] date_time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") if out_sub_dir is None: out_dir = os.path.join(out_dir, date_time_str) else: # will use the indicated sub dir name instead of date out_dir = os.path.join(out_dir, out_sub_dir) if master_process: os.makedirs(out_dir, exist_ok=True) print_with_time_master(f"logging output here: {out_dir}") # force shard data for simplicity # assert local_data_shard_dir is not None dist_barrier() if allow_data_shard_reuse: assert local_data_shard_dir is not None for fn in [ tokenizer_filename, val_filename, val_info_filename, val_metas_filename, ]: assert os.path.isfile(os.path.join(local_data_shard_dir, fn)), os.path.join( local_data_shard_dir, fn ) if local_data_shard_dir is not None and not allow_data_shard_reuse and ddp_local_rank == 0: print_with_time_master("sharding data...") shutil.rmtree(local_data_shard_dir, ignore_errors=True) os.makedirs(local_data_shard_dir) # copy over tokenizer for fn in [tokenizer_filename]: shutil.copyfile( os.path.join(data_dir, fn), os.path.join(local_data_shard_dir, fn), ) # 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 with open(os.path.join(data_dir, val_info_filename)) as f: val_info = json.load(f) # assemble shard info file new_val_info = {} new_idx_offset = 0 orig_idx_seq = [] for dset_name, info in val_info.items(): if "idx_list" not in info: print_with_time_master(f"{dset_name} doesn't have idx_list") idx_set = set() for idx_k, idx_v in info["idx_map"].items(): idx_set.add(int(idx_k)) for idx_v_v in idx_v: idx_set.add(int(idx_v_v)) idx_list = sorted(list(idx_set)) else: idx_list = [idx for idx in info["idx_list"]] from_n_sample = int(round(from_frac * len(idx_list))) to_n_sample = int(round(to_frac * len(idx_list))) keep_idx_list = idx_list[from_n_sample:to_n_sample] new_val_info[dset_name] = { "idx_list": list(range(new_idx_offset, new_idx_offset + len(keep_idx_list))) } orig_idx_seq.extend(keep_idx_list) new_idx_offset += len(keep_idx_list) print_with_time_master(f"shard size: {len(orig_idx_seq):,}") with open(os.path.join(local_data_shard_dir, val_info_filename), "w") as f: json.dump(new_val_info, f) print_with_time_master("done with info shard") # assemble shard metas file 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) val_metas = read_jsonl(os.path.join(data_dir, val_metas_filename), parse_idx_set=set(orig_idx_seq)) assert len(val_data) == len(val_metas) new_val_metas = [] for idx in orig_idx_seq: m = {k: v for k, v in val_metas[idx].items()} m["orig_idx"] = idx new_val_metas.append(m) assert not any(m is None for m in new_val_metas) write_jsonl(new_val_metas, os.path.join(local_data_shard_dir, val_metas_filename)) print_with_time_master("done with metas shard") # write shards to disk new_val_data = np.memmap( os.path.join(local_data_shard_dir, val_filename), dtype=np.uint16, mode="w+", shape=(1, t_memmap, semantic_n_codebooks + coarse_n_codebooks), ) new_idx = 0 orig_idx_seq_chunks = list(funcy.chunks(100_000, orig_idx_seq)) for n_chunk, orig_idx_seq_chunk in enumerate(orig_idx_seq_chunks): new_val_data = np.memmap( os.path.join(local_data_shard_dir, val_filename), dtype=np.uint16, mode="r+", shape=( new_idx + len(orig_idx_seq_chunk), t_memmap, semantic_n_codebooks + coarse_n_codebooks, ), ) # sorting drastically increases read speed cause of memory pagination for n1, n2 in sorted( list( zip( range(new_idx, new_idx + len(orig_idx_seq_chunk)), orig_idx_seq_chunk, ) ), key=lambda x: x[-1], ): new_val_data[n1] = val_data[n2] new_idx += len(orig_idx_seq_chunk) new_val_data.flush() print_with_time_master(f"processed memmap chunk {n_chunk+1}/{len(orig_idx_seq_chunks)}") del new_val_data, new_val_metas, new_val_info del val_metas, val_data, val_info gc.collect() print_with_time_master("done sharding.") if local_data_shard_dir is not None: random.seed(6006) data_dir = local_data_shard_dir dist_barrier() # load data print_with_time_master("loading data...") val_dataset_names = [] val_data_idx_list_flat = [] 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 # also here: https://github.com/karpathy/nanoGPT/commit/f68ac2200df82d59a1b916fb8236b4b8bc68baa7 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 = read_jsonl(os.path.join(data_dir, val_metas_filename)) assert len(val_data) == len(val_metas), (len(val_data), len(val_metas)) for dset_name, info in val_info.items(): val_dataset_names.append(dset_name) if "idx_list" not in info: print_with_time_master(f"{dset_name} doesn't have idx_list") idx_set = set() for idx_k, idx_v in info["idx_map"].items(): idx_set.add(int(idx_k)) for idx_v_v in idx_v: idx_set.add(int(idx_v_v)) idx_list = sorted(list(idx_set)) else: idx_list = [idx for idx in info["idx_list"]] val_data_idx_list_flat.extend(idx_list) val_data_idx_lists.append(idx_list) val_data_weights.append(len(val_data_idx_lists[-1])) # # If the job crashed and need to reload a list of processed indices if pre_processed_idx_path: val_data_idx_list_flat = [] with open(pre_processed_idx_path, "r") as fp: processed_idx = json.load(fp) processed_idx = set(processed_idx) for idx, val_meta in enumerate(val_metas): if val_meta["orig_idx"] not in processed_idx: val_data_idx_list_flat.append(idx) weights_norm = np.sum(val_data_weights) val_data_weights = [v / weights_norm for v in val_data_weights] val_artist_to_songs = defaultdict(list) print_with_time_master( f"{len(val_data):,} lines of val loaded, need to process {len(val_data_idx_list_flat)}" ) print_with_time_master("done loading data") dist_barrier() # make sure everything fits assert ( t_text + t_memmap + semantic_n_codebooks * semantic_shift_factor + (coarse_n_codebooks - 1) * coarse_shift_factor <= block_size ) # TODO: if no data sharding we need to do something like the below # # make work item list based on GPU (duplicate last sample to fill world) # n_samples_per_worker = int( # math.ceil(math.ceil(len(val_data_idx_list) / world_size) / batch_size) * batch_size # ) # val_data_idx_list += val_data_idx_list[-1:] * ( # n_samples_per_worker * world_size - len(val_data_idx_list) # ) # worker_row_idx_list = val_data_idx_list[ # ddp_rank * n_samples_per_worker : (ddp_rank + 1) * n_samples_per_worker # ] # tot_iter_num = int(n_samples_per_worker / batch_size) # if we have a fixed target list, just load and overwrite if pre_fixed_idx_path: val_data_idx_list_flat = [] with open(pre_fixed_idx_path, "r") as fp: val_data_idx_list_flat = json.load(fp) # 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, text_vocab_size=text_vocab_size, text_codebook_size=text_codebook_size, text_pad_token=text_pad_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, rope_theta=rope_theta, use_qk_norm=use_qk_norm, activation_f=activation_f, embed_scale_factor=embed_scale_factor, attention_sliding_window_size=attention_sliding_window_size, global_every_n_layers=global_every_n_layers, ) train_model_args = dict(dropout=dropout, attention_type=attention_type, layer_init=layer_init) if preload_checkpoint is not None and not preload_checkpoint.endswith(".pt"): # verification for old style we do later on checkpoint load verify_preload_model_args(model_args, preload_checkpoint) # init a new model from scratch print_with_time_master("Initializing a new model from scratch") gptconf = GPTConfig(**model_args) gpttrainconf = GPTTrainConfig(**train_model_args) model = GPT(gptconf, gpttrainconf) if model_as_bfloat16: model.to(torch.bfloat16) if not fsdp: model.to(device) cfg = model.config train_cfg = model.train_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() # import torch._dynamo # torch._dynamo.config.cache_size_limit = 512 #64 # compile the model if compile: print_with_time_master("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, mode="default") dist_barrier() else: print_with_time_master("not compiling model.") # load old single-file checkpoint if preload_checkpoint is not None and preload_checkpoint.endswith(".pt"): load_old_state_dict( model_args, preload_checkpoint, model, local_cache_dir, preload_strict=preload_strict ) dist_barrier() # 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: print_with_time_master("wrapping model in FSDP ....") 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), ) else: # both DDP and single-worker # optimizer if ddp: print_with_time_master("wrapping model in DDP") model = DDP(model, device_ids=[ddp_local_rank]) torch.cuda.empty_cache() dist_barrier() # load new distributed checkpoint if preload_checkpoint is not None and not preload_checkpoint.endswith(".pt"): _ = load_checkpoint( preload_checkpoint, False, model, None, ) dist_barrier() print_with_time_master("model setup done") data_sampling_info = { "cfg": cfg, "train_cfg": train_cfg, "batch_size": batch_size, "batch_size_tokens": batch_size_tokens, "tokenizer_fp": os.path.join(data_dir, tokenizer_filename), "device": device, "device_type": device_type, "val": { "data": val_data, "metas": val_metas, "infos": val_info, "artist_to_songs": val_artist_to_songs, "names": val_dataset_names, "weights": val_data_weights, "idx_lists": val_data_idx_lists, }, } # eval loop print_with_time_master("evaluating...") out_fp = os.path.join(out_dir, "loss.jsonl") # init the loss jsonl if it doesn't exists if master_process and not os.path.exists(out_fp): os.makedirs(out_dir, exist_ok=True) print_with_time_master(f"logging results here: {out_dir}") with open(out_fp, "w") as f: f.write("") logged_infos = [ "idx", "loss_sem", "loss_sem_10s", "loss_sem_last_20s", "loss_sem_no_text", "loss_sem_no_text_10s", "loss_sem_no_text_last_20s", "loss_coarse", "loss_coarse_10s", "loss_coarse_last_20s", "loss_coarse_no_text", "loss_coarse_no_text_10s", "loss_coarse_no_text_last_20s", ] for nn in range(coarse_n_codebooks): logged_infos.append(f"loss_coarse_{nn}") f.write(json.dumps(logged_infos) + "\n") print_with_time_master(f"logging results here: {out_dir}") loss_data = [] t_start = time.time() # absolute time since starting to train t0 = time.time() mfu = 0 tokens_per_s = 0 iter_num = 0 est_tot_iter_num = int(math.ceil(len(val_data_idx_list_flat)) / (batch_size * n_gpus_per_node)) print_with_time_master( f"estimated toal number of iterations {est_tot_iter_num}, size of the validation data {len(val_data_idx_list_flat)}, size of batch {batch_size}, ddp_rank is {ddp_rank}" ) while True: # prep containers for loss and index for all workers idx_mat_list = [torch.zeros(batch_size, device=device, dtype=torch.int64) for _ in range(world_size)] loss_mat_list = [ torch.zeros((batch_size, 12 + coarse_n_codebooks), device=device, dtype=torch.float32) for _ in range(world_size) ] # get batch and loss global_batch_row_idx_list = val_data_idx_list_flat[ iter_num * batch_size * n_gpus_per_node : (iter_num + 1) * batch_size * n_gpus_per_node ] # subslice only the part needed eval loss batch_row_idx_list = global_batch_row_idx_list[ ddp_local_rank * batch_size : (ddp_local_rank + 1) * batch_size ] real_batch_size = len(batch_row_idx_list) # make local index matrix idx_mat = -torch.ones(batch_size, device=device, dtype=torch.int64) for n in range(real_batch_size): idx_mat[n] = val_metas[batch_row_idx_list[n]]["orig_idx"] # prep local loss matrix loss_mat = torch.zeros((batch_size, 12 + coarse_n_codebooks), device=device, dtype=torch.float32) # fill to full batch if needed batch_row_idx_list += [0] * (batch_size - real_batch_size) X, Y, seq_lens = get_batch( data_sampling_info, "val", row_idx=batch_row_idx_list, min_text_offs=0, suppress_text=False, use_private=use_private, dummy_data=dummy_data, mask_padding=True, inference=True, ) X = X.to(device) Y = Y.to(device) X_2, Y_2, seq_lens_2 = get_batch( data_sampling_info, "val", row_idx=batch_row_idx_list, min_text_offs=0, suppress_text=True, use_private=use_private, dummy_data=dummy_data, mask_padding=True, inference=True, ) X_2 = X_2.to(device) Y_2 = Y_2.to(device) with torch.no_grad(), ctx: logits_sem, logits_coarse = model( X, seq_lens=seq_lens, return_logits=True, last_only=False, ) logits_sem_2, logits_coarse_2 = model( X_2, seq_lens=seq_lens_2, return_logits=True, last_only=False, ) # fill local loss matrix for n in range(real_batch_size): meta_row = val_metas[batch_row_idx_list[n]] max_dur_audio_idx = max( 1, int(math.floor((meta_row["end_s"] - meta_row["start_s"]) * semantic_rate_hz)) ) max_10s_audio_idx = min(max_dur_audio_idx, int(round(10 * semantic_rate_hz))) last_20s_audio_idx = max(0, max_dur_audio_idx - int(round(20 * semantic_rate_hz))) for nn in range(semantic_n_codebooks): # semantic loss_mat[n, 0] += F.cross_entropy( logits_sem[n, nn, :max_dur_audio_idx, :].reshape(-1, logits_sem.size(-1)), Y[n, nn, :max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # semantic, first 10s loss_mat[n, 1] += F.cross_entropy( logits_sem[n, nn, :max_10s_audio_idx, :].reshape(-1, logits_sem.size(-1)), Y[n, nn, :max_10s_audio_idx].reshape(-1), ignore_index=-1, ) # semantic, last 20s loss_mat[n, 2] += F.cross_entropy( logits_sem[n, nn, last_20s_audio_idx:max_dur_audio_idx, :].reshape( -1, logits_sem.size(-1) ), Y[n, nn, last_20s_audio_idx:max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # semantic, no text loss_mat[n, 3] += F.cross_entropy( logits_sem_2[n, nn, :max_dur_audio_idx, :].reshape(-1, logits_sem_2.size(-1)), Y_2[n, nn, :max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # semantic, no text, first 10s loss_mat[n, 4] += F.cross_entropy( logits_sem_2[n, nn, :max_10s_audio_idx, :].reshape(-1, logits_sem_2.size(-1)), Y_2[n, nn, :max_10s_audio_idx].reshape(-1), ignore_index=-1, ) # semantic, no text, last 20s loss_mat[n, 5] += F.cross_entropy( logits_sem_2[n, nn, last_20s_audio_idx:max_dur_audio_idx, :].reshape( -1, logits_sem_2.size(-1) ), Y_2[n, nn, last_20s_audio_idx:max_dur_audio_idx].reshape(-1), ignore_index=-1, ) for nn in range(coarse_n_codebooks): nn2 = semantic_n_codebooks + nn # coarse loss_mat[n, 6] += F.cross_entropy( logits_coarse[n, nn, :max_dur_audio_idx, :].reshape(-1, logits_coarse.size(-1)), Y[n, nn2, :max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # coarse, first 10s loss_mat[n, 7] += F.cross_entropy( logits_coarse[n, nn, :max_10s_audio_idx, :].reshape(-1, logits_coarse.size(-1)), Y[n, nn2, :max_10s_audio_idx].reshape(-1), ignore_index=-1, ) # coarse, last 20s loss_mat[n, 8] += F.cross_entropy( logits_coarse[n, nn, last_20s_audio_idx:max_dur_audio_idx, :].reshape( -1, logits_coarse.size(-1) ), Y[n, nn2, last_20s_audio_idx:max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # coarse, no text loss_mat[n, 9] += F.cross_entropy( logits_coarse_2[n, nn, :max_dur_audio_idx, :].reshape(-1, logits_coarse_2.size(-1)), Y_2[n, nn2, :max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # coarse, no text, first 10s loss_mat[n, 10] += F.cross_entropy( logits_coarse_2[n, nn, :max_10s_audio_idx, :].reshape(-1, logits_coarse_2.size(-1)), Y_2[n, nn2, :max_10s_audio_idx].reshape(-1), ignore_index=-1, ) # coarse, no text, last 20s loss_mat[n, 11] += F.cross_entropy( logits_coarse_2[n, nn, last_20s_audio_idx:max_dur_audio_idx, :].reshape( -1, logits_coarse_2.size(-1) ), Y_2[n, nn2, last_20s_audio_idx:max_dur_audio_idx].reshape(-1), ignore_index=-1, ) # save each coarse codebook's loss loss_mat[n, 12 + nn] += F.cross_entropy( logits_coarse[n, nn, :max_dur_audio_idx, :].reshape(-1, logits_coarse.size(-1)), Y[n, nn2, :max_dur_audio_idx].reshape(-1), ignore_index=-1, ) for n in range(6): loss_mat[:, n] = loss_mat[:, n] / semantic_n_codebooks for n in range(6, 12): loss_mat[:, n] = loss_mat[:, n] / coarse_n_codebooks # collect and put together if ddp: torch.distributed.all_gather(idx_mat_list, idx_mat) torch.distributed.all_gather(loss_mat_list, loss_mat) else: idx_mat_list = [idx_mat] loss_mat_list = [loss_mat] n_empty_slots = 0 for idx_mat, loss_mat in zip(idx_mat_list, loss_mat_list): idx_mat = idx_mat.detach().cpu().numpy() loss_mat = loss_mat.detach().cpu().numpy() for n in range(batch_size): if idx_mat[n] < 0: n_empty_slots += 1 continue loss_info_list = [ int(idx_mat[n]), ] for curr_i in range(12 + coarse_n_codebooks): loss_info_list.append(round(float(loss_mat[n][curr_i]), 5)) if master_process: loss_data.append(loss_info_list) if n_empty_slots > 0: # print(f"debug: rank {ddp_rank}, n_empty: {n_empty_slots}") print_with_time_master(f"{n_empty_slots} empty slots") is_done = n_empty_slots == batch_size * world_size # timing and logging t1 = time.time() dt = t1 - t0 t0 = t1 if master_process and (iter_num % log_interval == 0 or is_done): tokens_per_s = batch_size * block_size / dt print_with_time_master( f"iter {iter_num}/{est_tot_iter_num}:" f" step_time {dt*1000:.1f}ms," f" throughput {tokens_per_s/1e3:,.0f}k tok/s/node," f" total time {t1 - t_start:.0f}s" ) if wandb_log: log_dict = { "iter": iter_num, "tokens_per_s": tokens_per_s, } wandb.log(log_dict) if master_process and (iter_num % write_interval == 0 or is_done): print_with_time_master(f"writing to {out_fp}, {len(loss_data)} rows") with open(os.path.join(out_dir, "loss.jsonl"), "a") as f: for row in loss_data: f.write(json.dumps(row) + "\n") # only clear loss when we are at write interval or done if iter_num % write_interval == 0 or is_done: # always clear loss_data loss_data = [] if is_done: # we are done break dist_barrier() iter_num += 1 dist_barrier() print_with_time_master("done.") if ddp: destroy_process_group()