scripts/widowx_checkpoint.py

"""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'])