from pathlib import Path import torch import numpy as np import socket import argparse import os import json import subprocess from hps import Hyperparams, parse_args_and_update_hparams, add_imle_arguments from helpers.utils import ( configure_inductor_for_low_memory_compile, is_dist_avail_and_initialized, logger, maybe_download, ) from data import mkdir_p from contextlib import contextmanager import torch.distributed as dist # from apex.optimizers import FusedAdam as AdamW from torch.optim import AdamW from models import IMLE from torch.nn.parallel.distributed import DistributedDataParallel from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR, SequentialLR import math class AbsoluteCosineAnnealingLR(CosineAnnealingLR): """CosineAnnealingLR using the absolute formula instead of the incremental one. PyTorch's default CosineAnnealingLR.get_lr() computes the next LR incrementally from the *current* optimizer group['lr']. Any external modification of the optimizer LR (e.g. a resample cool-down factor) therefore poisons all subsequent steps, collapsing the schedule to ~0. This subclass replaces get_lr() with the closed-form cosine formula that depends only on base_lrs, eta_min, last_epoch, and T_max. """ def get_lr(self): return [ self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * self.last_epoch / self.T_max)) / 2 for base_lr in self.base_lrs ] import random from helpers.utils import is_main_process, get_world_size, get_rank from torch.nn.parallel import DistributedDataParallel as DDP import torch.nn as nn def update_ema(imle, ema_imle, ema_rate): for p1, p2 in zip(imle.parameters(), ema_imle.parameters()): p2.data.mul_(ema_rate) p2.data.add_(p1.data * (1 - ema_rate)) def as_plain_nn(model): """Returns the model without optimization wrappers.""" if isinstance(model, torch._dynamo.eval_frame.OptimizedModule): return as_plain_nn(model._orig_mod) elif isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel): return as_plain_nn(model.module) elif isinstance(model, nn.DataParallel): return model.module else: return model def map_saved_by_type(x): if isinstance(x, nn.Module): return as_plain_nn(x).state_dict() elif hasattr(x, "state_dict"): return x.state_dict() else: return x def _atomic_torch_save(obj, filepath): """Crash-safe save: write to tmp, fsync data + parent dir, then rename. On Lustre/network FS, torch.save() returning is not enough -- the data may still be in page cache. fsync() forces the bytes to stable storage before the rename swaps the file in. Without this, a node crash mid-checkpoint can leave a zero-byte / truncated file even though .tmp -> filepath was atomic. """ tmp = filepath + ".tmp" with open(tmp, "wb") as f: torch.save(obj, f) f.flush() try: os.fsync(f.fileno()) except OSError: pass os.replace(tmp, filepath) try: d = os.open(os.path.dirname(filepath) or ".", os.O_RDONLY) try: os.fsync(d) finally: os.close(d) except OSError: pass def _verify_torch_load(filepath): """Returns True iff the file can be loaded back by torch (basic integrity).""" try: torch.load(filepath, map_location="cpu") return True except Exception: return False def save_model(path, imle, ema_imle, optimizer, scheduler, scaler, H): model_state = map_saved_by_type(imle) ema_state = map_saved_by_type(ema_imle) optim_state = map_saved_by_type(optimizer) sched_state = map_saved_by_type(scheduler) scaler_state = map_saved_by_type(scaler) # EMA first: if anything below crashes, the *checkpoint that matters for # eval* is already on disk and valid. _atomic_torch_save(ema_state, f"{path}-model-ema.th") _atomic_torch_save(model_state, f"{path}-model.th") _atomic_torch_save(optim_state, f"{path}-opt.th") _atomic_torch_save(sched_state, f"{path}-sched.th") _atomic_torch_save(scaler_state, f"{path}-scaler.th") # Post-write verification: if any of the freshly-written checkpoints is # unreadable, raise loudly so we can investigate instead of silently # corrupting the rolling 'latest-*' set. for suffix in ("model-ema.th", "model.th", "opt.th", "sched.th", "scaler.th"): fp = f"{path}-{suffix}" if not _verify_torch_load(fp): raise RuntimeError(f"Checkpoint verification FAILED for {fp}") from_log = os.path.join(H.save_dir, 'log.jsonl') to_log = f'{os.path.dirname(path)}/{os.path.basename(path)}-log.jsonl' subprocess.check_output(['cp', from_log, to_log]) def accumulate_stats(stats, frequency): z = {} for k in stats[-1]: if k in ['distortion_nans', 'rate_nans', 'skipped_updates', 'gcskip', 'loss_nans']: z[k] = np.sum([a[k] for a in stats[-frequency:]]) elif k == 'grad_norm': vals = [a[k] for a in stats[-frequency:]] finites = np.array(vals)[np.isfinite(vals)] if len(finites) == 0: z[k] = 0.0 else: z[k] = np.max(finites) elif k == 'loss': vals = [a[k] for a in stats[-frequency:]] finites = np.array(vals)[np.isfinite(vals)] z['loss'] = np.mean(vals) z['loss_filtered'] = np.mean(finites) elif k == 'iter_time': z[k] = stats[-1][k] if len(stats) < frequency else np.mean([a[k] for a in stats[-frequency:]]) else: z[k] = np.mean([a[k] for a in stats[-frequency:]]) return z def linear_warmup(warmup_iters): def f(iteration): return 1.0 if iteration > warmup_iters else iteration / warmup_iters return f def distributed_maybe_download(path, local_rank, mpi_size): if not path.startswith('gs://'): return path filename = path[5:].replace('/', '-') with first_rank_first(local_rank, mpi_size): fp = maybe_download(path, filename) return fp @contextmanager def first_rank_first(local_rank, mpi_size): if mpi_size > 1 and local_rank > 0: dist.barrier() try: yield finally: if mpi_size > 1 and local_rank == 0: dist.barrier() def setup_save_dirs(H): H.save_dir = os.path.join(H.save_dir, H.desc) mkdir_p(H.save_dir) H.logdir = os.path.join(H.save_dir, 'log') def set_seed(seed): np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) random.seed(seed) def set_up_hyperparams(s=None): H = Hyperparams() parser = argparse.ArgumentParser() parser = add_imle_arguments(parser) parse_args_and_update_hparams(H, parser, s=s) setup_save_dirs(H) set_seed(H.seed) logprint = logger(H.logdir) np.random.seed(H.seed) torch.manual_seed(H.seed) torch.cuda.manual_seed(H.seed) random.seed(H.seed) return H, logprint def restore_params(model, path, local_rank, mpi_size, map_ddp=True, map_cpu=False, strict=True): state_dict = torch.load(distributed_maybe_download(path, local_rank, mpi_size), map_location='cpu') if map_ddp: new_state_dict = {} l = len('module.') for k in state_dict: if k.startswith('module.'): new_state_dict[k[l:]] = state_dict[k] else: new_state_dict[k] = state_dict[k] state_dict = new_state_dict # torch.compile wraps the module in OptimizedModule whose state_dict is # prefixed with '_orig_mod.'. Saved checkpoints from this codebase have # the prefix stripped (see save_model -> map_saved_by_type -> as_plain_nn), # but defensively normalise both sides so a checkpoint can be loaded into # any wrapping order (compile(DDP(base)), DDP(compile(base)), or plain). PFX = '_orig_mod.' if any(k.startswith(PFX) for k in state_dict): state_dict = {(k[len(PFX):] if k.startswith(PFX) else k): v for k, v in state_dict.items()} # Match the inverse case too: model expects '_orig_mod.' but checkpoint # has bare keys. Detect by sniffing one model parameter name. try: any_model_key = next(iter(model.state_dict().keys()), '') if any_model_key.startswith(PFX) and not any(k.startswith(PFX) for k in state_dict): state_dict = {PFX + k: v for k, v in state_dict.items()} except Exception: pass model.load_state_dict(state_dict, strict=strict) def restore_log(path, local_rank, mpi_size): loaded = [json.loads(l) for l in open(distributed_maybe_download(path, local_rank, mpi_size))] try: cur_eval_loss = float('inf') for z in loaded: if 'type' in z and z['type'] == 'train_loss' and 'best_fid' in z: cur_eval_loss = min(cur_eval_loss, z['best_fid']) except: cur_eval_loss = float('inf') starting_epoch = max([z['epoch'] for z in loaded if 'type' in z and z['type'] == 'train_loss']) iterate = max([z['step'] for z in loaded if 'type' in z and z['type'] == 'train_loss']) return cur_eval_loss, iterate, starting_epoch def load_imle(H, logprint): local_rank = get_rank() device = torch.device("cuda") imle = IMLE(H) imle.to(device) if H.restore_path: if(is_main_process()): logprint(f'Restoring imle from {H.restore_path}') restore_params(imle, H.restore_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size, strict=H.load_strict) ema_imle = IMLE(H) ema_imle = ema_imle.to(device) # Move to the correct device. if H.restore_ema_path: if(is_main_process()): logprint(f'Restoring ema imle from {H.restore_ema_path}') try: restore_params(ema_imle, H.restore_ema_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size, strict=H.load_strict) except Exception as e: if is_main_process(): logprint(f'WARNING: Failed to load EMA from {H.restore_ema_path}: {type(e).__name__}: {e}. Falling back to main model weights.') ema_imle.load_state_dict(imle.state_dict()) else: ema_imle.load_state_dict(imle.state_dict()) ema_imle.requires_grad_(False) ema_imle.eval() ddp_dev = torch.cuda.current_device() if(is_dist_avail_and_initialized()): imle = DDP(imle, device_ids=[ddp_dev], output_device=ddp_dev, gradient_as_bucket_view=True, static_graph=True ) if H.compile: configure_inductor_for_low_memory_compile() imle = torch.compile(imle) ema_imle = torch.compile(ema_imle) return imle, ema_imle def load_opt(H, imle, logprint): optimizer = AdamW(imle.parameters(), weight_decay=H.wd, lr=H.lr, betas=(H.adam_beta1, H.adam_beta2), eps=H.adam_eps) scheduler1 = LambdaLR(optimizer, lr_lambda=linear_warmup(H.warmup_iters)) total_iters = getattr(H, 'total_iters', None) if total_iters is None: default_train_size = 50000 if H.dataset == 'cifar10' else max(H.n_batch, 1) subset_len = H.subset_len if getattr(H, 'subset_len', -1) and H.subset_len > 0 else default_train_size steps_per_epoch = max(1, (subset_len + H.n_batch - 1) // H.n_batch) accum_steps = max(1, H.accumulation_steps) total_iters = max(H.warmup_iters + 1, (H.num_epochs * steps_per_epoch + accum_steps - 1) // accum_steps) H.total_iters = total_iters cosine_iters = max(1, total_iters - H.warmup_iters) scheduler2 = AbsoluteCosineAnnealingLR(optimizer, T_max=cosine_iters, eta_min=0.1 * H.lr) scheduler = SequentialLR(optimizer, schedulers=[scheduler1, scheduler2], milestones=[H.warmup_iters]) scaler = torch.GradScaler(device="cuda") if H.restore_optimizer_path: if(is_main_process()): logprint(f'Restoring optimizer from {H.restore_optimizer_path}') try: optimizer.load_state_dict( torch.load(H.restore_optimizer_path, map_location='cpu')) except Exception as e: if is_main_process(): logprint(f'WARNING: Failed to load optimizer from {H.restore_optimizer_path}: {type(e).__name__}: {e}. Using fresh optimizer.') if H.restore_scheduler_path: if(is_main_process()): logprint(f'Restoring scheduler from {H.restore_scheduler_path}') try: scheduler.load_state_dict( torch.load(H.restore_scheduler_path, map_location='cpu', weights_only=False)) except Exception as e: if is_main_process(): logprint(f'WARNING: Failed to load scheduler from {H.restore_scheduler_path}: {type(e).__name__}: {e}. Using fresh scheduler.') if H.restore_scaler_path: if(is_main_process()): logprint(f'Restoring scaler from {H.restore_scaler_path}') try: scaler.load_state_dict( torch.load(H.restore_scaler_path, map_location='cpu')) except Exception as e: if is_main_process(): logprint(f'WARNING: Failed to load scaler from {H.restore_scaler_path}: {type(e).__name__}: {e}. Using fresh scaler.') if H.restore_log_path: cur_eval_loss, iterate, starting_epoch = restore_log(H.restore_log_path, H.local_rank, H.mpi_size) else: cur_eval_loss, iterate, starting_epoch = float('inf'), 0, 0 logprint('starting at epoch', starting_epoch, 'iterate', iterate, 'eval loss', cur_eval_loss) return optimizer, scheduler, scaler, cur_eval_loss, iterate, starting_epoch def save_latents(H, outer, split_ind, latents, name='latents'): Path("{}/latent/".format(H.save_dir)).mkdir(parents=True, exist_ok=True) _atomic_torch_save(latents, '{}/latent/{}-{}-{}.npy'.format(H.save_dir, outer, split_ind, name)) def save_snoise(H, outer, snoise): Path("{}/latent/".format(H.save_dir)).mkdir(parents=True, exist_ok=True) for sn in snoise: _atomic_torch_save(sn, '{}/latent/snoise-{}-{}.npy'.format(H.save_dir, outer, sn.shape[2])) def save_latents_latest(H, split_ind, latents, name='latest'): Path("{}/latent/".format(H.save_dir)).mkdir(parents=True, exist_ok=True) _atomic_torch_save(latents, '{}/latent/{}-{}.npy'.format(H.save_dir, split_ind, name))