File size: 10,818 Bytes
68eed61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """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'])
|