import shutil from contextlib import nullcontext import datetime import logging import math import os import random import time import gc import copy import numpy as np from tqdm import tqdm import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy from torch.distributed.device_mesh import DeviceMesh from torch.distributed import destroy_process_group, init_process_group from torch.utils.data import DataLoader from data_utils import SamplingParams, BCTDataset from audioloader import AudioLoaderDataset 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 from utils.bct import BlockUsageStatistics from utils.helpers import ( dist_barrier, hash_string_to_number, load_checkpoint, load_old_state_dict, load_old_optimizer_state_dict, print_with_time, print_with_time_master, save_checkpoint, save_old_checkpoint, verify_preload_model_args, ) from utils.logging import build_gpu_memory_monitor, Color, NoColor from utils.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling 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) os.umask(0o003) # set umask to 0o003 to allow group write for created directories data_dir = None out_dir = None enable_profiling = False dump_folder = "/app/suno/gpt_profiling/" save_traces_folder = "traces" profile_freq = 50 enable_memory_snapshot = False save_memory_snapshot_folder = "memory_snapshots" master_addr = "localhost" master_port = 12355 train_filename = None train_metas_filename = "metas_tr.jsonl" train_info_filename = None val_filename = None val_metas_filename = "metas_val.jsonl" val_info_filename = None tokenizer_filename = "tokenizer_60k.json" debug_val_only = False mock_data = False # fake batches to test get_batch overhead save_debug_build_text = False # save build_text output for debugging debug_output_dir = "/tmp/debug_output" # directory to save debug files preload_checkpoint = None 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 suppress_text = False checkpoint_save_old_format = True # vocab/time constants text_vocab_size = 60_032 # (multiple of 64) text_codebook_size = 60_001 semantic_n_codebooks = 1 # default to 1 for backwards compatibility, set to 4+ for multi-codebook RVQ semantic_vocab_size = 4032 # (multiple of 64) semantic_codebook_size = 4000 semantic_rate_hz = 25 semantic_shift_factor = 5 # temporal shift factor for multi-codebook hierarchical patterns semantic_type = "mert" # semantic encoder type: starts with "mert" or "musicfm" # model params block_size = 32_000 # ~20 minutes of audio t_text = None # not used, just for backwards compatibility t_audio = None # not used, just for backwards compatibility use_rotary_pos_emb = True rope_theta = 500_000 use_qk_norm = True activation_f = "silu" embed_scale_factor = 10.0 # data gradient_accumulation_steps = 1 # used to simulate larger batch sizes batch_size = 2 # if gradient_accumulation_steps > 1, this is the micro-batch size batch_store_size = 4 # number of samples to store in memory # train params semantic_codebook_weight = 1.0 last_codebook_weight = 1.0 z_loss_factor = 1e-5 mask_padding = False pack = True layer_init = True allow_infill = True allow_artist = True allow_cover = True allow_overpaint = True allow_underpaint = True allow_vox = True allow_stem = True allow_sample = True allow_sfx = True prob_sample = 0.1 # Probability of using sample conditioning (0.0-1.0) prob_sample_from_stems = 0.95 # Probability of using stems vs full song for sampling (0.0-1.0) sample_permutation_prob = ( 0.3 # Probability of applying pitch/rate permutation to extracted samples (0.0-1.0) ) prob_stem = 0.5 # Probability of using stem conditioning (0.0-1.0) prob_text_conditioning_pairs = ( 0.5 # Probability of using text conditioning pairs vs original blocks (0.0-1.0) ) prob_token_dropout = 0.0 # Probability of applying token dropout to discrete output blocks (0.0-1.0) token_dropout_pct = 0.5 # Percentage of tokens to mask when token dropout is applied (0.0-1.0) dropout_codebook_pct = 0.1 # probability of masking out higher-level semantic codebooks # TODO (viertthio): optimize contours calculation and warping algorithm to reduce computation overhead in data loading prob_warp_contours = 0.05 # Probability of applying time-warping to contour features (0.0-1.0) # Contour feature calculation probabilities (to reduce computation overhead) prob_loudness_25hz = 0.02 # Probability of calculating loudness_25hz for activity tags prob_contour_loudness_seq = 0.02 # Probability of calculating loudness contour sequence prob_contour_spectral_centroid_seq = 0.02 # Probability of calculating spectral centroid prob_contour_spectral_complexity_seq = 0.02 # Probability of calculating spectral complexity allow_skip = False allow_playlist = True allow_mumble = True allow_lyrics_randomize = True stem_active_sections_weight = 4 # Weight multiplier for metas with stem_active_sections use_text_loss = False prob_text_loss = 0.1 # Probability of text reconstruction when use_text_loss=True use_mmbert = True prob_use_mmbert = 0.05 # Probability of using mmBERT when use_mmbert=True prob_repa_hoot = 0.1 prob_repa_midi = 0.1 use_vae_input = False output_paradigm = "gpt" # "gpt" or "diffusion" output_distribution = "semantic" # "semantic" or "vae" use_hoot = True use_ditto = True max_audio_duration_s = 60 * 30 n_samples_per_meta = 1 use_continuous_semantic_input = False use_noncausal_input = False use_block_type_embeddings = False repa_semantic = False # Add continuous semantic as auxiliary target for layer-specific loss repa_mixed_semantic = False # Add mixed semantic as auxiliary target for stem conditioning repa_hoot = False # Add hoot encoder embeddings as auxiliary target repa_midi = False # Add midi encoder embeddings as auxiliary target repa_layer = -1 # Layer for all repa heads: -1 = final layer, 0-23 = specific layer # eval items custom_seed_offset = 0 eval_interval = 2000 log_interval = 25 eval_iters = 50 eval_only = False # if True, script exits right after the first eval debug_gradients = False # wandb logging wandb_log = False wandb_project = "suno-test" wandb_run_name = "test" wandb_dir = None # 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? interleave_probability = 0.0 # adamw optimizer learning_rate = 5e-4 # max learning rate min_lr = None # None is max_lr/10 max_iters = 100_000 # total number of training iterations warmup_iters = None # None is max_iters/10 lr_decay_iters = None # should be ~= max_iters per Chinchilla step_save_iters = 20_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 # 1 is off # 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 = "full_shard" # ----------------------------------------------------------------------------- config_keys = [ k for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str, type(None))) ] exec(open("configurator.py").read()) # overrides from command line or config file config = {k: globals()[k] for k in config_keys} # will be useful for logging # ----------------------------------------------------------------------------- if val_filename is None and os.path.exists(os.path.join(data_dir, "data_val.bin")): assert train_filename is None train_filename = "data_tr.bin" val_filename = "data_val.bin" use_raw_audio = val_filename is None if min_lr is None: min_lr = learning_rate / 10 if warmup_iters is None: warmup_iters = int(round(max_iters / 10)) # auto set a few params batch_size_tokens = block_size * batch_size if pack: batch_size = 1 # just for infer compatibility if t_text is None: t_text = int(round(block_size * 0.1 / 256) * 256) if t_audio is None: t_audio = int(round(block_size * 0.85 / 256) * 256) 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=2 * 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. local_master_process = ddp_local_rank == 0 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 local_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}.") # Create device mesh for FSDP2 if ddp and fsdp: dp_mesh = DeviceMesh("cuda", list(range(world_size))) # make sure we offset seeds in a clever way seed_offset += ( custom_seed_offset * world_size + 0 if preload_checkpoint is None else hash_string_to_number(preload_checkpoint) ) torch.manual_seed(6006 + seed_offset) random.seed(6006 + seed_offset) np.random.seed(6006 + seed_offset) torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn device_type = "cuda" if "cuda" in device else "cpu" # for later use in torch.autocast ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype] ctx = ( nullcontext() if device_type == "cpu" or fsdp else torch.amp.autocast(device_type=device_type, dtype=ptdtype) ) # fairly arbitrary loss averaging here, eg: # [0.40] + [0.10, 0.09, 0.09, 0.08, 0.07, 0.06, 0.06, 0.05] loss_discount_facs = np.linspace(1.0, last_codebook_weight, semantic_n_codebooks) loss_discount_facs = loss_discount_facs / loss_discount_facs.sum() * semantic_codebook_weight print_with_time_master(f"loss discounts for codebooks: {loss_discount_facs.round(3)}") loss_discount_map = {} if use_text_loss: loss_discount_map["text_output"] = 1.0 if output_distribution == "vae": loss_discount_map["vae_output"] = 1.0 if output_distribution == "semantic": if output_paradigm == "gpt": # Per-codebook loss weights for n in range(semantic_n_codebooks): loss_discount_map[f"semantic_output_{n}"] = loss_discount_facs[n] # Add auxiliary continuous semantic loss for layer-specific target (repa) if repa_semantic: loss_discount_map["repa_semantic_output"] = 1.0 # Add auxiliary mixed semantic loss for stem conditioning (repa) if repa_mixed_semantic: loss_discount_map["repa_mixed_semantic_output"] = 1.0 # Add auxiliary hoot embedding loss (repa) if repa_hoot: loss_discount_map["repa_hoot_output"] = 1.0 # Add auxiliary midi embedding loss (repa) if repa_midi: loss_discount_map["repa_midi_output"] = 1.0 if output_paradigm == "diffusion": rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=ddp_rank) else: rng = None if output_distribution == "continuous_semantic": if output_paradigm == "diffusion": loss_discount_map["continuous_semantic_output"] = 1.0 for k, v in loss_discount_map.items(): print_with_time_master(f"loss weight for {k}: {v:.3f}") 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}") # logging if wandb_log and master_process: import wandb import importlib.metadata import sys wandb_log_cfg = {} wandb_log_cfg["run_config"] = {k: v for k, v in config.items()} wandb_log_cfg["world_size"] = world_size wandb_log_cfg["slurm_id"] = os.environ.get("SLURM_JOB_ID") wandb_log_cfg["slurm_name"] = os.environ.get("SLURM_JOB_NAME") wandb_log_cfg["slurm_script_path"] = os.environ.get("SLURM_SCRIPT_PATH") wandb_log_cfg["checkpoint_dir"] = out_dir wandb_log_cfg["pip_freeze"] = { dist.metadata["Name"]: dist.version for dist in importlib.metadata.distributions() } wandb_log_cfg["python_path"] = sys.executable wandb.init(project=wandb_project, name=wandb_run_name, config=wandb_log_cfg, dir=wandb_dir) wandb.run.log_code(".") print_with_time_master(f"Total world size {world_size}") if not use_raw_audio: raise NotImplementedError("Raw audio not implemented") 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, 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, semantic_type=semantic_type, t_text=t_text, t_audio=t_audio, 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, 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_continuous_semantic_input=use_continuous_semantic_input, use_block_type_embeddings=use_block_type_embeddings, repa_semantic=repa_semantic, repa_mixed_semantic=repa_mixed_semantic, repa_hoot=repa_hoot, repa_midi=repa_midi, repa_layer=repa_layer, ) # 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) gpu_memory_monitor = build_gpu_memory_monitor() # 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 not fsdp: model.to(device) cfg = model.config train_cfg = model.train_config model_args = gptconf.__dict__ print(model_args) # 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() # compile will be applied per-block after AC and before FSDP if compile: import torch._dynamo torch._dynamo.config.cache_size_limit = 512 # 64 iter_num = 0 total_tokens_processed = 0 rel_tokens_processed = 0 # necessary for preloaded checkpoints and throughput calculations best_val_loss = 1e9 if local_cache_dir is not None and local_master_process: shutil.rmtree(local_cache_dir, ignore_errors=True) os.makedirs(local_cache_dir, exist_ok=True) os.chmod(local_cache_dir, 0o774) # 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 FSDP2...") # CORRECT ORDER (following torchtitan): # 1. Activation Checkpointing # 2. Compile per-block # 3. FSDP wrapping # Step 1: Apply activation checkpointing BEFORE compile and FSDP if grad_checkpointing: apply_fsdp_checkpointing(model) # Step 2: Compile per-block BEFORE FSDP (if enabled) if compile: print_with_time_master("compiling each transformer block...") # Use register_module like torchtitan to properly replace the module for layer_id in range(len(model.transformer["h"])): transformer_block = model.transformer["h"][layer_id] # Compile each block individually for better efficiency with repeated structure # Note: fullgraph=True can cause OOM, using default mode for memory efficiency compiled_block = torch.compile(transformer_block, fullgraph=True) model.transformer["h"][layer_id] = compiled_block print_with_time_master("model compilation complete") # Step 3: Apply FSDP2 # Create mixed precision policy mp_policy = MixedPrecisionPolicy( param_dtype=ptdtype, reduce_dtype=ptdtype, ) # Create FSDP config dict fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy} # Add CPU offload if enabled (note: cpu_offload variable not in config, using commented line as reference) # if cpu_offload: # fsdp_config["offload_policy"] = CPUOffloadPolicy() # Reshard after forward (equivalent to FULL_SHARD strategy) reshard_after_forward = sharding_strategy == "full_shard" # Apply FSDP2 to each transformer block for layer_id, transformer_block in enumerate(model.transformer["h"]): fully_shard( transformer_block, **fsdp_config, reshard_after_forward=reshard_after_forward, ) # Apply FSDP2 to the root model fully_shard(model, **fsdp_config, reshard_after_forward=reshard_after_forward) 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}%)" ) 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, preload_checkpoint if local_cache_dir is None else os.path.join(local_cache_dir, "ckpt.pt"), ) # 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) # ++ Create SamplingParams for training ++ train_sampling_params = SamplingParams( inference=False, # Not inference mode during training mock_data=mock_data, suppress_text=suppress_text, mask_padding=mask_padding, pack=pack, allow_infill=allow_infill, # Use the script's config flags allow_artist=allow_artist, allow_cover=allow_cover, allow_stem=allow_stem, allow_sample=allow_sample, prob_sample=prob_sample, prob_sample_from_stems=prob_sample_from_stems, sample_permutation_prob=sample_permutation_prob, max_num_audio_samples=10, # Extract up to 10 audio samples for richer conditioning allow_overpaint=allow_overpaint, allow_underpaint=allow_underpaint, allow_vox=allow_vox, prob_stem=prob_stem, allow_skip=allow_skip, allow_playlist=allow_playlist, allow_mumble=allow_mumble, allow_sfx=allow_sfx, allow_lyrics_randomize=allow_lyrics_randomize, text_loss=use_text_loss, prob_text_loss=prob_text_loss, use_mmbert=use_mmbert, prob_use_mmbert=prob_use_mmbert, use_vae_input=use_vae_input, output_paradigm=output_paradigm, output_distribution=output_distribution, noise_rng=rng, use_hoot=use_hoot, use_ditto=use_ditto, max_audio_duration_s=max_audio_duration_s, n_samples_per_meta=n_samples_per_meta, use_continuous_semantic_input=use_continuous_semantic_input, use_noncausal_input=use_noncausal_input, repa_semantic=repa_semantic, repa_mixed_semantic=repa_mixed_semantic, repa_hoot=repa_hoot, repa_midi=repa_midi, prob_text_conditioning_pairs=prob_text_conditioning_pairs, prob_warp_contours=prob_warp_contours, prob_loudness_25hz=prob_loudness_25hz, prob_contour_loudness_seq=prob_contour_loudness_seq, prob_contour_spectral_centroid_seq=prob_contour_spectral_centroid_seq, prob_contour_spectral_complexity_seq=prob_contour_spectral_complexity_seq, prob_repa_hoot=prob_repa_hoot, prob_repa_midi=prob_repa_midi, prob_token_dropout=prob_token_dropout, token_dropout_pct=token_dropout_pct, dropout_codebook_pct=dropout_codebook_pct, interleave_probability=interleave_probability, ) val_sampling_params = copy.deepcopy(train_sampling_params) val_sampling_params.inference = True print_with_time_master(f"train_sampling_params: {train_sampling_params}") block_usage_stats = BlockUsageStatistics() def make_data_iter(split, is_eval=False): sampling_params = val_sampling_params if is_eval else train_sampling_params metas_filename = train_metas_filename if split == "train" else val_metas_filename info_filename = train_info_filename if split == "train" else val_info_filename oracle_dataset = AudioLoaderDataset( cfg, train_cfg, os.path.join(data_dir, metas_filename), batch_size_tokens, os.path.join(data_dir, tokenizer_filename), device, split=split, dataset_idx=None, info_path=None if info_filename is None else os.path.join(data_dir, info_filename), sampling_params=sampling_params, stem_active_sections_weight=stem_active_sections_weight, ) oracle_dataloader = DataLoader( oracle_dataset, shuffle=False, num_workers=3 if is_eval else 6, prefetch_factor=batch_store_size * (2 if is_eval else 4), batch_size=None, worker_init_fn=lambda x: random.seed(x + seed_offset * 1000), ) oracle_dataloader_iter = iter(oracle_dataloader) dataset = BCTDataset( oracle_dataloader_iter, batch_size_tokens, cfg, sampling_params=sampling_params, device=device, tokenizer_fp=os.path.join(data_dir, tokenizer_filename), save_debug_build_text=save_debug_build_text, debug_output_dir=debug_output_dir, ) dataloader = DataLoader( dataset, shuffle=False, num_workers=0, # no multiprocessing so its in sync across workers batch_size=None, worker_init_fn=lambda x: random.seed(x + seed_offset * 1000), ) def cached_dataloader_iter_fn(dataloader): # preload batch_store_size batches in memory. # #this way we make new batches every 10 steps, smoothing out the load dataloader_iter = iter(dataloader) batch_store = [] while True: if not batch_store: for _ in tqdm(range(batch_store_size), desc="preloading batches", disable=True): try: batch = next(dataloader_iter) except StopIteration: print("dataloader_iter exhausted, resetting") dataloader_iter = iter(dataloader) batch = next(dataloader_iter) batch_store.append(batch) batch = batch_store.pop(0) yield batch return cached_dataloader_iter_fn(dataloader) tr_dataloader_iter = make_data_iter("train", is_eval=False) val_dataloaders = {} for k in ["train", "val"]: val_dataloaders[k] = {} val_dataloaders[k][0] = make_data_iter(k, is_eval=True) train_dataset_names = ["main"] val_dataset_names = ["main"] @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 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 _ in range(effective_eval_iters): packed_sequences = next(val_dataloaders[split][dataset_idx]) with ctx: loss_dict = model(packed_sequences, device=device) 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] = 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)} # 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/")])) model.train() for name, module in model.named_modules(): assert module.training # make sure we can undo everything return out gc.disable() # manually gc to avoid slowdowns https://imbue.com/research/70b-infrastructure/ # training loop print_with_time_master("training...") t0 = time.time() t00 = time.time() t_start = time.time() # absolute time since starting to train t_data = 0 # time spent loading data t_wait = 0 # time spent waiting in loop t_model = 0 # time spent in master node model fw+bw 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 = [] # with torch.autograd.set_detect_anomaly(True): gpu_memory_monitor.reset_peak_stats() with ( maybe_enable_profiling( enable_profiling, dump_folder, save_traces_folder, profile_freq, global_step=iter_num ) as torch_profiler, maybe_enable_memory_snapshot( enable_memory_snapshot, dump_folder, save_memory_snapshot_folder, profile_freq, global_step=iter_num, ) as memory_profiler, ): 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}: 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, ) 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() # Ensure gradients are clean after eval to prevent non-finite gradients optimizer.zero_grad(set_to_none=True) # 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 avg_seq_length = 0.0 t_data = 0 t_model = 0 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 t0_tmp = time.time() packed_sequences = next(tr_dataloader_iter) if micro_step == 0: avg_seq_length = packed_sequences.avg_seq_length_before_crop t_data += time.time() - t0_tmp t0_tmp = time.time() # compute average length of a sequence avg_sequence_length = packed_sequences.n_tokens / len(packed_sequences) with grad_sync_context(): with ctx: # Get per-block losses when updating statistics if micro_step == 0: loss_dict, per_block_losses = model( packed_sequences, device=device, return_block_losses=True ) block_usage_stats.update(packed_sequences, per_block_losses) else: loss_dict = model(packed_sequences, device=device) loss = sum( v * loss_discount_map[k] for k, v in loss_dict.items() if k not in ["z_loss", "repa_reg"] and k in loss_discount_map ) # Handle case where no valid losses exist (sum returns int 0) if isinstance(loss, int): loss = torch.tensor(0.0, device=device, requires_grad=True) loss_val = loss.item() # loss as float. this is a CPU-GPU sync if "z_loss" in loss_dict: # use 1/10 of paper cause of multi-codebook loss = loss + z_loss_factor * loss_dict["z_loss"] if "repa_reg" in loss_dict: # Add repa regularization (already weighted in model) loss = loss + loss_dict["repa_reg"] if wandb_log and master_process: d = { "iter": iter_num, "n_tokens": total_tokens_processed, "avg_sequence_length": avg_sequence_length, **{f"train/{k}": v for k, v in loss_dict.items()}, } wandb.log(d) loss = loss / gradient_accumulation_steps total_tokens_processed += packed_sequences.n_tokens * world_size rel_tokens_processed += packed_sequences.n_tokens * world_size 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) # backward pass loss.backward() running_loss.append(loss_val) t_model += time.time() - t0_tmp # clip the gradient t0_tmp = time.time() if grad_clip != 0.0: # FSDP2 uses standard PyTorch gradient clipping for both DDP and FSDP 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) t_wait = time.time() - t0_tmp 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) / max(1, int(round(world_size / n_gpus_per_node))) ) pct_wait = t_wait / dt * 100 pct_data = t_data / dt * 100 pct_model = t_model / dt * 100 pct_overhead = 100 - pct_wait - pct_data - pct_model avg_loss = np.mean(running_loss) 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}" ) print_with_time_master(block_usage_stats) if wandb_log: log_dict = { "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, "avg_seq_length": avg_seq_length, "perf/pct_wait": pct_wait, "perf/pct_data": pct_data, "perf/pct_model": pct_model, "perf/pct_overhead": pct_overhead, "perf/t_data": t_data, "perf/t_model": t_model, "perf/t_wait": t_wait, "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, } # Add block loss and usage metrics log_dict.update(block_usage_stats.get_wandb_metrics()) wandb.log(log_dict) gpu_memory_monitor.reset_peak_stats() iter_num += 1 local_iter_num += 1 # signals the profiler that the next profiling step has started if torch_profiler: torch_profiler.step() if memory_profiler: memory_profiler.step() # termination conditions if iter_num >= max_iters: print_with_time_master("done.") break dist_barrier() print_with_time_master("removing cache dir.") if local_cache_dir is not None and local_master_process: shutil.rmtree(local_cache_dir, ignore_errors=True) dist_barrier() print_with_time_master("done.") dist_barrier() if ddp: destroy_process_group()