code/helpers/train_helpers.py

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))