"""Resumable Zeva deltas with data cursors; checkpoints are NEVER auto-deleted.""" import os import time import shutil import warnings from pathlib import Path import torch import torch.distributed as dist from cosmos_framework.checkpoint.dcp import DistributedCheckpointer def _is_rank0() -> bool: return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 def _full_cpu_param(param): """Collect one possibly-FSDP parameter on every rank, returning CPU data.""" # DTensor.full_tensor() is collective; every rank must execute it in the # same parameter order, even though only rank 0 retains the result. if hasattr(param, 'full_tensor'): param = param.full_tensor() return param.detach().to(device='cpu').clone() class WidowXDeltaCheckpointer(DistributedCheckpointer): def _directory(self): path = Path(os.environ['WIDOWX_RUN_DIR']) / 'checkpoints' path.mkdir(parents=True, exist_ok=True) return path def save(self, model, optimizer, scheduler, grad_scaler, iteration): world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 rank0 = _is_rank0() optimizers = getattr(optimizer, 'optimizers', [optimizer]) trainable = {id(p) for opt in optimizers for group in opt.param_groups for p in group['params']} params = {} # All ranks participate in DTensor all-gathers. Only rank 0 writes a # full delta, so the artifact is independent of the world size. for name, param in model.named_parameters(): if id(param) not in trainable: continue value = _full_cpu_param(param) if rank0: params[name] = value if rank0 and not params: raise RuntimeError('No optimizer parameters selected for checkpoint') # Convert the sharded optimizer state to a rank-0 CPU state. This is # the FSDP2-supported path and preserves FP32 FusedAdam master weights. try: from torch.distributed.checkpoint.state_dict import StateDictOptions, get_optimizer_state_dict optimizer_state = get_optimizer_state_dict( model, optimizers, options=StateDictOptions(full_state_dict=True, cpu_offload=True, ignore_frozen_params=True), ) except Exception: # Keep a useful single-rank fallback for older framework builds. if world_size > 1: raise optimizer_state = optimizer.state_dict() if rank0 else None rng_states = [torch.get_rng_state()] cuda_rng_states = [[state.cpu() for state in torch.cuda.get_rng_state_all()]] if world_size > 1: gathered_rng = [None for _ in range(world_size)] gathered_cuda_rng = [None for _ in range(world_size)] dist.all_gather_object(gathered_rng, torch.get_rng_state()) dist.all_gather_object( gathered_cuda_rng, [state.cpu() for state in torch.cuda.get_rng_state_all()], ) rng_states = gathered_rng cuda_rng_states = gathered_cuda_rng directory = self._directory() path = directory / f'step_{iteration:09d}.pt' from scripts.widowx_contract import training_contract cursor = getattr(model, '_widowx_data_cursor', None) cursors = [cursor] if world_size > 1: cursors = [None] * world_size dist.all_gather_object(cursors, cursor) payload = dict(format='zeva_widowx_delta_v3', iteration=iteration, model=params, optimizer=optimizer_state, scheduler=scheduler.state_dict(), grad_scaler=grad_scaler.state_dict(), rng=rng_states, cuda_rng=cuda_rng_states, base=os.environ['BASE_CHECKPOINT_PATH'], data_cursors=cursors, data_contract=training_contract()) if rank0: if path.exists(): raise FileExistsError(f'Refusing to overwrite an existing checkpoint: {path}') if shutil.disk_usage(directory).free < 2 * 1024**3: raise OSError('Less than 2 GiB free for checkpoint; no checkpoints were auto-deleted') # Use a unique temporary name and retry transient filesystem errors. A # fixed ``.tmp`` path can be left behind by an interrupted save and # PyTorch's zip writer then aborts the whole distributed job on the # next checkpoint. last_error = None for attempt in range(3): temporary = directory / f'{path.name}.tmp.{os.getpid()}.{attempt}' try: if temporary.exists(): temporary.unlink() torch.save(payload, temporary) if temporary.stat().st_size < 1024: raise IOError(f'checkpoint file is unexpectedly small: {temporary}') temporary.replace(path) last_error = None break except (OSError, RuntimeError) as exc: last_error = exc try: temporary.unlink() except FileNotFoundError: pass if attempt < 2: time.sleep(2 ** attempt) if last_error is not None: raise RuntimeError(f'failed to save checkpoint after 3 attempts: {path}') from last_error print(f'WIDOWX_CHECKPOINT step={iteration} path={path}', flush=True) # Retention is user-controlled. Never prune even evaluated snapshots. if world_size > 1: dist.barrier() def load(self, model, optimizer=None, scheduler=None, grad_scaler=None): iteration = super().load(model, optimizer, scheduler, grad_scaler) paths = sorted(self._directory().glob('step_*.pt')) choice = os.environ.get('WIDOWX_RESUME', 'auto') if choice == 'none': if paths: raise ValueError('Fresh training requires a new/empty checkpoint directory; existing files are protected') return iteration if choice == 'auto' and not paths: return iteration selected = paths[-1] if choice == 'auto' else Path(choice).resolve() payload = torch.load(selected, map_location='cpu', weights_only=False) if Path(payload['base']).resolve() != Path(os.environ['BASE_CHECKPOINT_PATH']).resolve(): raise ValueError('Delta checkpoint uses a different base model') from scripts.widowx_contract import training_contract from scripts import widowx_data_resume contract = training_contract() if payload.get('data_contract') != contract: if os.environ.get('WIDOWX_ALLOW_LEGACY_RESUME') != '1': raise ValueError('Checkpoint data contract is missing/different; use a fresh run, not silent resume') warnings.warn('Explicit legacy/data-change resume: optimizer is restored but sample continuity is NOT guaranteed') cursors = payload.get('data_cursors') rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 world = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 if cursors and any(c is not None for c in cursors): if len(cursors) != world: if os.environ.get('WIDOWX_ALLOW_DATA_RESHARD') != '1': raise ValueError('GPU count changed; exact data resume is unavailable. Explicitly opt into data reshuffling.') next_epoch = max(c[0] for c in cursors if c is not None) + 1 widowx_data_resume.PENDING_CURSOR = [next_epoch, 0, 0, rank, world] warnings.warn('GPU count changed: beginning a NEW shuffled data epoch, not exact sample resume') else: widowx_data_resume.PENDING_CURSOR = cursors[rank] named = dict(model.named_parameters()) with torch.no_grad(): for name, value in payload['model'].items(): p = named[name] if hasattr(p, 'device_mesh') and hasattr(p, 'placements'): # Re-shard a rank-0 full tensor into this rank's DTensor # placement. The full tensor is available on every rank # because the checkpoint is a shared local file. from torch.distributed.tensor import DTensor, Replicate full = DTensor.from_local( value.to(device=p.device), p.device_mesh, [Replicate()] * len(p.placements), run_check=False, ) local = full.redistribute(p.device_mesh, p.placements).to_local() p.to_local().copy_(local) else: p.copy_(value.to(device=p.device, dtype=p.dtype)) if optimizer is not None and payload.get('optimizer') is not None: optimizer_state = payload['optimizer'] try: from torch.distributed.checkpoint.state_dict import StateDictOptions, set_optimizer_state_dict set_optimizer_state_dict( model, getattr(optimizer, 'optimizers', [optimizer]), optimizer_state, options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True), ) except Exception: if world > 1: raise optimizer.load_state_dict(optimizer_state) if scheduler is not None: scheduler.load_state_dict(payload['scheduler']) if grad_scaler is not None: grad_scaler.load_state_dict(payload['grad_scaler']) rng = payload['rng'] cuda_rng = payload['cuda_rng'] if isinstance(rng, list) and rng and torch.is_tensor(rng[0]): rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 rng = rng[min(rank, len(rng) - 1)] cuda_rng = cuda_rng[min(rank, len(cuda_rng) - 1)] torch.set_rng_state(rng) # Each torchrun process owns one active CUDA device. Physical device # order can change on restart without mapping RNG to the wrong GPU. if torch.cuda.is_available(): device = torch.cuda.current_device() torch.cuda.set_rng_state(cuda_rng[min(device, len(cuda_rng)-1)], device=device) print(f'WIDOWX_RESUMED step={payload["iteration"]} path={selected}', flush=True) return int(payload['iteration'])