from collections import OrderedDict import contextlib import datetime import json import os import pickle import re import shutil import tempfile import time import random import boto3 import torch import torch.distributed as dist import wandb from g2p_en import G2p from colorama import Fore, Style from torch.distributed import barrier, is_initialized from torch.distributed.checkpoint.state_dict import StateDictOptions from torch.nn.parameter import Parameter FULL_PRECISION_KEY_FRAGMENTS = ("pos_emb", "inv_freq") # ARPAbet → dictionary-style respelling map ARPABET_TO_RESPELLING = { "AA": "ah", # father "AE": "a", # cat "AH": "uh", # about "AO": "aw", # caught "AW": "ow", # now "AY": "eye", # my "B": "b", "CH": "ch", "D": "d", "DH": "th", # this "EH": "eh", # bed "ER": "ur", # bird "EY": "ay", # say "F": "f", "G": "g", "HH": "h", "IH": "i", # bit "IY": "ee", # seat "JH": "j", # jam "K": "k", "L": "l", "M": "m", "N": "n", "NG": "ng", "OW": "oh", # go "OY": "oy", # boy "P": "p", "R": "r", "S": "s", "SH": "sh", "T": "t", "TH": "th", # thin "UH": "oo", # book (best guess for most readers) "UW": "oo", # food "V": "v", "W": "w", "Y": "y", "Z": "z", "ZH": "zh", # genre — if confusing, could use "zher" } COMMON_MISPRONOUNCED_RESPELLINGS = { # Foods & Ingredients "acai": "ahsighee", "quinoa": "keenwah", "gnocchi": "nyokee", "bruschetta": "broosketta", "jalapeno": "halapainyo", "croissant": "cwahsahnt", "tortilla": "torteeya", "hors d'oeuvre": "orderv", "pho": "fuh", "gyro": "yeeroh", "paella": "pah-ehya", # consider "pahaya" # Foreign loanwords (general) "rendezvous": "rahndayvoo", "faux": "foe", "bourgeois": "boorzhwah", "genre": "zhahnruh", "liaison": "leeayzon", "debacle": "daybahkul", "ennui": "onwee", "niche": "neesh", "coup": "koo", "déjà vu": "dayzhahvoo", "au revoir": "oh ruhvwah", "carte blanche": "cart blawnsh", "fiancé": "feeahnsay", "résumé": "rezuhmay", # Place names "worcestershire": "wustersheer", "gloucester": "gloster", "leicester": "lester", "edinburgh": "edinburrah", "new orleans": "nawlins", "thames": "tems", "cairo": "kairo", "dover": "dohver", "reykjavik": "raykyaveek", "melbourne": "melbun", # People & brands "nike": "nikey", "adidas": "ahdeedas", "colonel": "kernel", "hermes": "airmez", "versace": "versahchee", "mozart": "mohzart", "goethe": "gurtuh", "chloe": "kloh-ee", "saoirse": "seer-shuh", "aoife": "ee-fuh", "siobhan": "shivawn", "rachel": "raychel", # Silent letters "bologna": "baloney", "debris": "duhbree", "salmon": "samon", "solder": "sodder", "aisle": "ile", "island": "iland", "plumber": "plummer", "subtle": "suttle", "receipt": "reseet", "honest": "onest", "heir": "air", "gnome": "nome", "knife": "nife", # Other weird ones "segue": "segway", "epitome": "uhpituhmee", "victuals": "vittles", "cache": "cash", "nausea": "nawzhuh", "gesture": "jeshchur", "nuclear": "nookyular", "pecan": "puhkahn", # regional variant "affluent": "aflooent", "et cetera": "etseteruh", "often": "offen", # silent 't' variant "almond": "ahmuhnd", # regional "clothes": "cloze", "wolf": "woof", "jewelry": "joolree", "Wednesday": "wensday", "February": "febyuary", "library": "librerry", "espresso": "espresso", # commonly mispronounced "expresso" "pronunciation": "pruhnunseeayshun", # mroe "anemone": "uhnemonee", "asterisk": "asterisk", # often mispronounced as "asterix" "balloon": "baloon", "banal": "buhnahl", # or "baynul", varies "baroque": "buhroke", "beignet": "benyay", "biopic": "biopic", # often misread as "bio-pick", but "bye-oh-pick" "celtic": "keltic", # for sports teams "chassis": "chassee", "chiaroscuro": "keeahroskuro", "coup de grâce": "koo duh grahss", "dais": "dayis", "demesne": "demeen", "duodenum": "dooahdeenum", # varies "echelon": "eshuhlon", "eider": "ider", "ensign": "ensun", "eschew": "eshchoo", "facetious": "fuhsee shus", "forte": "fortay", # debated but common usage "gaffe": "gaf", "gif": "jif", # controversial but relevant "grotesque": "grotessk", "guillotine": "geeuhteen", "halcyon": "halseeun", "harass": "huh rass", "hegemony": "hehjuhmonee", "hirsute": "hersoot", "hyperbole": "hyperbuhlee", "indict": "indite", "inquiry": "inquiree", "leisure": "leezhur", "lingerie": "lahnzhuray", "macabre": "muhkahbruh", "mauve": "moav", "melee": "maylay", "minutiae": "minoosheeay", "mischief": "misschif", "moi": "mwah", "nauseous": "nawshus", "neither": "nee ther", "omen": "ohmen", "onomatopoeia": "onuhmatuhpeeuh", "paradigm": "paradime", "phlegm": "flem", "precocious": "pruhcoshus", "prescient": "preshunt", "prologue": "prolawg", "pseudonym": "soodunim", "ptarmigan": "tarmigan", "queue": "kyoo", "rapport": "rapor", "reconnaissance": "rekonuhsuns", "reprise": "reprize", "reveille": "revuhlee", "rhetoric": "retuhric", "risqué": "riskay", "rout": "rowt", "sacrosanct": "sackrosankt", "sangfroid": "sangfwah", "schematic": "skeematic", "schism": "sizm", "sergeant": "sarjunt", "sinew": "sinyoo", "soiree": "swahray", "sovereign": "sovrun", "subtlety": "suttultee", "syrup": "sirup", "taut": "tawt", "timbre": "tambr", "tome": "tohm", "treatise": "tretis", "viscous": "viskus", "vogue": "vohg", "yacht": "yot", "zealot": "zellut", } def strip_stress(phoneme): return phoneme.strip("0").strip("1").strip("2") def phonemes_to_respelling(phonemes): stripped = [strip_stress(p) for p in phonemes] return "".join(ARPABET_TO_RESPELLING.get(p, "") for p in stripped) def respell_random_words_in_text(text: str, N: int = 3): g2p = G2p() text = text.replace("'", "") # Tokenize words and punctuation separately words = re.findall(r"\b\w+\b|\W+", text) word_indices = [i for i, w in enumerate(words) if re.match(r"\b\w+\b", w)] selected_indices = random.sample(word_indices, min(N, len(word_indices))) for i in selected_indices: word = words[i] # Check if the word is in our common mispronounced words dictionary if word.lower() in COMMON_MISPRONOUNCED_RESPELLINGS: respelling = COMMON_MISPRONOUNCED_RESPELLINGS[word.lower()] else: # If not in dictionary, use phoneme-based respelling phonemes = g2p(word) respelling = phonemes_to_respelling(phonemes) if respelling: words[i] = respelling return "".join(words) def load_checkpoint(checkpoint_path, use_ema_if_exists=True): raw_state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False) maybe_model_args = None if "model_args" in raw_state_dict.keys(): maybe_model_args = raw_state_dict["model_args"] if "state_dict" in raw_state_dict.keys(): raw_state_dict = raw_state_dict["state_dict"] if "ema_model" in raw_state_dict.keys() and use_ema_if_exists: raw_state_dict = raw_state_dict["ema_model"] elif "model" in raw_state_dict.keys(): raw_state_dict = raw_state_dict["model"] # remove _orig_mod if needed (from compile) tmp_state_dict = {} for k, v in raw_state_dict.items(): k = re.sub(r"_orig_mod\.", "", k) k = re.sub(r"^diffusion\.", "", k) k = re.sub(r"^diffusion\_ema\.", "", k) tmp_state_dict[k] = v state_dict = tmp_state_dict del tmp_state_dict # see if EMA (and use if exists) if any("ema_model." in k for k in state_dict.keys()) and use_ema_if_exists: overwrite_state_dict = {} tmp_state_dict = {} for k, v in state_dict.items(): if k.startswith("ema_model."): overwrite_state_dict[k.replace("ema_model.", "model.")] = v else: tmp_state_dict[k] = v for k, v in overwrite_state_dict.items(): tmp_state_dict[k] = v state_dict = tmp_state_dict del tmp_state_dict, overwrite_state_dict if all("ff.1" not in k for k in state_dict.keys()): # TODO: this fixes legacy stable audio repo checkpoints tmp_state_dict = {} for k, v in state_dict.items(): k = re.sub(r"^(model\.)+", "", k) k = re.sub(r"\.ff\.2\.", ".ff.1.", k) k = re.sub(r"conditioner\.conditioners\.tags_and_lyrics", "text_conditioner", k) k = re.sub(r"conditioner\.conditioners\.semantic_codes", "semantic_conditioner", k) k = re.sub(r"conditioner\.conditioners\.phonemes", "phoneme_conditioner", k) k = re.sub(r"conditioner\.conditioners\.latent\_context", "ctx_conditioner", k) if k in [ "initted", "step", "text_conditioner.proj_out.weight", "text_conditioner.proj_out.bias", "to_global_embed.0.weight", "to_global_embed.2.weight", ]: continue tmp_state_dict[k] = v state_dict = tmp_state_dict del tmp_state_dict # TODO: remove gross hack if ( "phoneme_conditioner.embedding.weight" in state_dict and "phoneme_conditioner.pos_embedding.scale" not in state_dict ): del state_dict["phoneme_conditioner.embedding.weight"] # TODO: remove gross hack if "vae_pad_embed" in state_dict and "ctx_conditioner.proj_out.weight" not in state_dict: del state_dict["vae_pad_embed"] return state_dict, maybe_model_args def _space_repl(m): s = m.group() n_newline = s.count("\n") if n_newline >= 2: return "\n\n" elif n_newline == 1: return "\n" return " " def simplify_whitespace(text, retain_newlines=True): """simplify while respecting up to 2 newlines""" if retain_newlines: text = re.sub(r"\s+", _space_repl, text).strip() else: text = re.sub(r"\s+", " ", text).strip() return text def is_ddp(): return int(os.environ.get("RANK", -1)) != -1 def is_master(): if is_ddp(): return int(os.environ["RANK"]) == 0 return True def print_with_time(content): """Print the content with the current time.""" print(f"[{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}]: {content}") def print_with_time_master(content): if is_master(): print_with_time(content) def dist_barrier(): if is_initialized(): barrier() def _downcast_state_dict(state_dict): return OrderedDict( { k: v if any(s in k for s in FULL_PRECISION_KEY_FRAGMENTS) else v.to(torch.bfloat16) for k, v in state_dict.items() } ) def _nonblocking_move(state_dict, device): futures = OrderedDict() for k, v in state_dict.items(): if isinstance(v, (torch.Tensor, Parameter)): futures[k] = v.to(device, non_blocking=True) else: futures[k] = v return futures def _move_state_dict_to_device(state_dict, device): """Fast way of moving a state dict to a device""" if isinstance(device, str): device = torch.device(device) if device.type == "cuda": with torch.cuda.stream(torch.cuda.Stream()): futures = _nonblocking_move(state_dict, device) torch.cuda.current_stream().synchronize() else: futures = _nonblocking_move(state_dict, device) return OrderedDict(futures) class FSDP_EMA: def __init__(self, model, decays=None, warmup_steps=None, update_every=None, cpu_offload=False): if decays is None: decays = [0.9999] if warmup_steps is None: warmup_steps = [0] * len(decays) if update_every is None: update_every = [1] * len(decays) if not (len(decays) == len(warmup_steps) == len(update_every)): raise ValueError("The lengths of decays, warmup_steps, and update_every must be the same.") self.model = model # Adjust decays based on update frequency self.decays = [decay**freq for decay, freq in zip(decays, update_every)] self.warmup_steps = warmup_steps self.update_every = update_every self.num_scales = len(decays) self.cpu_offload = cpu_offload self.device = None # Track which parameters don't need EMA updates self.skip_keys = set() if dist.get_rank() == 0: for name, param in model.named_parameters(): if not param.requires_grad: self.skip_keys.add(name.replace("_fsdp_wrapped_module.", "")) for name, _ in model.named_buffers(): self.skip_keys.add(name.replace("_fsdp_wrapped_module.", "")) # Initialize EMA states for each scale # Note: this stores everything in float32 for better precision self.ema_states = [] model_state = torch.distributed.checkpoint.state_dict.get_model_state_dict( self.model, options=StateDictOptions( full_state_dict=True, cpu_offload=True, broadcast_from_rank0=False, # Don't need broadcast since we only use rank 0 ), ) if dist.get_rank() == 0: self.ema_states = [ {k: v.clone() for k, v in model_state.items()} for _ in range(self.num_scales) ] # remove grads just in case for ema_state in self.ema_states: for param in ema_state.values(): param.requires_grad_(False) def update(self, step=0): # Check if any updates are needed this step updates_needed = [step % freq == 0 for freq in self.update_every] if not any(updates_needed): return with torch.no_grad(): model_state = torch.distributed.checkpoint.state_dict.get_model_state_dict( self.model, options=StateDictOptions( full_state_dict=True, # Need full state dict since we're doing EMA cpu_offload=self.cpu_offload, # Keep on CPU broadcast_from_rank0=False, # Don't need broadcast since we only use rank 0 ), ) if dist.get_rank() == 0: for i, needs_update in enumerate(updates_needed): if not needs_update: continue if self.device is None: self.device = next(iter(model_state.values())).device if not self.cpu_offload: self.ema_states[i] = _move_state_dict_to_device(self.ema_states[i], self.device) decay = self.decays[i] warmup = self.warmup_steps[i] current_decay = decay if step > warmup else 0.0 for name, ema_param in self.ema_states[i].items(): if name in model_state: model_param = model_state[name] assert model_param.shape == ema_param.shape if current_decay == 0.0 or name in self.skip_keys: ema_param.copy_(model_param) else: ema_param.mul_(current_decay).add_(model_param * (1 - current_decay)) if not self.cpu_offload: self.ema_states[i] = _move_state_dict_to_device(self.ema_states[i], "cpu") def state_dict(self): states = [] for i, ema_state in enumerate(self.ema_states): model_state = _downcast_state_dict( {k: v.detach().clone().cpu() for k, v in ema_state.items()} ) states.append( { "state_dict": model_state, "decay": self.decays[i], "warmup_steps": self.warmup_steps[i], "update_every": self.update_every[i], } ) return states # Logging functions def log_training_metrics( master_process, loss, lr_scheduler, epoch, tot_elapsed_time, elapsed_time, batch, world_size, global_step, std_data, grad_norm, dpo_losses={}, # dictionary that carries more loss terms distill_losses={}, # dictionary that carries more loss terms ): metrics = torch.tensor([loss.item(), std_data, grad_norm]).to(loss.device) for dpo_loss_name, dpo_loss_value in dpo_losses.items(): metrics = torch.cat([metrics, torch.tensor([dpo_loss_value]).to(loss.device)]) for distill_loss_name, distill_loss_value in distill_losses.items(): metrics = torch.cat([metrics, torch.tensor([distill_loss_value]).to(loss.device)]) dist.all_reduce(metrics, op=dist.ReduceOp.SUM) metrics = metrics / world_size # Extract base metrics base_metrics_count = 3 # loss, std_data, grad_norm avg_loss, avg_std_data, avg_grad_norm = metrics[:base_metrics_count] # Extract additional metrics dynamically offset = base_metrics_count avg_dpo_losses = metrics[offset : offset + len(dpo_losses)] if dpo_losses else [] offset += len(dpo_losses) avg_distill_losses = metrics[offset:] if distill_losses else [] if master_process: # TODO: in prefix seq_len here is everything, not just main audio # Handle both memmap format (batch[0] is tensor) and dynamic format (batch is dict) if isinstance(batch, dict): # Dynamic format - batch is already processed bs, _, seq_len = batch["diffusion_input"].shape else: # Memmap format - batch[0] is the diffusion_input tensor bs, _, seq_len = batch[0].shape log_data = { "train/loss": avg_loss.item(), "train/lr": lr_scheduler.get_last_lr()[0], "train/epoch": epoch, "train/tot_elapsed_time": round(tot_elapsed_time), "train/iterations_per_second": round(1 / elapsed_time, 2), "train/samples_per_second": round(bs * world_size / elapsed_time, 2), "train/k_tokens_per_second": round(bs * seq_len * world_size / elapsed_time / 1000, 2), "train/std_data": round(avg_std_data.item(), 3), "train/grad_norm": round(avg_grad_norm.item(), 4), "trainer/global_m_tokens": round(bs * seq_len * world_size * global_step / 1000000, 2), "trainer/global_samples": bs * world_size * global_step, "trainer/global_step": global_step, } for dpo_loss_name, dpo_loss_value in zip(dpo_losses.keys(), avg_dpo_losses): log_data[f"train/{dpo_loss_name}"] = dpo_loss_value.item() for distill_loss_name, distill_loss_value in zip(distill_losses.keys(), avg_distill_losses): log_data[f"train/{distill_loss_name}"] = distill_loss_value.item() wandb.log(log_data) print_with_time_master( f"{Fore.GREEN}Training metrics:{Style.RESET_ALL} " + ", ".join( [ ( f"{Fore.YELLOW}{k.replace('train/', '')}:{Style.RESET_ALL} {v:.6f}" if isinstance(v, float) else f"{Fore.YELLOW}{k.replace('train/', '')}:{Style.RESET_ALL} {v}" ) for k, v in log_data.items() ] ) ) def log_validation_metrics( master_process, avg_val_loss, batch, world_size, global_step, tot_elapsed_time, dpo_losses={}, # dictionary that carries more loss terms ): if master_process: # Handle both memmap format (batch[0] is tensor) and dynamic format (batch is dict) if isinstance(batch, dict): # Dynamic format - batch is already processed bs, _, seq_len = batch["diffusion_input"].shape else: # Memmap format - batch[0] is the diffusion_input tensor bs, _, seq_len = batch[0].shape val_log_data = { "val/loss": avg_val_loss, "trainer/global_tokens": bs * seq_len * world_size * global_step, "trainer/global_samples": bs * world_size * global_step, "trainer/global_step": global_step, "train/tot_elapsed_time": round(tot_elapsed_time), } for dpo_loss_name, dpo_loss_value in dpo_losses.items(): val_log_data[f"val/{dpo_loss_name}"] = dpo_loss_value wandb.log(val_log_data) print_with_time_master( f"{Fore.MAGENTA}Validation metrics:{Style.RESET_ALL} " + ", ".join( [ ( f"{Fore.YELLOW}{k}:{Style.RESET_ALL} {v:.4f}" if isinstance(v, float) else f"{Fore.YELLOW}{k}:{Style.RESET_ALL} {v}" ) for k, v in val_log_data.items() ] ) ) def _clean_state_dict(state_dict): return { k.replace("_orig_mod.", "").replace("_fsdp_wrapped_module.", ""): v for k, v in state_dict.items() } def save_checkpoint( out_dir, model, ema_model, optimizer, best_val_loss, current_val_loss, step_save_iters, run_config=None, model_args=None, iter_num=None, n_tokens=None, save_best_ckpt=True, save_last_ckpt=True, ): dist_barrier() t_save = time.time() model_state, optim_state = torch.distributed.checkpoint.state_dict.get_state_dict( model, optimizer, options=StateDictOptions( full_state_dict=True, # Need full state dict since we're doing EMA cpu_offload=True, # Keep on CPU broadcast_from_rank0=False, # Don't need broadcast since we only use rank 0 ), ) model_state = _downcast_state_dict(model_state) # Manually handle EMA state dict ema_states = None if is_master() and ema_model is not None: ema_states = ema_model.state_dict() if is_master(): # only write state dicts on rank 0 assert model_state is not None, "some sort of distributed bug" assert optim_state is not None, "some sort of distributed bug" checkpoint = { "model": model_state, "optimizer": optim_state, "run_config": run_config, "model_args": model_args, "iter_num": iter_num, "n_tokens": n_tokens, "best_val_loss": current_val_loss, } if ema_states is not None: checkpoint["ema_models"] = ema_states print_with_time(f"saving checkpoint to {out_dir}") os.makedirs(out_dir, exist_ok=True) if current_val_loss < best_val_loss and save_best_ckpt: torch.save(checkpoint, os.path.join(out_dir, "best_ckpt.pt")) shutil.copy( os.path.join(out_dir, "best_ckpt.pt"), os.path.join(out_dir, "last_ckpt.pt"), ) elif save_last_ckpt: torch.save(checkpoint, os.path.join(out_dir, "last_ckpt.pt")) # make infer with ema over base, dtype casting and no optimizer infer_sd = {k: checkpoint[k] for k in ["model_args", "best_val_loss"]} if "ema_models" in checkpoint: # take middle entry from list n_middle = (len(checkpoint["ema_models"]) - 1) // 2 infer_sd["model"] = checkpoint["ema_models"][n_middle]["state_dict"] else: infer_sd["model"] = checkpoint["model"] # probably always want to save the last infer checkpoint torch.save(infer_sd, os.path.join(out_dir, "last_ckpt_infer.pt")) if iter_num % step_save_iters == 0 and iter_num > 0: torch.save( checkpoint, os.path.join(out_dir, f"step_{iter_num}_ckpt.pt"), ) torch.save(infer_sd, os.path.join(out_dir, f"step_{iter_num}_infer.pt")) estimation_time = time.time() - t_save # del model_state, optim_state, original_osd torch.cuda.empty_cache() if current_val_loss < best_val_loss: best_val_loss = current_val_loss dist_barrier() print_with_time_master(f"saving took {estimation_time:.1f} seconds.") return best_val_loss # the number of warmup steps before the active step in each profiling cycle WARMUP = 3 # how much memory allocation/free ops to record in memory snapshots MEMORY_SNAPSHOT_MAX_ENTRIES = 100000 @contextlib.contextmanager def maybe_enable_profiling( enable_profiling: bool, dump_folder: str, save_traces_folder: str, profile_freq: int, *, global_step: int = 0, ): if enable_profiling: trace_dir = os.path.join(dump_folder, save_traces_folder) rank = torch.distributed.get_rank() def trace_handler(prof): curr_trace_dir_name = "iteration_" + str(prof.step_num) curr_trace_dir = os.path.join(trace_dir, curr_trace_dir_name) if not os.path.exists(curr_trace_dir): os.makedirs(curr_trace_dir, exist_ok=True) print_with_time_master(f"Dumping traces at step {prof.step_num}") begin = time.monotonic() trace_path = f"{curr_trace_dir}/rank{rank}_trace.json.gz" prof.export_chrome_trace(trace_path) print_with_time_master(f"Finished dumping traces in {time.monotonic() - begin:.2f} seconds") print_with_time_master(f"Trace saved at: {trace_path}") # Profiling is a heavy operation which could cost very different amount of time # across all ranks. Insert a barrier to make sure all ranks have finished profiling # before moving on. # TODO: Can we find a cleaner way? torch.distributed.barrier() print_with_time_master(f"Profiling active. Traces will be saved at {trace_dir}") if not os.path.exists(trace_dir): os.makedirs(trace_dir, exist_ok=True) warmup, active = WARMUP, 1 wait = profile_freq - (active + warmup) assert wait >= 0, "profile_freq must be greater than or equal to warmup + active" with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active), on_trace_ready=trace_handler, record_shapes=True, ) as torch_profiler: torch_profiler.step_num = global_step yield torch_profiler else: torch_profiler = contextlib.nullcontext() yield None @contextlib.contextmanager def maybe_enable_memory_snapshot( enable_snapshot: bool, dump_folder: str, save_memory_snapshot_folder: str, profile_freq: int, *, global_step: int = 0, ): if enable_snapshot: snapshot_dir = os.path.join(dump_folder, save_memory_snapshot_folder) if not os.path.exists(snapshot_dir): os.makedirs(snapshot_dir, exist_ok=True) rank = torch.distributed.get_rank() class MemoryProfiler: def __init__(self, step_num: int, freq: int): torch.cuda.memory._record_memory_history(max_entries=MEMORY_SNAPSHOT_MAX_ENTRIES) # when resume training, we start from the last step self.step_num = step_num self.freq = freq def step(self, exit_ctx: bool = False): self.step_num += 1 if not exit_ctx and self.step_num % self.freq != 0: return if not exit_ctx: curr_step = self.step_num dir_name = f"iteration_{curr_step}" else: # dump as iteration_0_exit if OOM at iter 1 curr_step = self.step_num - 1 dir_name = f"iteration_{curr_step}_exit" curr_snapshot_dir = os.path.join(snapshot_dir, dir_name) if not os.path.exists(curr_snapshot_dir): os.makedirs(curr_snapshot_dir, exist_ok=True) print_with_time_master(f"Dumping memory snapshot at step {curr_step}") begin = time.monotonic() snapshot_path = f"{curr_snapshot_dir}/rank{rank}_memory_snapshot.pickle" with open(snapshot_path, "wb") as output: pickle.dump(torch.cuda.memory._snapshot(), output) print_with_time_master( f"Finished dumping memory snapshot to {snapshot_path} in {time.monotonic() - begin:.2f} seconds" ) torch.distributed.barrier() print_with_time_master(f"Memory profiler active. Snapshot will be saved at {snapshot_dir}") profiler = MemoryProfiler(global_step, profile_freq) try: yield profiler except torch.OutOfMemoryError as e: profiler.step(exit_ctx=True) else: yield None def read_jsonl(filepath): data = [] with open(filepath) as f: for line in f: line = line.strip() if len(line) == 0: continue m = json.loads(line) data.append(m) return data def write_jsonl(data, filepath, do_append=False): openarg = "a" if do_append else "w" with open(filepath, openarg) as f: for d in data: f.write(json.dumps(d, ensure_ascii=False) + "\n") def get_filename(filepath, keep_ext=True): if "http" in filepath: clean_filepath = filepath.split("?")[0] else: clean_filepath = filepath filename = clean_filepath.split("/")[-1] if "." not in filename: raise ValueError("filename does not seem to contain a period.") m = re.search(r"(.+)\.([^\.]+)$", filename) if not m: raise ValueError(f"filename could not be parsed for `{filepath}`") filename = m.group(1) file_ext = m.group(2).lower() if len(file_ext) > 10: raise ValueError(f"file extension suspiciously long for `{filepath}`") if keep_ext: filename = filename + "." + file_ext return filename S3_BUCKET_PATH_RE = r"s3\:\/\/(.+?)\/" def _parse_s3_filepath(s3_filepath): bucket_name = re.search(S3_BUCKET_PATH_RE, s3_filepath).group(1) rel_s3_filepath = re.sub(S3_BUCKET_PATH_RE, "", s3_filepath) return bucket_name, rel_s3_filepath def download_s3_file( from_s3_filepath, to_local_filepath, ): bucket_name, from_rel_s3_filepath = _parse_s3_filepath(from_s3_filepath) client = boto3.client("s3") client.download_file(bucket_name, from_rel_s3_filepath, to_local_filepath) def _get_file_ext(filepath): filename = get_filename(filepath, keep_ext=True) file_ext = filename.split(".")[-1] return file_ext def read_from_s3(filepath, read_f): with tempfile.TemporaryDirectory() as tmp_dir: tmp_filepath = os.path.join(tmp_dir, f"file.{_get_file_ext(filepath)}") download_s3_file(filepath, tmp_filepath) data = read_f(tmp_filepath) return data