3467f8befb
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
142 lines
5.3 KiB
Python
142 lines
5.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import random
|
|
from collections import Counter
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed.checkpoint.stateful
|
|
from torch.distributed.checkpoint.state_dict import (StateDictOptions, get_model_state_dict, get_optimizer_state_dict,
|
|
set_model_state_dict, set_optimizer_state_dict)
|
|
|
|
|
|
class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
|
|
|
|
def __init__(self, model: torch.nn.Module) -> None:
|
|
self.model = model
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
state_dict = get_model_state_dict(self.model)
|
|
|
|
param_requires_grad = {
|
|
k.replace("._checkpoint_wrapped_module.", ".")
|
|
for k, v in self.model.named_parameters() if v.requires_grad
|
|
}
|
|
if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
|
|
param_requires_grad = {key.removeprefix("module.") for key in param_requires_grad}
|
|
|
|
filtered_state_dict = {k: v for k, v in state_dict.items() if k in param_requires_grad}
|
|
|
|
return filtered_state_dict
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
set_model_state_dict(
|
|
self.model,
|
|
model_state_dict=state_dict,
|
|
options=StateDictOptions(strict=False),
|
|
)
|
|
|
|
|
|
class OptimizerWrapper(torch.distributed.checkpoint.stateful.Stateful):
|
|
|
|
def __init__(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
return get_optimizer_state_dict( # type: ignore[no-any-return]
|
|
self.model,
|
|
self.optimizer,
|
|
options=StateDictOptions(flatten_optimizer_state_dict=True),
|
|
)
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
set_optimizer_state_dict(
|
|
self.model,
|
|
self.optimizer,
|
|
optim_state_dict=state_dict,
|
|
options=StateDictOptions(flatten_optimizer_state_dict=True),
|
|
)
|
|
|
|
|
|
class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful):
|
|
|
|
def __init__(self, scheduler) -> None:
|
|
self.scheduler = scheduler
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
return {"scheduler": _scheduler_state_for_dcp(self.scheduler.state_dict())}
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
self.scheduler.load_state_dict(_scheduler_state_from_dcp(state_dict["scheduler"]))
|
|
|
|
|
|
def _scheduler_state_for_dcp(value: Any, *, name: str = "") -> Any:
|
|
"""Use stable string keys for scheduler mappings flattened by DCP."""
|
|
if isinstance(value, Mapping):
|
|
if name == "milestones":
|
|
return {str(key): item for key, item in value.items()}
|
|
return {key: _scheduler_state_for_dcp(item, name=str(key)) for key, item in value.items()}
|
|
if isinstance(value, list):
|
|
return [_scheduler_state_for_dcp(item) for item in value]
|
|
return value
|
|
|
|
|
|
def _scheduler_state_from_dcp(value: Any, *, name: str = "") -> Any:
|
|
"""Restore MultiStepLR milestones after DCP stringifies mapping keys."""
|
|
if isinstance(value, Mapping):
|
|
if name == "milestones":
|
|
milestones: dict[int, Any] = {}
|
|
for key, item in value.items():
|
|
milestones[int(key)] = item
|
|
return Counter(milestones)
|
|
return {key: _scheduler_state_from_dcp(item, name=str(key)) for key, item in value.items()}
|
|
if isinstance(value, list):
|
|
return [_scheduler_state_from_dcp(item) for item in value]
|
|
return value
|
|
|
|
|
|
class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful):
|
|
|
|
def __init__(self, noise_generator: torch.Generator | None = None) -> None:
|
|
self.noise_generator = noise_generator
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
state = {
|
|
"torch_rng_state": torch.get_rng_state(),
|
|
"numpy_rng_state": np.random.get_state(),
|
|
"python_rng_state": random.getstate(),
|
|
}
|
|
|
|
if torch.cuda.is_available():
|
|
state["cuda_rng_state"] = torch.cuda.get_rng_state()
|
|
if torch.cuda.device_count() > 1:
|
|
state["cuda_rng_state_all"] = torch.cuda.get_rng_state_all()
|
|
|
|
if self.noise_generator is not None:
|
|
state["noise_generator_state"] = self.noise_generator.get_state()
|
|
|
|
return state
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
if "torch_rng_state" in state_dict:
|
|
torch.set_rng_state(state_dict["torch_rng_state"])
|
|
|
|
if "numpy_rng_state" in state_dict:
|
|
np.random.set_state(state_dict["numpy_rng_state"])
|
|
|
|
if "python_rng_state" in state_dict:
|
|
random.setstate(state_dict["python_rng_state"])
|
|
|
|
# Restore CUDA random state
|
|
if torch.cuda.is_available():
|
|
if "cuda_rng_state" in state_dict:
|
|
torch.cuda.set_rng_state(state_dict["cuda_rng_state"])
|
|
if "cuda_rng_state_all" in state_dict:
|
|
torch.cuda.set_rng_state_all(state_dict["cuda_rng_state_all"])
|
|
|
|
# Restore noise generator state
|
|
if "noise_generator_state" in state_dict and self.noise_generator is not None:
|
|
self.noise_generator.set_state(state_dict["noise_generator_state"])
|