# DPO training # We will define positive labels to have odd indices in the dataset # And the negative labels have even indices in the dataset from collections import defaultdict import shutil from contextlib import nullcontext import datetime import funcy import functools import json import logging import math import os import random import time import gc 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.nn import functional as F from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed import destroy_process_group, init_process_group from torch.utils.data import DataLoader from data_utils_mmap import CustomAudioDataset, SamplingParams, read_jsonl, write_jsonl from utils.dpo_data_utils import get_batch, preference_loss from modules.base import ( apply_fsdp_checkpointing, configure_optimizers as base_configure_optimizers, estimate_mfu_no_model, LayerNorm, CausalSelfAttention, MLP, ) 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, load_old_optimizer_state_dict, print_with_time, print_with_time_master, save_checkpoint, save_old_checkpoint, suppress_logging, verify_preload_model_args, ) from utils.logging import build_gpu_memory_monitor, Color, NoColor color = Color if True else NoColor # 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) master_addr = "localhost" master_port = 12355 do_ipo = True dpo_beta = 5.0 # temperature parameter for DPO loss sft_loss_scale = 0.0 # probably want sth like 0.001 since loss is like ~3 loss_clip_threshold = 4.0 # skip sft loss for samples with pi_pos_loss data_dir = None local_data_shard_dir = None allow_data_shard_reuse = False 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" debug_val_only = False dummy_data = False # fake batches to test overhead preload_checkpoint = None model_cache_loss_name = "base_model" preload_optimizer = False 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 grad_checkpointing = False weights_multiplier = None # eg "genius_lyrics:2;genius_hq:0.5" is_finetune = False suppress_text = False checkpoint_save_old_format = True # vocab/time constants text_vocab_size = 60_032 # (multiple of 64) text_codebook_size = 60_001 text_pad_token = text_codebook_size semantic_n_codebooks = 1 semantic_vocab_size = 4032 # (multiple of 64) semantic_codebook_size = 4000 semantic_rate_hz = 25 semantic_shift_factor = 50 coarse_vocab_size = 2112 # (multiple of 64) coarse_codebook_size = 2048 coarse_n_codebooks = 12 data_coarse_n_codebooks = 12 coarse_rate_hz = 25 coarse_shift_factor = 5 t_text = 1152 t_audio = 3136 t_memmap = 6016 t_data_memmap = 6016 # for cover, artist, we expand the memmap file to a larger matrix to fit 4 mins block_size = 4288 use_rotary_pos_emb = True rope_theta = 500_000 use_qk_norm = True activation_f = "silu" embed_scale_factor = 1.0 # train params semantic_codebook_weight = 4.0 last_codebook_weight = 0.5 mask_padding = True pack = False layer_init = False infill_augment = False dropout_semantic = False allow_artist = False allow_cover = False use_text_loss = False use_mmbert = False use_vae_input = False output_paradigm = "gpt" # "gpt" or "diffusion" output_distribution = "semantic" # "semantic" or "vae" use_hoot = False use_ditto = False # 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 model_as_bfloat16 = False # only really for eval debug_gradients = False # wandb logging wandb_log = False wandb_project = "suno-test" wandb_run_name = "test" wandb_dir = None # 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 eval_loss_batch_size = 8 # this can be much larger than normal batch size # 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? # adamw optimizer learning_rate = 8e-4 # max learning rate min_lr = 8e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla max_iters = 100_000 # total number of training iterations warmup_iters = 10_000 # how many steps to warm up for lr_decay_iters = None # should be ~= max_iters per Chinchilla step_save_iters = 50_000 # at this checkpoint we save the model weight_decay = 1e-1 beta1 = 0.9 beta2 = 0.9 grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0 attention_type = "tao" attention_sliding_window_size = 1024 global_every_n_layers = 1 shuffle_data = False # load the data deterministically, or shuffle it local_shuffle_data = False # shuffle the data on each card # reward filtering filtered_indices_path = None # path to JSON with filtered train indices (filters train set only) # 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" # ----------------------------------------------------------------------------- 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 # ----------------------------------------------------------------------------- # auto set a few params batch_size_tokens = block_size * batch_size eval_loss_batch_size_tokens = block_size * eval_loss_batch_size if pack: batch_size = 1 assert t_text + t_audio <= block_size assert dtype in ("bfloat16", "float32") if debug_val_only or eval_only: train_filename = val_filename train_metas_filename = val_metas_filename train_info_filename = val_info_filename if not eval_only: wandb_log = False eval_iters = int(eval_iters * gradient_accumulation_steps) if lr_decay_iters is None: lr_decay_iters = max_iters # set up distributed variables os.environ["MASTER_ADDR"] = str(master_addr) os.environ["MASTER_PORT"] = str(master_port) if "SLURM_PROCID" in os.environ: # Running on SLURM if int(os.environ["SLURM_NTASKS_PER_NODE"]) != torch.cuda.device_count(): raise ValueError( f"SLURM_NTASKS_PER_NODE ({os.environ['SLURM_NTASKS_PER_NODE']}) does not match" f" the number of CUDA devices ({torch.cuda.device_count()}) on node {os.environ['HOSTNAME']}" ) ddp_rank = int(os.environ["SLURM_PROCID"]) ddp_local_rank = int(os.environ["SLURM_LOCALID"]) world_size = int(os.environ["SLURM_JOB_NUM_NODES"]) * int(os.environ["SLURM_NTASKS_PER_NODE"]) else: # Running locally ddp_rank = 0 ddp_local_rank = 0 world_size = 1 os.environ["RANK"] = str(ddp_rank) os.environ["LOCAL_RANK"] = str(ddp_local_rank) # 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: try: init_process_group( backend="nccl", timeout=datetime.timedelta(seconds=24 * 60 * 60), rank=ddp_rank, world_size=world_size, device_id=torch.device(f"cuda:{ddp_local_rank}"), ) except Exception as e: print(f"Distributed error on rank {ddp_rank} with host {os.environ['HOSTNAME']}") raise e 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: # if not ddp, we are running on a single gpu, and one process master_process = True seed_offset = 1 n_gpus_per_node = torch.cuda.device_count() dist_barrier() print_with_time_master(f"ddp init: world size {world_size} ddp_rank {ddp_rank}.") gc.disable() # manually gc to avoid slowdowns https://imbue.com/research/70b-infrastructure/ 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.0, last_codebook_weight, semantic_n_codebooks) * semantic_codebook_weight, np.linspace(1.0, 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] # 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(".") print_with_time_master(f"Total world size {world_size}") 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_master(f"logging checkpoint here: {out_dir}") # shard data if necessary 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, train_filename, train_info_filename, train_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 and val for fn in [tokenizer_filename, val_filename, val_metas_filename, val_info_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, train_info_filename)) as f: train_info = json.load(f) # make sure we turn into int since keys in json get auto turned into strings for dset_name in train_info.keys(): if "idx_map" in train_info[dset_name]: train_info[dset_name]["idx_map"] = { int(k): v for k, v in train_info[dset_name]["idx_map"].items() } # assemble shard info file new_train_info = {} new_idx_offset = 0 orig_idx_seq = [] for dset_name, info in train_info.items(): if info.get("task", "default") == "default": assert "idx_list" in info idx_list = 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_train_info[dset_name] = { "idx_list": list(range(new_idx_offset, new_idx_offset + len(keep_idx_list))), "task": "default", } orig_idx_seq.extend(keep_idx_list) new_idx_offset += len(keep_idx_list) elif info["task"] == "covers": idx_map_list = [(k, v) for k, v in info["idx_map"].items()] from_n_sample = int(round(from_frac * len(idx_map_list))) to_n_sample = int(round(to_frac * len(idx_map_list))) keep_idx_list = [] new_idx_map = defaultdict(list) for k, v in idx_map_list[from_n_sample:to_n_sample]: keep_idx_list.append(k) keep_idx_list.extend(v) new_idx_map[new_idx_offset] = list( range(new_idx_offset + 1, new_idx_offset + 1 + len(v)) ) new_idx_offset += 1 + len(v) new_train_info[dset_name] = {"idx_map": dict(new_idx_map), "task": "covers"} orig_idx_seq.extend(keep_idx_list) else: raise ValueError(f"unknown task for {dset_name} in info file") print_with_time_master(f"shard size: {len(orig_idx_seq):,}") with open(os.path.join(local_data_shard_dir, train_info_filename), "w") as f: json.dump(new_train_info, f) print_with_time_master("done with info shard") # assemble shard metas file train_data = np.memmap(os.path.join(data_dir, train_filename), dtype=np.uint16, mode="r") train_data = train_data.reshape(-1, t_data_memmap, semantic_n_codebooks + data_coarse_n_codebooks) train_metas = read_jsonl( os.path.join(data_dir, train_metas_filename), parse_idx_set=set(orig_idx_seq) ) assert len(train_data) == len(train_metas) new_train_metas = [train_metas[idx] for idx in orig_idx_seq] assert not any(m is None for m in new_train_metas) write_jsonl(new_train_metas, os.path.join(local_data_shard_dir, train_metas_filename)) print_with_time_master("done with metas shard") # write shards to disk new_train_data = np.memmap( os.path.join(local_data_shard_dir, train_filename), dtype=np.uint16, mode="w+", shape=(1, t_data_memmap, semantic_n_codebooks + data_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_train_data = np.memmap( os.path.join(local_data_shard_dir, train_filename), dtype=np.uint16, mode="r+", shape=( new_idx + len(orig_idx_seq_chunk), t_data_memmap, semantic_n_codebooks + data_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_train_data[n1] = train_data[n2] new_idx += len(orig_idx_seq_chunk) new_train_data.flush() print_with_time_master(f"processed memmap chunk {n_chunk + 1}/{len(orig_idx_seq_chunks)}") del new_train_data, new_train_metas, new_train_info del train_metas, train_data, train_info gc.collect() print_with_time_master("done sharding.") if local_data_shard_dir is not None: data_dir = local_data_shard_dir dist_barrier() # load data print_with_time_master("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, } else: weights_multiplier_map = { k.split(":")[0].strip(): float(k.split(":")[1].strip()) for k in weights_multiplier.strip(";").split(";") if ":" in k } def load_dataset( data_dir: str, filename: str, info_filename: str, metas_filename: str, weights_multiplier_map: dict, is_finetune: bool, ) -> list: dataset_names = [] data_idx_lists = [] # used to randomly sample from the dataset data_weights = [] data = np.memmap(os.path.join(data_dir, filename), dtype=np.uint16, mode="r") data = data.reshape(-1, t_data_memmap, semantic_n_codebooks + data_coarse_n_codebooks) assert data[:100, :, :semantic_n_codebooks].max() <= semantic_vocab_size if data_coarse_n_codebooks > 0: assert data[:100, :, semantic_n_codebooks:].max() <= coarse_vocab_size with open(os.path.join(data_dir, info_filename)) as f: infos = json.load(f) metas = read_jsonl(os.path.join(data_dir, metas_filename)) assert len(data) == len(metas), (len(data), len(metas)) artist_to_songs = defaultdict(list) for i, m in enumerate(metas): if "artist" in m: artist_to_songs[f"{m['dataset']}__{m['artist']}"].append(i) artist_to_songs = {k: v for k, v in artist_to_songs.items() if len(v) > 1} if allow_artist: assert len(artist_to_songs) > 0, "no artist data found" print_with_time_master(f"found {len(artist_to_songs):,} samples with artists on main process") idx_set = set() # for checking that we don't have any duplicates has_cover = False # make sure we turn into int since keys in json get auto turned into strings for dset_name in sorted(infos.keys()): if "idx_map" in infos[dset_name]: infos[dset_name]["idx_map"] = {int(k): v for k, v in infos[dset_name]["idx_map"].items()} for dset_name in sorted(infos.keys()): info = infos[dset_name] dataset_names.append(dset_name) if info.get("task", "default") == "default": assert "idx_list" in info idx_list = info["idx_list"][:] idx_set |= set(idx_list) idx_list = sorted([idx for idx in info["idx_list"]]) elif info["task"] == "covers": assert pack, "for now pack needs to be active to do covers" assert batch_size_tokens >= t_memmap * 2 + t_text, "for covers we need double the blocksize" has_cover = True # dict from original idx to list of covers idx # use original idxs as the idx_list idx_list = [] n_covers = 0 for idx, child_idx_l in info["idx_map"].items(): idx_list.append(int(idx)) idx_set.add(int(idx)) idx_set |= set(child_idx_l) n_covers += len(child_idx_l) print_with_time_master( f"found {len(idx_list):,} samples with {n_covers:,} total covers on main process" ) else: raise ValueError(f"unknown task for {dset_name} in info file") # DO NOT Shuffle DPO # random.shuffle(idx_list) data_idx_lists.append(idx_list) data_weights.append(len(idx_list) * weights_multiplier_map.get(dset_name, 1.0)) if allow_cover: assert has_cover, "no cover data found" weights_norm = np.sum(data_weights) data_weights = [v / weights_norm for v in data_weights] if not is_finetune: print_with_time_master(f"indexed {len(idx_set) / len(data) * 100:.1f}% of data") for k in weights_multiplier_map.keys(): assert k in dataset_names del idx_set shard_info = "" if local_data_shard_dir is None else " (sharded)" print_with_time_master(f"{len(data):,} lines of {filename} loaded.{shard_info}") assert len(data) == len(metas) assert len(infos) == len(dataset_names) == len(data_weights) == len(data_idx_lists) return ( dataset_names, data_idx_lists, data_weights, data, metas, infos, artist_to_songs, ) ( val_dataset_names, val_data_idx_lists, val_data_weights, val_data, val_metas, val_info, val_artist_to_songs, ) = load_dataset( data_dir, val_filename, val_info_filename, val_metas_filename, weights_multiplier_map, is_finetune, ) ( train_dataset_names, train_data_idx_lists, train_data_weights, train_data, train_metas, train_info, train_artist_to_songs, ) = load_dataset( data_dir, train_filename, train_info_filename, train_metas_filename, weights_multiplier_map, is_finetune, ) 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_master(weights_str) print_with_time_master("done loading data") dist_barrier() # 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_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_rate_hz=coarse_rate_hz, coarse_shift_factor=coarse_shift_factor, t_text=t_text, t_audio=t_audio, 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, use_text_loss=use_text_loss, use_mmbert=use_mmbert, use_vae_input=use_vae_input, output_paradigm=output_paradigm, output_distribution=output_distribution, use_hoot=use_hoot, use_ditto=use_ditto, ) # extra args that we don't need to save with the model 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, preload_strict=preload_strict) elif preload_checkpoint is not None: ckpt = torch.load(preload_checkpoint, mmap=True, weights_only=False) if "model_args" in ckpt: model_args = ckpt["model_args"] print_with_time_master(f"loaded model args from checkpoint: {model_args}") print_with_time_master( f"Original model args with block_size, t_text, t_audio: {model_args['block_size']}, {model_args['t_text']}, {model_args['t_audio']}" ) model_args["block_size"] = block_size model_args["t_text"] = t_text model_args["t_audio"] = t_audio print_with_time_master( f"Overriding model args with block_size, t_text, t_audio: {model_args['block_size']}, {model_args['t_text']}, {model_args['t_audio']}" ) gpu_memory_monitor = build_gpu_memory_monitor() if preload_checkpoint is None: raise ValueError("DOING DPO without pre-checkpoint?") # init a new model from scratch print_with_time_master("Initializing train 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 print_with_time_master("finish init train model") # 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() all_param = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print_with_time_master( f"trainable params: {trainable_params:,d} || " f"all params: {all_param:,d} || " f"trainable%: {100 * trainable_params / all_param:.4f}" ) # for name, module in model.named_children(): # if "lm_heads" in name: # continue # for param in module.parameters(): # param.requires_grad = False all_param = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print( f"trainable params: {trainable_params:,d} || " f"all params: {all_param:,d} || " f"trainable%: {100 * trainable_params / all_param:.4f}" ) # 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) dist_barrier() else: print_with_time_master("not compiling model.") iter_num = 0 total_tokens_processed = 0 rel_tokens_processed = 0 # necessary for preloaded checkpoints and throughput calculations best_val_loss = 1e9 # load old single-file checkpoint if preload_checkpoint is not None and preload_checkpoint.endswith(".pt"): print_with_time_master("start loading state dict") load_old_state_dict( model_args, preload_checkpoint, model, local_cache_dir, preload_strict=preload_strict, use_mmap=True, ) print_with_time_master("finish loading state dict") 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), ) gpu_mem_stats = gpu_memory_monitor.get_peak_stats() print_with_time_master( f"GPU memory usage for model: " f"{gpu_mem_stats.max_reserved_gib:.2f}GiB" f"({gpu_mem_stats.max_reserved_pct:.2f}%)" ) 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, ) else: # both DDP and single-worker # optimizer optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type) 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 old single-file checkpoint if preload_checkpoint is not None and preload_checkpoint.endswith(".pt") and preload_optimizer: iter_num, total_tokens_processed, best_val_loss = load_old_optimizer_state_dict( model, optimizer, local_cache_dir, ) # load new distributed checkpoint if preload_checkpoint is not None and not preload_checkpoint.endswith(".pt"): iter_num, total_tokens_processed, best_val_loss = load_checkpoint( preload_checkpoint, preload_optimizer, model, optimizer, ) dist_barrier() print_with_time_master("model setup done") # 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) # Apply reward filtering to train set only if filtered_indices_path is not None: print_with_time_master(f"Loading filtered indices from {filtered_indices_path}...") with open(filtered_indices_path, "r") as f: filtered_data = json.load(f) train_filtered_indices = filtered_data.get("train", None) if train_filtered_indices is not None: filtered_set = set(train_filtered_indices) print_with_time_master( f"Loaded {len(train_filtered_indices)} filtered train indices " f"({len(train_filtered_indices)//2} pairs)" ) # Filter idx_lists: keep only indices that exist in filtered set original_count = sum(len(idx_list) for idx_list in train_data_idx_lists) train_data_idx_lists = [ [idx for idx in idx_list if idx in filtered_set] for idx_list in train_data_idx_lists ] filtered_count = sum(len(idx_list) for idx_list in train_data_idx_lists) # Safety check: For DPO, dataset 0 (negative) and dataset 1 (positive) must have same length if len(train_data_idx_lists) >= 2: if len(train_data_idx_lists[0]) != len(train_data_idx_lists[1]): raise ValueError( f"DPO pair mismatch after filtering! " f"Dataset 0 (negative) has {len(train_data_idx_lists[0])} samples, " f"Dataset 1 (positive) has {len(train_data_idx_lists[1])} samples. " f"They must be equal to maintain pair correspondence." ) print_with_time_master( f"Filtered train dataset: {filtered_count} indices out of {original_count} " f"({filtered_count/original_count*100:.1f}% kept)" ) # Print per-task breakdown after filtering by looking at metadata print_with_time_master("\nFiltered train dataset breakdown by task:") task_counts = defaultdict(int) # Count tasks from all filtered indices across all datasets for idx_list in train_data_idx_lists: for idx in idx_list: meta = train_metas[idx] task = meta.get("task", "default") task_counts[task] += 1 # Print sorted by count for task, count in sorted(task_counts.items(), key=lambda x: -x[1]): frac = count / filtered_count if filtered_count > 0 else 0 print_with_time_master(f" {task}: {count:,} samples ({frac*100:.1f}%)") print_with_time_master("") 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, "train": { "data": train_data, "metas": train_metas, "infos": train_info, "artist_to_songs": train_artist_to_songs, "names": train_dataset_names, "weights": train_data_weights, "idx_lists": train_data_idx_lists, "all_idx_lists": sorted([idx for sublist in train_data_idx_lists for idx in sublist]), }, "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, "all_idx_lists": sorted([idx for sublist in val_data_idx_lists for idx in sublist]), }, } # replace the eval loss data sampling info with the larger batch size eval_loss_data_sampling_info = data_sampling_info.copy() eval_loss_data_sampling_info["batch_size"] = eval_loss_batch_size eval_loss_data_sampling_info["batch_size_tokens"] = eval_loss_batch_size_tokens @torch.no_grad() def estimate_loss(): 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, LayerNorm, 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" 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 tmp_dpo = defaultdict(list) 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): idxs, X, Y, loss_start_index_list = get_batch( data_sampling_info, split, dataset_idx=dataset_idx, inference=True, dummy_data=dummy_data, suppress_text=suppress_text, n_offs=None, load_dpo_pair=True, return_idx=True, ) with ctx: # print("DEBUG", text_offset, X.shape, Y.shape) assert (X >= 0).all() Y_pos = Y[1::2] Y_neg = Y[::2] idxs_pos = idxs[1::2] idxs_neg = idxs[::2] ref_pos_loss = get_weighted_loss_from_cache(split, idxs_pos) ref_neg_loss = get_weighted_loss_from_cache(split, idxs_neg) with torch.no_grad(): model.eval() pi_logits_text, pi_logits_semantic, pi_logits_coarse = model( X, y=Y, return_logits=True, last_only=False ) pos_pi_logits_semantic = pi_logits_semantic[1::2] pos_pi_logits_coarse = ( pi_logits_coarse[1::2] if pi_logits_coarse is not None else None ) pi_pos_loss = get_weighted_loss_from_batch_logits( pos_pi_logits_semantic, pos_pi_logits_coarse, Y_pos, loss_start_index_list[1::2], ) neg_pi_logits_semantic = pi_logits_semantic[::2] neg_pi_logits_coarse = ( pi_logits_coarse[::2] if pi_logits_coarse is not None else None ) pi_neg_loss = get_weighted_loss_from_batch_logits( neg_pi_logits_semantic, neg_pi_logits_coarse, Y_neg, loss_start_index_list[::2], ) curr_device = pi_pos_loss.device ref_pos_loss = ref_pos_loss.to(curr_device) ref_neg_loss = ref_neg_loss.to(curr_device) loss, chosen_rewards, rejected_rewards = preference_loss( policy_chosen_logps=-pi_pos_loss, policy_rejected_logps=-pi_neg_loss, reference_chosen_logps=-ref_pos_loss, reference_rejected_logps=-ref_neg_loss, beta=dpo_beta, label_smoothing=0.0, ipo=do_ipo, kto=False, ) print_with_time_master( f"Eval -- {idxs} Policy positive loss: " f"{pi_pos_loss}, {pi_pos_loss.dtype}, " f"Policy negative loss: {pi_neg_loss}, " f"{pi_neg_loss.dtype}" ) print_with_time_master( f"Eval -- Reference positive loss: " f"{ref_pos_loss}, {ref_pos_loss.dtype}, " f"Reference negative loss: {ref_neg_loss}, " f"{ref_neg_loss.dtype}" ) print_with_time_master( f"Eval -- Preference loss: {loss}, " f"{loss.dtype}, Chosen rewards: " f"{chosen_rewards}, {chosen_rewards.dtype}, " f"Rejected rewards: {rejected_rewards}, " f"{rejected_rewards.dtype}" ) loss = loss.mean() reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() loss_val = loss.item() # loss as float. this is a CPU-GPU sync chosen_reward_val = chosen_rewards.mean().float().cpu().numpy() rejected_reward_val = rejected_rewards.mean().float().cpu().numpy() reward_accuracies_val = reward_accuracies.float().cpu().numpy() tmp_dpo[f"{split}/chosen_reward"].append(chosen_reward_val) tmp_dpo[f"{split}/rejected_reward"].append(rejected_reward_val) tmp_dpo[f"{split}/reward_accuracies"].append(reward_accuracies_val) tmp_dpo[f"{split}/dpo_loss"].append(loss_val) with torch.no_grad(): model.eval() loss_dict = model(X, y=Y) # For debugging loss_dict -- in crow the training skips context # target_metas = val_metas if split == "val" else train_metas # data_items = [target_metas[i] for i in idxs_pos] # data_item_tasks = [e.get("task", "") for e in data_items] # # Compute a list of cross-entropy losses, one for every 250 tokens # pi_pos_loss_prev_list: list[float] = [] # seq_len = Y_pos.shape[-1] # step_size = 25 * 30 # for start_idx in range(0, seq_len, step_size): # end_idx = min(start_idx + step_size, seq_len) # logits_slice = pos_pi_logits_semantic[..., start_idx:end_idx, :] # y_slice = Y_pos[..., start_idx:end_idx] # # Flatten for cross_entropy: (batch, tokens, vocab) -> (batch*tokens, vocab) # logits_flat = logits_slice.reshape(-1, logits_slice.size(-1)) # y_flat = y_slice.reshape(-1) # loss = F.cross_entropy( # logits_flat, # y_flat, # ignore_index=-1, # ).item() # pi_pos_loss_prev_list.append(loss) # print( # f"current world {torch.distributed.get_rank()} " # f"loss_dict: {loss_dict}, " # f"data_item_tasks: {data_item_tasks}, " # f"pos_idxs: {idxs_pos}, " # f"loss start index list: {loss_start_index_list[1::2]}, " # f"pi_pos_loss: {pi_pos_loss}, " # f"pi_pos_loss_prev_list: {pi_pos_loss_prev_list}" # ) losses.append([loss_dict[k].item() for k in loss_discount_map.keys()]) # Clean up GPU memory after each eval iteration del X, Y, Y_pos, Y_neg del pi_logits_text, pi_logits_semantic, pi_logits_coarse del pos_pi_logits_semantic, pos_pi_logits_coarse del neg_pi_logits_semantic, neg_pi_logits_coarse del pi_pos_loss, pi_neg_loss, ref_pos_loss, ref_neg_loss del loss, chosen_rewards, rejected_rewards, loss_dict for n, loss_name in enumerate(loss_discount_map.keys()): loss_tensor[n_loss_entry] = float(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)} # update the dpo dict tmp_dpo = {k: np.mean(v) for k, v in tmp_dpo.items()} # add extra loss items out = {k: v for k, v in tmp_out.items()} out[f"train/{loss_prefix}"] = float( np.mean([v for k, v in tmp_out.items() if k.startswith("train/")]) ) out[f"val/{loss_prefix}"] = float(np.mean([v for k, v in tmp_out.items() if k.startswith("val/")])) out.update(tmp_dpo) model.train() for name, module in model.named_modules(): assert module.training # make sure we can undo everything # Clear GPU cache after evaluation to free memory before training resumes torch.cuda.empty_cache() return out train_sampling_params = SamplingParams( inference=True, # DPO IS ALWAYS INFERENCE MODE dummy_data=dummy_data, suppress_text=suppress_text, mask_padding=mask_padding, pack=pack, allow_infill=infill_augment, # Use the script's config flags allow_artist=allow_artist, allow_cover=allow_cover, ) tr_dataset = CustomAudioDataset( data_sampling_info, "train", sampling_params=train_sampling_params, ) tr_dataloader = DataLoader( tr_dataset, shuffle=False, num_workers=0, ) tr_dataloader_iter = iter(tr_dataloader) print_with_time_master(f"Validate random number: {random.random()}") # these are dicts that are # index of the local data file # maps to a dict of losses print_with_time_master("Evaluating") local_ref_loss_lookups = {} local_ref_loss_lookups["train"] = {} local_ref_loss_lookups["val"] = {} all_ref_loss_lookups = [None for _ in range(world_size)] all_seen_idxs = [None for _ in range(world_size)] def get_loss_stored( ref_pos_logits_sem, ref_pos_logits_coarse, curr_y, idxs: list, split: str, loss_start_index_list: list, ): """This will update the index to loss lookup dictionary under the hood. curr_y: batch, ncode, time """ for curr_index, curr_idx in enumerate(idxs): loss_dict = {} loss_dict["z_loss"] = 0 loss_start_index = loss_start_index_list[curr_index] # or 0 for n_semantic in range(semantic_n_codebooks): logits = ref_pos_logits_sem[curr_index][n_semantic] loss_dict[f"semantic_{n_semantic}"] = ( F.cross_entropy( logits.reshape(-1, logits.size(-1))[loss_start_index:], curr_y[curr_index, n_semantic, :].reshape(-1)[loss_start_index:], ignore_index=-1, ) .detach() .float() .cpu() .numpy() .item() ) loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean().detach().float().cpu() for n_coarse in range(coarse_n_codebooks): if ref_pos_logits_coarse is None: continue logits = ref_pos_logits_coarse[curr_index][n_coarse] n2 = semantic_n_codebooks + n_coarse loss_dict[f"coarse_{n_coarse}"] = ( F.cross_entropy( logits.reshape(-1, logits.size(-1))[loss_start_index:], curr_y[curr_index, n2, :].reshape(-1)[loss_start_index:], ignore_index=-1, ) .detach() .float() .cpu() .numpy() .item() ) loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean().detach().float().cpu() loss_dict["orig_idx"] = curr_idx loss_dict["z_loss"] = loss_dict["z_loss"].numpy().item() local_ref_loss_lookups[split][curr_idx] = loss_dict def get_weighted_loss_from_batch_logits( ref_pos_logits_sem, ref_pos_logits_coarse, curr_y, loss_start_index_list ): """Compute the loss, not average across the batch.""" # TODO: somehow using this func kills gradient... curr_device = ref_pos_logits_sem.device batch_loss = torch.zeros(ref_pos_logits_sem.shape[0], requires_grad=True).to(curr_device) loss_discount_facs_tensor = torch.from_numpy(loss_discount_facs) for curr_index in range(ref_pos_logits_sem.shape[0]): loss_start_index = loss_start_index_list[curr_index] # or 0 for n_semantic in range(semantic_n_codebooks): logits = ref_pos_logits_sem[curr_index][n_semantic] cross_entropy_loss = F.cross_entropy( logits.reshape(-1, logits.size(-1))[loss_start_index:], curr_y[curr_index, n_semantic, :].reshape(-1)[loss_start_index:], ignore_index=-1, ) # print_with_time_master(f"semantic loss: {cross_entropy_loss}") batch_loss[curr_index] += cross_entropy_loss * (loss_discount_facs_tensor[n_semantic]) if ref_pos_logits_coarse is None: continue for n_coarse in range(coarse_n_codebooks): logits = ref_pos_logits_coarse[curr_index][n_coarse] n2 = semantic_n_codebooks + n_coarse cross_entropy_loss = F.cross_entropy( logits.reshape(-1, logits.size(-1))[loss_start_index:], curr_y[curr_index, n2, :].reshape(-1)[loss_start_index:], ignore_index=-1, ) # print_with_time_master(f"coarse loss {n_coarse}: {cross_entropy_loss}") batch_loss[curr_index] += cross_entropy_loss * (loss_discount_facs_tensor[n2]) # keep the loss high preicision? # print_with_time_master(f"batch_loss {batch_loss}") return batch_loss def get_weighted_loss_from_cache(split: str, idxs: list): batch_loss = torch.zeros(len(idxs)) for curr_index, curr_idx in enumerate(idxs): loss_dict = ref_loss_lookups[split][curr_idx] # print_with_time_master(f"curr_idx {curr_idx} loss_dict {loss_dict}") curr_loss = sum( v * loss_discount_map[k] for k, v in loss_dict.items() if k != "z_loss" and k != "orig_idx" ) batch_loss[curr_index] = curr_loss return batch_loss # start an evaluation loop first print_with_time_master("Start the ref model loss eval loop.") def update_loss_dict_lookup(split: str): """Pre compute the evaluation loss. Typically we don't shard data for dpo. If we don't shard data -- we will use the world size for loss compute. Otherwise, we need to use local world size for the loss computation. """ print_with_time_master(f"Start evaluating loss for split {split}") if local_data_shard_dir is None: # if we don't shard data -- which is typically true for dpo # we use the global rank... eval_ddp_rank = ddp_rank eval_world_size = world_size else: # for sharded data, we want the local rank eval_ddp_rank = ddp_local_rank eval_world_size = n_gpus_per_node curr_data_idx_lists = train_data_idx_lists if split == "train" else val_data_idx_lists data_idx_lists_flat = sorted([idx for sublist in curr_data_idx_lists for idx in sublist]) est_eval_tot_iter_num = int( math.ceil(len(data_idx_lists_flat) / (eval_loss_batch_size * eval_world_size)) ) print_with_time_master( f"estimated total number of iterations {est_eval_tot_iter_num}," + f"size of the data {len(data_idx_lists_flat)}," + f"size of batch {eval_loss_batch_size}," + f"ddp_rank is {eval_ddp_rank}" ) t0 = time.time() for eval_iter_num in range(est_eval_tot_iter_num): global_batch_row_idx_list = data_idx_lists_flat[ eval_iter_num * eval_loss_batch_size * eval_world_size : (eval_iter_num + 1) * eval_loss_batch_size * eval_world_size ] # subslice only the part needed eval loss batch_row_idx_list = global_batch_row_idx_list[ eval_ddp_rank * eval_loss_batch_size : (eval_ddp_rank + 1) * eval_loss_batch_size ] real_batch_size = len(batch_row_idx_list) # fill to full batch if needed batch_row_idx_list += [0] * (eval_loss_batch_size - real_batch_size) _, X, Y, loss_start_index_list = get_batch( eval_loss_data_sampling_info, split, min_text_offs=0, suppress_text=False, dummy_data=dummy_data, inference=True, return_idx=True, abs_row_idx=batch_row_idx_list, # use the abs idx of the whole dataset ) X = X.to(device) Y = Y.to(device) with torch.no_grad(): model.eval() ref_logits_text, ref_logits_smenatic, ref_logits_coarse = model( X, y=Y, return_logits=True, last_only=False ) # print_with_time_master(f"original idx list {batch_row_idx_list}, output_idx_list {output_idx_list}") t1 = time.time() dt = t1 - t0 t0 = t1 if master_process and (eval_iter_num % log_interval == 0): tokens_per_s = eval_loss_batch_size * block_size / dt print_with_time_master( f"iter {eval_iter_num}/{est_eval_tot_iter_num}:" + f" step_time {dt * 1000:.1f}ms," + f" throughput {tokens_per_s / 1e3:,.0f}k tok/s/node," ) get_loss_stored( ref_logits_smenatic, ref_logits_coarse, curr_y=Y, idxs=batch_row_idx_list, split=split, loss_start_index_list=loss_start_index_list, ) # distribute compute the loss and update local copies model_cache_loss_path = os.path.join(data_dir, f"{model_cache_loss_name}_cached_loss.json") if os.path.exists(model_cache_loss_path): print_with_time_master(f"Loading pre-computed cache loss: {model_cache_loss_path}") with open(model_cache_loss_path, "r") as fp: str_key_ref_loss_lookups = json.load(fp) # remap the json keys...so stupid that json saves integers as strings ref_loss_lookups = {} ref_loss_lookups["train"] = {} ref_loss_lookups["val"] = {} for ref_loss_idx, ref_loss_v in str_key_ref_loss_lookups["train"].items(): ref_loss_lookups["train"][int(ref_loss_idx)] = ref_loss_v for ref_loss_idx, ref_loss_v in str_key_ref_loss_lookups["val"].items(): ref_loss_lookups["val"][int(ref_loss_idx)] = ref_loss_v else: print_with_time_master("Pre-compute cache loss") update_loss_dict_lookup("train") update_loss_dict_lookup("val") dist_barrier() torch.distributed.all_gather_object(all_ref_loss_lookups, local_ref_loss_lookups) print_with_time_master(len(all_ref_loss_lookups)) ref_loss_lookups = {} ref_loss_lookups["train"] = {} ref_loss_lookups["val"] = {} # uppack the dict for sub_loss_lookups in all_ref_loss_lookups: ref_loss_lookups["train"].update(sub_loss_lookups["train"]) ref_loss_lookups["val"].update(sub_loss_lookups["val"]) if master_process: with open(model_cache_loss_path, "w") as fp: json.dump(ref_loss_lookups, fp) print_with_time_master(f"Saving pre-computed cache loss: {model_cache_loss_path}") dist_barrier() print_with_time_master( "Check if loaded correctly: " f"\n train {len(ref_loss_lookups['train'])} cached losses, {len(train_metas)} total metas" f"\n val {len(ref_loss_lookups['val'])} cached losses, {len(val_metas)} total metas" ) # When filtering is applied, we only need losses for the filtered indices # Check that all indices we'll actually use have cached losses if filtered_indices_path is not None: missing_train_indices = [] for idx_list in train_data_idx_lists: for idx in idx_list: if idx not in ref_loss_lookups["train"]: missing_train_indices.append(idx) if missing_train_indices: raise ValueError( f"Filtered training indices are missing from cached losses! " f"Missing {len(missing_train_indices)} indices. " f"First few: {missing_train_indices[:10]}" ) print_with_time_master( f"✓ All {sum(len(idx_list) for idx_list in train_data_idx_lists)} filtered train indices " f"have cached reference losses" ) else: # Without filtering, we need losses for all indices assert len(train_metas) == len( ref_loss_lookups["train"] ), f"Mismatch: {len(train_metas)} metas vs {len(ref_loss_lookups['train'])} cached losses" assert len(val_metas) == len( ref_loss_lookups["val"] ), f"Mismatch: {len(val_metas)} val metas vs {len(ref_loss_lookups['val'])} cached losses" # training loop print_with_time_master(f"Validate random number: {random.random()}") print_with_time_master("training...") 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 effective_tokens_per_s_per_node = 0 running_loss = [] running_chosen_reward = [] running_rejected_reward = [] running_reward_accuracy = [] # with torch.autograd.set_detect_anomaly(True): local_seen_idxs = set() data_idx_lists_flat = sorted([idx for sublist in train_data_idx_lists for idx in sublist]) # cause two kinds of data, pos and neg, we need to divide by 2 local_idx_size = int(math.ceil(len(data_idx_lists_flat) / world_size / 2)) # these are the fixed local idxs, make sure it doesn't overflow local_idxs_fixed = [ (local_idx_size * ddp_rank + i) % len(data_idx_lists_flat) for i in range(local_idx_size) ] # note that since this is a fixed length list, we can't train more than 1 epochs without duplicating! if local_shuffle_data: random.seed(custom_seed_offset) random.shuffle(local_idxs_fixed) # note that we need to divide by 2 to account for pos # get batch load only negative (and positive is a different dataset name) # so the row idx is half of the batch size input_ids = local_idxs_fixed[0 : batch_size // 2] idxs, X, Y, loss_start_index_list = get_batch( data_sampling_info, "train", dummy_data=dummy_data, row_idx=None if shuffle_data else input_ids, suppress_text=suppress_text, load_dpo_pair=True, return_idx=True, ) # fetch the very first batch print_with_time_master(f"First input_ids: {input_ids}") for idx in idxs: local_seen_idxs.add(idx) gpu_memory_monitor.reset_peak_stats() # save X and Y into a npz file # if master_process: # np.savez( # os.path.join("/home/tony", f"dpo_batch_iter_{iter_num}.npz"), # X=X.cpu().numpy(), # Y=Y.cpu().numpy(), # loss_start_index_list=loss_start_index_list, # ) # print_with_time_master(f"Saved DPO batch data for iteration {iter_num}") while True: # determine and set the learning rate for this iteration lr = get_lr(iter_num) 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: dist_barrier() 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_master( f"loss estimation took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)" ) print_with_time_master( 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, "n_tokens": total_tokens_processed, } for k, v in losses.items(): log_dict[k] = v wandb.log(log_dict) dist_barrier() if iter_num > 0: if checkpoint_save_old_format: best_val_loss = save_old_checkpoint( out_dir, model, optimizer, best_val_loss, losses["val/loss"], step_save_iters, time_since_last_loss, model_args=model_args, iter_num=iter_num, n_tokens=total_tokens_processed, debug_val_only=debug_val_only, save_best_ckpt=False, # speed up cause we don't need these for DPO anyways... save_last_ckpt=False, # speed up cause we don't need these for DPO anyways... ) else: best_val_loss = save_checkpoint( out_dir, model, optimizer, best_val_loss, losses["val/loss"], step_save_iters, time_since_last_loss, model_args=model_args, iter_num=iter_num, n_tokens=total_tokens_processed, debug_val_only=debug_val_only, ) dist_barrier() # end if eval test only if eval_only: print_with_time_master("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: # X: Batch, [text, semantic, codec], time # Y: Batch, [semantic, codec], time # positive are odd indices, negative are even indices Y_pos = Y[1::2] Y_neg = Y[::2] idxs_pos = idxs[1::2] idxs_neg = idxs[::2] # will need to return logits, and then use # https://github.com/eric-mitchell/direct-preference-optimization/blob/main/trainers.py#L90 # print_with_time_master(f"Current memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB") ref_pos_loss = get_weighted_loss_from_cache("train", idxs_pos) ref_neg_loss = get_weighted_loss_from_cache("train", idxs_neg) # Alternative... pi_logits_text, pi_logits_smenatic, pi_logits_coarse = model( X, y=Y, return_logits=True, last_only=False ) pos_pi_logits_semantic = pi_logits_smenatic[1::2] pos_pi_logits_coarse = pi_logits_coarse[1::2] if pi_logits_coarse is not None else None neg_pi_logits_semantic = pi_logits_smenatic[::2] neg_pi_logits_coarse = pi_logits_coarse[::2] if pi_logits_coarse is not None else None pi_pos_loss = get_weighted_loss_from_batch_logits( pos_pi_logits_semantic, pos_pi_logits_coarse, Y_pos, loss_start_index_list[1::2] ) pi_neg_loss = get_weighted_loss_from_batch_logits( neg_pi_logits_semantic, neg_pi_logits_coarse, Y_neg, loss_start_index_list[::2] ) curr_device = pi_pos_loss.device ref_pos_loss = ref_pos_loss.to(curr_device) ref_neg_loss = ref_neg_loss.to(curr_device) # print_with_time_master( # f"Policy positive -- normal loss: {pi_pos_loss}, " + # f"reference loss: {ref_pos_loss}, " # ) # print_with_time_master( # f"Policy negative -- normal loss: {pi_neg_loss}, " + # f"reference loss: {ref_neg_loss}, " # ) # need to convert to negative log likelihood, sign flip loss, chosen_rewards, rejected_rewards = preference_loss( policy_chosen_logps=-pi_pos_loss, policy_rejected_logps=-pi_neg_loss, reference_chosen_logps=-ref_pos_loss, reference_rejected_logps=-ref_neg_loss, beta=dpo_beta, label_smoothing=0.0, ipo=do_ipo, kto=False, ) # print_with_time_master(f"Preference loss at {iter_num} is {loss}, {chosen_rewards}, {rejected_rewards}.") reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() # just check the positive loss threshold valid_mask = pi_pos_loss <= loss_clip_threshold if valid_mask.any(): loss = loss[valid_mask].mean() else: print(f"Warning: All samples have pi_pos_loss > loss_clip_threshold, {pi_pos_loss}") loss = loss.mean() # add sft loss: https://github.com/princeton-nlp/SimPO/blob/main/scripts/simpo_trainer.py if sft_loss_scale != 0.0: # Filter out anomalously large losses (likely data anomalies) # Apply at sample level - some samples contribute, others don't if valid_mask.any(): # Only backprop on non-anomalous samples sft_loss = pi_pos_loss[valid_mask].mean() * sft_loss_scale loss += sft_loss loss_val = loss.item() # loss as float. this is a CPU-GPU sync chosen_reward_val = chosen_rewards.mean().float().cpu().numpy() rejected_reward_val = rejected_rewards.mean().float().cpu().numpy() reward_accuracies_val = reward_accuracies.float().cpu().numpy() # Delete tensors immediately after extracting values del chosen_rewards, rejected_rewards, reward_accuracies loss = loss / gradient_accumulation_steps total_tokens_processed += X.shape[0] * X.shape[-1] * world_size rel_tokens_processed += X.shape[0] * X.shape[-1] * world_size # Save current batch info before prefetching next batch current_idxs = idxs current_input_ids = input_ids # immediately async prefetch next batch while model is doing the forward pass on the GPU # input_ids = [ # # divide by 2 to account for pos and neg # local_idx_size * ddp_rank + batch_size // 2 * (local_iter_num + 1) + i # for i in range(batch_size // 2) # ] iter_retrieval_start_idx = (local_iter_num + 1) * batch_size // 2 iter_retrieval_end_idx = (local_iter_num + 2) * batch_size // 2 # Use modulo to wrap around when we reach the end of the list input_ids = [ local_idxs_fixed[i % len(local_idxs_fixed)] for i in range(iter_retrieval_start_idx, iter_retrieval_end_idx) ] idxs, X, Y, loss_start_index_list = get_batch( data_sampling_info, "train", dummy_data=dummy_data, row_idx=None if shuffle_data else input_ids, suppress_text=suppress_text, load_dpo_pair=True, return_idx=True, ) print_with_time_master( f"iter {iter_num}, local_iter_num: {local_iter_num}, idxs: {current_idxs}, input_ids: {current_input_ids}" ) for idx in current_idxs: local_seen_idxs.add(idx) if debug_gradients and wandb_log and master_process: d = { "iter": iter_num, "n_tokens": total_tokens_processed, "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) # print_with_time_master(f"loss {loss}, dtype: {loss.dtype}") # Delete intermediate tensors no longer needed del pi_logits_text, pi_logits_smenatic, pi_logits_coarse del pos_pi_logits_semantic, pos_pi_logits_coarse del neg_pi_logits_semantic, neg_pi_logits_coarse # Only on very first backward, clear cache if iter_num == 0 and micro_step == 0: torch.cuda.empty_cache() # backward pass # print_with_time_master(f"Current memory allocated before loss backward: {torch.cuda.memory_allocated() / 1e9:.2f} GB, {torch.cuda.max_memory_allocated() / 1e9:.2f} GB") # torch.cuda.reset_max_memory_allocated() # print_with_time_master(f"Preference loss pre-backward at {iter_num} is {loss}.") loss.backward() # print_with_time_master(f"Current memory allocated after loss backward: {torch.cuda.memory_allocated() / 1e9:.2f} GB, {torch.cuda.max_memory_allocated() / 1e9:.2f} GB") # torch.cuda.reset_max_memory_allocated() running_loss.append(loss_val) running_chosen_reward.append(chosen_reward_val) running_rejected_reward.append(rejected_reward_val) running_reward_accuracy.append(reward_accuracies_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) # print_with_time_master(f"Current memory allocated after optimizer step: {torch.cuda.memory_allocated() / 1e9:.2f} GB, {torch.cuda.max_memory_allocated() / 1e9:.2f} GB") # torch.cuda.reset_max_memory_allocated() if master_process and wandb_log and grad_norm is not None: wandb.log( { "iter": iter_num, "n_tokens": total_tokens_processed, "misc/grad_norm": grad_norm, } ) if iter_num % 300 == 1: # manually collect garbage in sync to avoid slowdowns over time # this slows down the iteration by ~30%, dont gc too often # do it on mod 1 to avoid it showing up in the logs assert not gc.isenabled() gc.collect() # timing and logging t1 = time.time() dt = t1 - t0 t0 = t1 if iter_num % log_interval == 0 or iter_num == max_iters - 1: if master_process: tok_per_batch = block_size if not pack else batch_size_tokens 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, ) mfu *= tok_per_batch / block_size tokens_per_s = world_size * batch_size * tok_per_batch * gradient_accumulation_steps / dt effective_tokens_per_s_per_node = ( rel_tokens_processed / (time.time() - t_start) / int(round(world_size / n_gpus_per_node)) ) avg_loss = np.mean(running_loss) avg_reward_accuracy = np.mean(running_reward_accuracy) running_loss = [] gpu_mem_stats = gpu_memory_monitor.get_peak_stats() print_with_time_master( f"{color.cyan}iter {iter_num}:" f"{color.green} avg_loss {avg_loss:.3f}," f"{color.yellow} step_time {dt * 1000:.1f}ms," f"{color.magenta} mfu {mfu * 100:.1f}%," f"{color.blue} throughput {tokens_per_s / 1e3:,.0f}k tok/s," f"{color.red} total time {time.time() - t_start:.0f}s," f"{color.yellow} memory {gpu_mem_stats.max_reserved_gib:5.2f}GiB" f"({gpu_mem_stats.max_reserved_pct:.2f}%)" f"{color.reset}" ) if wandb_log: wandb.log( { "iter": iter_num, "n_tokens": total_tokens_processed, "avg_running_loss": avg_loss, "lr": lr, "mfu": mfu * 100, "eff_tok/s/node": effective_tokens_per_s_per_node, "tok/s": tokens_per_s, "chosen_reward": np.mean(running_chosen_reward), "rejected_rewards": np.mean(running_rejected_reward), "reward_accuracies": avg_reward_accuracy, "memory/max_active(GiB)": gpu_mem_stats.max_active_gib, "memory/max_active(%)": gpu_mem_stats.max_active_pct, "memory/max_reserved(GiB)": gpu_mem_stats.max_reserved_gib, "memory/max_reserved(%)": gpu_mem_stats.max_reserved_pct, "memory/num_alloc_retries": gpu_mem_stats.num_alloc_retries, "memory/num_ooms": gpu_mem_stats.num_ooms, } ) gpu_memory_monitor.reset_peak_stats() running_chosen_reward = [] running_rejected_reward = [] running_reward_accuracy = [] iter_num += 1 local_iter_num += 1 # termination conditions if iter_num >= max_iters: print_with_time_master("done.") break dist_barrier() torch.distributed.all_gather_object(all_seen_idxs, local_seen_idxs) # print_with_time_master(f"check this... {all_seen_idxs}") all_seen_idxs = set.union(*all_seen_idxs) print_with_time_master(f"all seen idxs: {len(all_seen_idxs)}") if ddp: destroy_process_group()