# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import collections import os import threading import traceback from typing import TYPE_CHECKING import subprocess import torch from torch.distributed.checkpoint.state_dict import ( StateDictOptions, get_optimizer_state_dict, set_model_state_dict, set_optimizer_state_dict, ) from imaginaire.model import ImaginaireModel from imaginaire.utils import callback, distributed, log, misc if TYPE_CHECKING: from imaginaire.config import CheckpointConfig, JobConfig from custom.utils.utils import get_resume, get_local_root # NOTE(bvh): this is actual used Checkpointer class class Checkpointer: """The checkpointer class. Supports checkpoint saving/loading to local disk.""" def __init__(self, config_checkpoint: CheckpointConfig, config_job: JobConfig, callbacks: callback.CallBackGroup): """Constructor of the checkpointer. Args: config_checkpoint (CheckpointConfig): The config object for the checkpointer. """ # Set the callback functions. self.callbacks = callbacks local_root = get_local_root(config_job) self.checkpoint_dir_local = f"{local_root}/{config_job.path_local}" self.strict_resume = config_checkpoint.strict_resume self.load_path = config_checkpoint.load_path or None self.load_training_state = config_checkpoint.load_training_state self.only_load_scheduler_state = config_checkpoint.only_load_scheduler_state self.keep_local = getattr(config_checkpoint, 'keep_local', False) self.allow_optimizer_reinit_on_mismatch = getattr( config_checkpoint, 'allow_optimizer_reinit_on_mismatch', False) self.save_thread = None # Prepare for resume or restart self.resume = config_checkpoint.resume if self.resume is not None: self.resume = get_resume(self.resume, local_root=local_root) self.pretrained = config_checkpoint.pretrained if self.pretrained is not None: self.pretrained = get_resume(self.pretrained, local_root=local_root) @misc.timer("cosmos_predict2/checkpointer.py:Checkpointer.save()") def save( self, model: ImaginaireModel, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, grad_scaler: torch.amp.GradScaler, iteration: int, num_samples: int, ) -> None: """Save network weights, optimizer parameters, scheduler parameters to a checkpoint (local only). S3 upload is handled by the trainer via directory sync. Args: model (ImaginaireModel): The PyTorch model. optimizer (torch.optim.Optimizer): The model optimizer. scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler. grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training). iteration (int): Current iteration number. num_samples (int): Number of samples processed so far. """ self.callbacks.on_save_checkpoint_start(model, iteration) distributed.barrier() # to avoid timeout # Also store number of samples in the filename checkpoint_file = f"iter_{iteration:09}_{num_samples:09}.pt" # Handle optimizer state dict if FSDP is enabled is_fsdp = model.config.fsdp_shard_size != 0 and distributed.get_world_size() > 1 if is_fsdp: # NOTE(bvh): a no_grad validation forward leaves FSDP2 modules that have # reshard_after_forward=False (last DiT block, root-managed params) UNSHARDED, so # model.named_parameters() returns the unsharded temporaries and no longer matches the # optimizer's sharded DTensor params by identity, which makes the optimizer state gather # below raise KeyError. Reshard everything first (no-op when already sharded); same # pattern as cosmos-predict2.5's ema_scope. from torch.distributed.fsdp import FSDPModule for m in model.modules(): if isinstance(m, FSDPModule): m.reshard() try: optimizer_state_dict = get_optimizer_state_dict( model, optimizer, options=StateDictOptions( full_state_dict=True, cpu_offload=True, ), ) except Exception as e: # NOTE(bvh): PyTorch DCP can crash (KeyError) when optimizer state is corrupted # (e.g. from NaN gradients). Fall back to saving without optimizer state. # CAUTION: this except is per-rank around a COLLECTIVE; if only a subset of ranks # raises, the others deadlock in the gather. rank0_only=False so the thrower is # visible regardless of rank (the gather is rank-symmetric since the optimizer # eagerly initializes state for every param, see fused_adam_dtensor.py). log.error(f'Failed to get FSDP optimizer state dict: {type(e).__name__}: {e}. ' f'Saving checkpoint without optimizer state.\n{traceback.format_exc()}', rank0_only=False) optimizer_state_dict = {} else: optimizer_state_dict = optimizer.state_dict() # Gather all the state dicts to be saved state_dicts_to_save = { "optim": optimizer_state_dict, "scheduler": scheduler.state_dict(), "trainer": { "grad_scaler": grad_scaler.state_dict(), "iteration": iteration, "num_samples": num_samples, }, } # Include DiT model if it's available if model.pipe.dit is not None: state_dicts_to_save["model"] = model.state_dict() # Include additional networks if there are any if len(model.pipe.extra_nets) > 0: state_dicts_to_save["extra"] = model.extra_state_dict() if distributed.get_rank() == 0: self.callbacks.on_save_checkpoint(model, state_dict=state_dicts_to_save) folders = state_dicts_to_save.keys() for folder in folders: state_dict = state_dicts_to_save[folder] with misc.timer(f"move {folder} state_dict to cpu"): if folder == 'extra': # Additional networks are saved in a dict state_dict = {key: misc.to(val, device="cpu") for key, val in state_dict.items()} else: # Other folders are saved directly state_dict = misc.to(state_dict, device="cpu") # Wait for previous saver thread to end. if self.save_thread: with misc.timer(f"join previous save_thread for {folder}"): self.save_thread.join() # Run the checkpoint saver in a separate thread. checkpoint_path = os.path.join(self.checkpoint_dir_local, folder, checkpoint_file) self.save_thread = threading.Thread( target=self._save_worker_local, daemon=False, args=(state_dict, checkpoint_path, distributed.get_rank()), ) self.save_thread.start() distributed.barrier() # to avoid timeout # Note: Checkpoints are saved on a separate thread and this callback is not accurate. # Please check logs from on_save_checkpoint_success() for better accuracy self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration) @misc.timer("cosmos_predict2/checkpointer.py:Checkpointer._save_worker_local()") def _save_worker_local(self, state_dict: dict[str, torch.Tensor], checkpoint_path: str, rank: int = 0) -> None: """Worker to save checkpoint to local disk, spawned with a child thread (runs in parallel with the training). S3 upload is handled by the trainer via directory sync after save completes. Args: state_dict (dict[str, torch.Tensor]): The state dict of the model/optimizer/scheduler. checkpoint_path (str): The path of the model checkpoint. rank (int): GPU device (default: 0). """ os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) checkpoint_file = os.path.basename(checkpoint_path) try: if rank == 0: torch.save(state_dict, checkpoint_path) self._write_latest_checkpoint_file(checkpoint_file) size_gb = os.path.getsize(checkpoint_path) / (1024 ** 3) log.success(f"Saved checkpoint (local): {checkpoint_path} ({size_gb:.2f} GiB)") iteration = int(checkpoint_file.split('_')[1]) self.callbacks.on_save_checkpoint_success(iteration=iteration) except Exception as e: # noqa: BLE001 log.exception(f"Checkpoint failed to save (local): {e}") @misc.timer("checkpoint loading") def load( self, model: ImaginaireModel, optimizer: torch.optim.Optimizer | None = None, scheduler: torch.optim.lr_scheduler.LRScheduler | None = None, grad_scaler: torch.amp.GradScaler | None = None, ): """Load network weights and optimizer states from a checkpoint in a single process. The priority of the checkpoint loading logic is: 1. Attempt to resume training if possible by looking for latest_checkpoint.txt under the same name. 2. If no latest checkpoint were found, it loads the model weights specified by config_checkpoint.path. - This is typically used for inference mode. - If config_checkpoint.load_optimizer_state is True, then also load the optimizer and scheduler states. 3. If none of the above, randomly initialize the model parameters and train from scratch. Args: model (ImaginaireModel): The PyTorch model. optimizer (torch.optim.Optimizer | None): The model optimizer (default: None). scheduler (torch.optim.lr_scheduler.LRScheduler | None): The optimization scheduler (default: None). grad_scaler (torch.amp.GradScaler | None): The gradient scaler (for mixed precision training). Returns: iteration (int): the iteration number to start/resume from. """ assert self.load_path is None, "load_path is not supported yet" self.callbacks.on_load_checkpoint_start(model) is_fsdp = model.config.fsdp_shard_size != 0 and distributed.get_world_size() > 1 if self.resume is not None: split = self.resume.split('/') checkpoint_dir, latest_checkpoint_file = '/'.join(split[:-2]), split[-1] elif self.pretrained is not None: split = self.pretrained.split('/') checkpoint_dir, latest_checkpoint_file = '/'.join(split[:-2]), split[-1] else: latest_checkpoint_file = None resuming = self.resume is not None if latest_checkpoint_file is not None: # 1. Resume training from latest_checkpoint.txt under the same name. model_checkpoint_path = os.path.join(checkpoint_dir, "model", latest_checkpoint_file) extra_checkpoint_path = os.path.join(checkpoint_dir, "extra", latest_checkpoint_file) # NOTE(bvh): missing files below are skipped silently, which on resume would mean # training from uninitialized weights (dit_path is auto-unset when resuming, see # scripts/train.py). A requested resume must find its model checkpoint. if resuming and not os.path.exists(model_checkpoint_path): raise FileNotFoundError(f'checkpoint.resume was set but the model checkpoint is ' f'missing: {model_checkpoint_path}') if resuming: optimizer_checkpoint_path = os.path.join(checkpoint_dir, "optim", latest_checkpoint_file) scheduler_checkpoint_path = os.path.join(checkpoint_dir, "scheduler", latest_checkpoint_file) trainer_checkpoint_path = os.path.join(checkpoint_dir, "trainer", latest_checkpoint_file) resume = True only_resume_scheduler = True else: optimizer_checkpoint_path = None scheduler_checkpoint_path = None trainer_checkpoint_path = None resume = False only_resume_scheduler = False else: model_checkpoint_path = None extra_checkpoint_path = None optimizer_checkpoint_path = None scheduler_checkpoint_path = None trainer_checkpoint_path = None resume = False only_resume_scheduler = False # Load checkpoint. if latest_checkpoint_file is not None: state_dicts_paths = { "model": model_checkpoint_path, "extra": extra_checkpoint_path, } if resuming: state_dicts_paths["optim"] = optimizer_checkpoint_path state_dicts_paths["scheduler"] = scheduler_checkpoint_path state_dicts_paths["trainer"] = trainer_checkpoint_path state_dicts_to_load = {} for key, checkpoint_path in state_dicts_paths.items(): # Load existing checkpoints if os.path.exists(checkpoint_path): log.info(f"Loading checkpoint (local): {checkpoint_path}") state_dicts_to_load[key] = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) log.success(f"Complete loading checkpoint (local): {checkpoint_path}") # NOTE(bvh): sattn_rope (see a4d_network.py) is dead right now so remove from loaded model if "model" in state_dicts_to_load: state_dicts_to_load["model"] = { k: v for k, v in state_dicts_to_load["model"].items() if "sattn_rope" not in k} self.callbacks.on_load_checkpoint(model, state_dict=state_dicts_to_load) # Load the state dicts. log.info("- Loading the model...") if is_fsdp: # If a model is wrapped with FSDP, its underlying weights will be DTensor. # However, Transformer Engine cannot load weights (the useless `extra_state`) into DTensor. # So we need to first remove the attention operators from Transformer Engine. # It will work correctly as long as the attention operators do not have any weights. # Cosmos-only: the Wan DiT (blocks at .wan.blocks, plain SDPA) needs no TE nulling / apply_cp. cosmos_te = model.pipe.dit is not None and hasattr(model.pipe.dit, "blocks") if model.pipe.dit is not None: if cosmos_te: for block in model.pipe.dit.blocks: block.self_attn.attn = None state_dicts_to_load_for_dit_reg = collections.OrderedDict() state_dicts_to_load_for_dit_ema = collections.OrderedDict() for key, val in state_dicts_to_load["model"].items(): if key.startswith("net."): state_dicts_to_load_for_dit_reg[key.replace("net.", "")] = val elif key.startswith("net_ema."): state_dicts_to_load_for_dit_ema[key.replace("net_ema.", "")] = val dit_strict = False if model.config.train_architecture == "lora" else True # Drop Wan alias dup keys (x_embedder/final_existing): FSDP2 dedup leaves them plain Tensors (mixed-DTensor copy error); the canonical wan.* keys load the same shared params. if not cosmos_te: alias_prefixes = ("x_embedder.", "final_existing.") for sd in (state_dicts_to_load_for_dit_reg, state_dicts_to_load_for_dit_ema): dropped = [k for k in sd if k.startswith(alias_prefixes)] for k in dropped: del sd[k] if dropped: log.info(f"Wan resume: dropped {len(dropped)} aliased dup keys " f"(loaded via canonical wan.*): {dropped}") dit_strict = False # Load Regular weights. set_model_state_dict( model.pipe.dit, state_dicts_to_load_for_dit_reg, options=StateDictOptions( full_state_dict=True, broadcast_from_rank0=True, strict=dit_strict, ), ) # Load EMA weights. if model.pipe.config.ema.enabled: set_model_state_dict( model.pipe.dit_ema, state_dicts_to_load_for_dit_ema, options=StateDictOptions( full_state_dict=True, broadcast_from_rank0=True, strict=dit_strict, ), ) # LFV: Load extra networks state_dicts_to_load_for_extra_reg = collections.OrderedDict() if "extra" in state_dicts_to_load.keys(): for extra_net_key, extra_net_val in state_dicts_to_load["extra"].items(): # Get state dict for each extra network for key, val in extra_net_val.items(): state_dicts_to_load_for_extra_reg[f'{extra_net_key}.{key}'] = val # LFV: Load extra network weights if len(state_dicts_to_load_for_extra_reg) > 0: set_model_state_dict( model.pipe.extra_nets, state_dicts_to_load_for_extra_reg, options=StateDictOptions( full_state_dict=True, broadcast_from_rank0=True, strict=dit_strict, ), ) # Restore the attention operators (cosmos TE only; no-op path for Wan). if cosmos_te: model.pipe.apply_cp() else: # Restoring DiT model if "model" in state_dicts_to_load: log.info(f"- Restoring DiT model...") model.load_state_dict(state_dicts_to_load["model"], strict=self.strict_resume) # Restoring additional networks if "extra" in state_dicts_to_load: for key, val in model.pipe.extra_nets.items(): log.info(f"- Restoring extra network {key}...") model.pipe.extra_nets[key].load_state_dict(state_dicts_to_load["extra"][key], strict=self.strict_resume) if resume or only_resume_scheduler: iteration = state_dicts_to_load["trainer"]["iteration"] num_samples = state_dicts_to_load["trainer"]["num_samples"] if "num_samples" in state_dicts_to_load["trainer"] else 0 assert scheduler log.info("- Loading the scheduler...") scheduler.load_state_dict(state_dicts_to_load["scheduler"]) scheduler.last_epoch = iteration else: iteration = 0 num_samples = 0 if resume: assert optimizer log.info("- Loading the optimizer...") try: if is_fsdp: set_optimizer_state_dict( model, optimizer, state_dicts_to_load["optim"], options=StateDictOptions( full_state_dict=True, broadcast_from_rank0=True, ), ) else: optimizer.load_state_dict(state_dicts_to_load["optim"]) except Exception as e: # noqa: BLE001 # NOTE(bvh): a pre-fix ckpt whose optimizer was saved as a 1/N FSDP shard # raises a size mismatch here. With the fix, new ckpts load fine; for OLD # ckpts, reinit the optimizer (lose only Adam momentum) instead of crashing. if not self.allow_optimizer_reinit_on_mismatch: raise log.critical( f"OPTIMIZER STATE NOT LOADED ({type(e).__name__}: {e}). Continuing with a FRESH " f"optimizer: model weights + iteration ({iteration}) + LR schedule are preserved, " f"but Adam momentum is reset (re-warms within a few hundred iters). This is expected " f"for a pre-fix checkpoint (optimizer saved as a 1/N shard).") log.info("- Loading the gradient scaler...") grad_scaler.load_state_dict(state_dicts_to_load["trainer"]["grad_scaler"]) log.success(f"Done with loading the checkpoint (iteration {iteration}).") else: log.success("Done with loading the checkpoint.") else: # Checkpoint not found and not specified. We will train everything from scratch. iteration = 0 num_samples = 0 log.info("Training from scratch.") torch.cuda.empty_cache() self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=model_checkpoint_path) return iteration, num_samples def _read_latest_checkpoint_file(self) -> str | None: """Get the file name of the latest saved checkpoint. If it doesn't exist, return None. Returns: checkpoint_file (str | None): file name of the latest saved checkpoint. """ checkpoint_file = None latest_path = os.path.join(self.checkpoint_dir_local, "latest_checkpoint.txt") if os.path.isfile(latest_path): checkpoint_file = open(latest_path).read().strip() return checkpoint_file def _write_latest_checkpoint_file(self, checkpoint_file: str) -> None: """Track the file name of the latest saved checkpoint. Args: checkpoint_file (str): file name of the latest saved checkpoint. """ content = f"{checkpoint_file}\n" latest_path = os.path.join(self.checkpoint_dir_local, "latest_checkpoint.txt") with open(latest_path, "w") as file: file.write(content) def _check_checkpoint_exists(self, checkpoint_path: str) -> None: """If the file checkpoint_path does not exist, raise an error. Args: checkpoint_path (str): full path to the checkpoint. """ if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"File not found (local): {checkpoint_path}") def finalize(self) -> None: """Finalize the checkpointer.""" if self.save_thread: self.save_thread.join()