File size: 4,411 Bytes
b9fc164 | 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 | from pathlib import Path
import re
import torch
from torch.distributed.checkpoint.state_dict import (StateDictOptions,
get_model_state_dict,
get_optimizer_state_dict,
set_model_state_dict,
set_optimizer_state_dict)
from torch.distributed.fsdp import FSDPModule
from torchdata.stateful_dataloader import StatefulDataLoader
from src.utils.logging import get_logger
_SAVE_OPTIONS = StateDictOptions(full_state_dict=True, cpu_offload=True)
_LOAD_OPTIONS = StateDictOptions(
full_state_dict=True, broadcast_from_rank0=True)
def resolve_checkpoint_path(path: str | Path) -> Path:
path = Path(path)
if path.is_file():
return path
if not path.is_dir():
raise FileNotFoundError(f"Checkpoint does not exist: {path}")
match = re.fullmatch(r"global_step(\d+)", path.name)
if match:
checkpoint = path / f"step_{int(match.group(1))}.pth"
if checkpoint.is_file():
return checkpoint
candidates = sorted(path.glob("*.pth"))
if not candidates:
raise FileNotFoundError(f"No .pth checkpoint found in {path}.")
if len(candidates) != 1:
raise ValueError(f"Multiple checkpoints found in {path}; specify the checkpoint file.")
return candidates[0]
def load_dataloader_state(dataloader, checkpoint_path: Path, step: int) -> None:
rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
path = checkpoint_path.parent / "dataloader" / f"dataloader_step{step}_rank{rank}.pth"
if not path.is_file():
raise FileNotFoundError(f"Requested dataloader resume state is missing: {path}")
state = torch.load(path, map_location="cpu", weights_only=False)
dataloader.load_state_dict(state["dataloader"])
def save_checkpoint(
model: torch.nn.Module | FSDPModule,
save_dir: str | Path,
step: int,
global_rank: int,
optimizer: torch.optim.Optimizer = None,
scheduler: torch.optim.lr_scheduler._LRScheduler = None,
dataloader: StatefulDataLoader = None
) -> None:
logger = get_logger()
save_path = Path(save_dir) / f"global_step{step}"
save_path.mkdir(parents=True, exist_ok=True)
state_dict = {
"step": step,
"model": get_model_state_dict(model, options=_SAVE_OPTIONS)
}
if optimizer is not None:
state_dict["optimizer"] = get_optimizer_state_dict(
model, optimizer, options=_SAVE_OPTIONS)
if scheduler is not None:
state_dict["scheduler"] = scheduler.state_dict()
if global_rank <= 0:
ckpt_path = save_path / f"step_{step}.pth"
logger.info(f"Saving checkpoint at step {step} to {ckpt_path}")
torch.save(state_dict, ckpt_path)
if dataloader is not None:
save_path = save_path / 'dataloader'
save_path.mkdir(parents=True, exist_ok=True)
dataloader_path = save_path / \
f"dataloader_step{step}_rank{global_rank}.pth"
state_dict = {"dataloader": dataloader.state_dict()}
torch.save(state_dict, dataloader_path)
def load_checkpoint(
model: torch.nn.Module | FSDPModule,
path: str | Path,
optimizer: torch.optim.Optimizer = None,
scheduler: torch.optim.lr_scheduler._LRScheduler = None,
dataloader: torch.utils.data.DataLoader = None,
) -> int:
logger = get_logger()
checkpoint_path = resolve_checkpoint_path(path)
logger.info(f"Loading checkpoint from {checkpoint_path}")
state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
for key, component in (("optimizer", optimizer), ("scheduler", scheduler)):
if component is not None and key not in state_dict:
raise KeyError(f"Requested resume state {key!r} is missing from {checkpoint_path}.")
set_model_state_dict(model, state_dict["model"], options=_LOAD_OPTIONS)
if optimizer is not None:
set_optimizer_state_dict(
model, optimizer, state_dict["optimizer"], options=_LOAD_OPTIONS)
if scheduler is not None:
scheduler.load_state_dict(state_dict["scheduler"])
if dataloader is not None:
load_dataloader_state(dataloader, checkpoint_path, state_dict.get("step", 0))
return state_dict.get("step", 0)
|