948 lines
43 KiB
Python
948 lines
43 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
"""
|
|
Distributed checkpoint (DCP) directory structure and storage backends.
|
|
|
|
The checkpointer saves model state in a sharded format across multiple processes:
|
|
|
|
self.save_dirname/
|
|
├── iter_000000005/ # Checkpoint at iteration 5
|
|
│ ├── model/ # Model state shards
|
|
│ │ ├── __0_0.distcp # Shard 0 from rank 0
|
|
│ │ └── __1_0.distcp # Shard 1 from rank 1
|
|
│ ├── optim/ # Optimizer state shards
|
|
│ │ ├── __0_0.distcp # Shard 0 from rank 0
|
|
│ │ └── __1_0.distcp # Shard 1 from rank 1
|
|
│ ├── scheduler/ # Learning rate scheduler state
|
|
│ │ ├── __0_0.distcp # Shard 0 from rank 0
|
|
│ │ └── __1_0.distcp # Shard 1 from rank 1
|
|
│ └── trainer/ # Additional training state
|
|
│ ├── __0_0.distcp # Shard 0 from rank 0
|
|
│ └── __1_0.distcp # Shard 1 from rank 1
|
|
└── latest_checkpoint.txt # Points to most recent checkpoint folder, e.g. iter_000000005
|
|
|
|
Storage path format:
|
|
self.save_dirname = "{config_job.path_local}/checkpoints"
|
|
|
|
The sharded format enables efficient distributed saving/loading by:
|
|
1. Parallelizing I/O across processes
|
|
2. Reducing memory usage per process
|
|
3. Supporting both local and cloud storage backends
|
|
"""
|
|
|
|
|
|
import enum
|
|
import functools
|
|
import multiprocessing
|
|
import os
|
|
import queue
|
|
import re
|
|
import time
|
|
from collections import namedtuple
|
|
from multiprocessing import get_context
|
|
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
|
import warnings
|
|
from typing import cast
|
|
|
|
import torch
|
|
import torch.distributed
|
|
import torch.distributed.checkpoint as dcp
|
|
from torch import nn
|
|
from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
|
|
from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
|
|
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.tensor import distribute_tensor
|
|
from torch.distributed.checkpoint.logger import _dcp_method_logger
|
|
from torch.distributed.checkpoint.utils import _api_bc_check, _profile, _DistWrapper
|
|
from torch.distributed.checkpoint._storage_utils import _storage_setup
|
|
from torch.distributed.checkpoint.stateful import Stateful
|
|
from torch.distributed.checkpoint.storage import StorageReader
|
|
|
|
from chronoedit._ext.imaginaire.checkpointer.base import AbstractCheckpointer
|
|
from chronoedit._ext.imaginaire.config import CheckpointConfig, JobConfig
|
|
from chronoedit._ext.imaginaire.model import ImaginaireModel
|
|
from chronoedit._ext.imaginaire.utils import callback, distributed, log, misc
|
|
from chronoedit._ext.imaginaire.utils.easy_io import easy_io
|
|
from chronoedit._src.models.wan_t2v_model import WANDiffusionModel as DiffusionModel
|
|
|
|
try:
|
|
"""
|
|
We override the default load function to skip loadding _extra_state keys created by transformer-engine.
|
|
"""
|
|
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner as _DefaultLoadPlanner
|
|
from torch.distributed.checkpoint.default_planner import (
|
|
DTensor,
|
|
LoadPlan,
|
|
_create_read_items,
|
|
_version,
|
|
flatten_state_dict,
|
|
)
|
|
from torch.distributed.checkpoint.metadata import Metadata, TensorStorageMetadata
|
|
|
|
def create_default_local_load_plan(
|
|
state_dict: dict[str, Any], metadata: Metadata, strict: bool = True, dcp_allow_mismatched_size: bool = False) -> LoadPlan:
|
|
requests = []
|
|
"""
|
|
Create the ``LoadPlan`` used by DefaultLoadPlanner.
|
|
|
|
It produces one read item per value in ``state_dict`` using the metadata in ``metadata``.
|
|
|
|
The default behavior is to match key exactly between state_dict and metadata.
|
|
It handles resharding by issuing multiple read requests against storage in order to match
|
|
load requirements.
|
|
"""
|
|
|
|
for fqn, obj in state_dict.items():
|
|
if fqn.endswith("._extra_state"): # dirty TE attention package!
|
|
continue
|
|
# ignore state_dict keys which do not exist in `state_dict` if strict=False
|
|
if fqn not in metadata.state_dict_metadata:
|
|
if strict:
|
|
raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
|
|
else:
|
|
log.warning(f"Local load plan: Missing key in checkpoint state_dict: {fqn}.")
|
|
continue
|
|
|
|
md = metadata.state_dict_metadata[fqn]
|
|
|
|
if not dcp_allow_mismatched_size:
|
|
if (
|
|
isinstance(md, TensorStorageMetadata)
|
|
and getattr(obj, "size", None) is not None
|
|
and md.size != obj.size()
|
|
):
|
|
if not strict:
|
|
log.critical(f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}")
|
|
continue
|
|
else:
|
|
raise ValueError(
|
|
f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}",
|
|
)
|
|
# Since DTensor supports submesh, adding extra check to ensure _create_read_items()
|
|
# gets called only when the current rank is part of the mesh for the corresponding DTensor.
|
|
if isinstance(obj, DTensor):
|
|
if obj.device_mesh.get_coordinate() is not None:
|
|
requests += _create_read_items(fqn, md, obj)
|
|
else:
|
|
requests += _create_read_items(fqn, md, obj)
|
|
|
|
return LoadPlan(requests)
|
|
|
|
class DefaultLoadPlanner(_DefaultLoadPlanner):
|
|
def set_partial_channel_weight(self, dcp_allow_mismatched_size: bool):
|
|
self.dcp_allow_mismatched_size = dcp_allow_mismatched_size
|
|
|
|
def create_local_plan(self) -> LoadPlan:
|
|
assert self.metadata is not None
|
|
if self.flatten_state_dict:
|
|
# To support checkpoints that are saved before v2.4, we have to
|
|
# differentiate if the missing keys are due to old checkpoints.
|
|
# The contracts are:
|
|
# 1. There are 3 cases when we found a missing key.
|
|
# 1.1 Actual missing key, but allow_partial_load is False
|
|
# 1.2 Actual missing key, but allow_partial load is True
|
|
# 1.3 Old checkpoint, but allow_partial_load is False
|
|
# 1.4 Old checkpoint, but allow_partial_load is True
|
|
# 2. If we found a missing key, we first convert the keys back to
|
|
# the key format of v2.3
|
|
# 3. If the previous missing keys are in the v2.3 keys, we assume
|
|
# this is a old checkpoint.
|
|
# 4. Pass the state_dict to `create_default_local_load_plan()`,
|
|
# which has the logic to check missing for allow_partial_load.
|
|
# So for 1.2 and 1.4 cases, we delegate allow_partial_load check to
|
|
# `create_default_local_load_plan()`. The logic here is to determine
|
|
# whether the checkpoint belong to 2.3 (or before) or 2.4 (or after).
|
|
current_keys = set(self.state_dict.keys())
|
|
load_keys = set(self.metadata.state_dict_metadata.keys())
|
|
missing_keys = load_keys - current_keys
|
|
if missing_keys:
|
|
_version._derived_version = "2_3"
|
|
old_state_dict, old_mappings = flatten_state_dict(self.original_state_dict)
|
|
old_keys = set(old_state_dict.keys())
|
|
if old_keys & missing_keys:
|
|
self.state_dict, self.mappings = old_state_dict, old_mappings
|
|
# _derived_version is only used by flatten_state_dict now.
|
|
# Set it back to None so that later we can save to a new version.
|
|
_version._derived_version = None
|
|
|
|
return create_default_local_load_plan(
|
|
self.state_dict,
|
|
self.metadata,
|
|
not self.allow_partial_load,
|
|
getattr(self, "dcp_allow_mismatched_size", False),
|
|
)
|
|
|
|
log.critical("for the back comptiable pytorch! New DefaultLoadPlanner class is created.")
|
|
|
|
@_dcp_method_logger(log_exceptions=True)
|
|
@_api_bc_check
|
|
def load(
|
|
state_dict: dict[str, Any],
|
|
*,
|
|
checkpoint_id: Union[str, os.PathLike, None] = None,
|
|
storage_reader: Optional[StorageReader] = None,
|
|
planner: Optional[DefaultLoadPlanner] = None,
|
|
process_group: Optional[torch.distributed.ProcessGroup] = None,
|
|
no_dist: bool = False,
|
|
) -> None:
|
|
"""
|
|
We override the default load function to perform a load plan check for mismatched/missing keys.
|
|
==========================Original Doc string=====================================
|
|
Load a checkpoint into a distributed state dict in SPMD style.
|
|
|
|
Each rank must have the same keys in their ``state_dict`` provided to this
|
|
API. Mismatched keys may result in hangs or errors. If unsure, you can use
|
|
the ``utils._assert_same_keys`` API to check (but may incur communication
|
|
costs).
|
|
|
|
Each rank will try to read the least amount of data necessary
|
|
to fullfill the requested `state_dict`. When loading :class:`ShardedTensor`
|
|
or :class:`DTensor` instances, each rank only reads data for their local shards.
|
|
|
|
For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``),
|
|
load will first call ``state_dict`` before attempting deserialization, followed by
|
|
``load_state_dict`` once the deserialization is complete.
|
|
For each non-``Stateful`` object, load will deserailize the object, and then replace
|
|
it in the ``state_dict`` with the deserialized object.
|
|
|
|
.. warning::
|
|
All tensors in ``state_dict`` must be allocated on their
|
|
destination device *prior to* calling this function.
|
|
|
|
All non-tensor data is loaded using `torch.load()` and modified in place
|
|
on state_dict.
|
|
|
|
.. warning::
|
|
Users must call `load_state_dict` on the root module to ensure load
|
|
pos-processing and non-tensor data properly propagates.
|
|
|
|
.. note:
|
|
If no process group is initialized, this function will assume the intent
|
|
is to load a checkpoint into the local process. This can be useful in the
|
|
case of local inference, and when using regular Tensors (as opposed to DTensor
|
|
or ShardedTensor)
|
|
|
|
.. note:
|
|
Rank 0 is assumed to be the coordinator rank.
|
|
|
|
Args:
|
|
state_dict (Dict[str, Any]): The state_dict to load the checkpoint into.
|
|
checkpoint_id (Union[str, os.PathLike, None]):
|
|
The ID of this checkpoint instance. The meaning of the checkpoint_id
|
|
depends on the storage. It can be a path to a folder or to a file.
|
|
It can also be a key if the storage is a key-value store.
|
|
(Default: ``None``)
|
|
storage_reader (Optional[StorageReader]):
|
|
Instance of StorageWriter used to perform reads. If this is not
|
|
specified, DCP will automatically infer the reader based on the
|
|
checkpoint_id. If checkpoint_id is also None, an exception will
|
|
be raised. (Default: ``None``)
|
|
planner (Optional[LoadPlanner]):
|
|
Instance of LoadPlanner. If this is not specificed, the default
|
|
planner will be used. (Default: ``None``)
|
|
process_group (Optional[ProcessGroup]):
|
|
ProcessGroup to be used for cross-rank synchronization.
|
|
(Default: ``None``)
|
|
no_dist (bool): If ``True``, this function will assume the intent is to load
|
|
a checkpoint without using cross-rank synchronization. (Default: ``False``)
|
|
Returns:
|
|
None.
|
|
|
|
Examples
|
|
>>> # xdoctest: +SKIP
|
|
>>> my_model = MyModule()
|
|
>>> optimizer = Adagrad(my_model.parameters())
|
|
>>> model_state_dict = my_model.state_dict()
|
|
>>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
|
|
... "/checkpoint/1"
|
|
... )
|
|
|
|
>>> torch.distributed.checkpoint.load_state_dict(
|
|
>>> state_dict=model_state_dict,
|
|
>>> storage_reader=fs_storage_reader,
|
|
>>> )
|
|
|
|
>>> # module.load_state_dict() function might have customized steps
|
|
>>> # to flush the state_dict, must call it to
|
|
>>> # ensure correct behavior.
|
|
>>> my_model.load_state_dict(model_state_dict)
|
|
|
|
.. note::
|
|
load_state_dict uses collectives to coordinate reads across ranks.
|
|
For NCCL-based process groups, internal tensor representations of
|
|
objects must be moved to the GPU device before communication takes place.
|
|
In this case, the device used is given by ``torch.cuda.current_device()``
|
|
and it is the user's responsibility to ensure that this is set so that each
|
|
rank has an individual GPU, via ``torch.cuda.set_device()``.
|
|
"""
|
|
|
|
no_dist = no_dist or (not torch.distributed.is_available()) or (not torch.distributed.is_initialized())
|
|
if no_dist:
|
|
warnings.warn(
|
|
"torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to load in a single process."
|
|
)
|
|
|
|
with _profile():
|
|
storage_reader = cast(
|
|
StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True)
|
|
)
|
|
|
|
# All ranks must have the same keys in their `state_dict` provided to
|
|
# this API. See documentation for more details.
|
|
# Here we simply sort the keys to ensure that all ranks load values in
|
|
# the same order.
|
|
keys = sorted(state_dict.keys())
|
|
|
|
statetful_sd = {}
|
|
for key in keys:
|
|
if key not in state_dict:
|
|
continue
|
|
elem = state_dict[key]
|
|
statetful_sd[key] = (
|
|
elem.state_dict() if isinstance(elem, Stateful) else elem
|
|
)
|
|
|
|
_load_state_dict(
|
|
state_dict=statetful_sd,
|
|
storage_reader=storage_reader,
|
|
process_group=process_group,
|
|
no_dist=no_dist,
|
|
planner=planner,
|
|
)
|
|
for key in keys:
|
|
if key not in state_dict:
|
|
continue
|
|
elem = state_dict[key]
|
|
if isinstance(elem, Stateful):
|
|
# If the state_dict is a Stateful object,
|
|
# DCP does an in-place load in the original state dict.
|
|
elem.load_state_dict(statetful_sd[key])
|
|
else:
|
|
# Otherwise, replace the state_dict with the loaded state_dict.
|
|
state_dict[key] = statetful_sd[key]
|
|
|
|
def _load_state_dict(
|
|
state_dict: dict[str, Any],
|
|
storage_reader: StorageReader,
|
|
process_group: Optional[torch.distributed.ProcessGroup] = None,
|
|
coordinator_rank: int = 0,
|
|
no_dist: bool = False,
|
|
planner: Optional[DefaultLoadPlanner] = None,
|
|
) -> None:
|
|
torch._C._log_api_usage_once("torch.distributed.checkpoint.load_state_dict")
|
|
|
|
distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
|
|
if planner is None:
|
|
planner = DefaultLoadPlanner()
|
|
|
|
ckpt_kwargs = {}
|
|
if (ckpt_id := getattr(storage_reader, "checkpoint_id", None)) is not None:
|
|
ckpt_kwargs["checkpoint_id"] = ckpt_id
|
|
ckpt_kwargs["process_group"] = distW.group
|
|
|
|
@_dcp_method_logger(**ckpt_kwargs)
|
|
def local_step():
|
|
assert planner is not None
|
|
metadata = storage_reader.read_metadata()
|
|
planner.set_up_planner(state_dict, metadata, distW.is_coordinator)
|
|
storage_reader.set_up_storage_reader(metadata, distW.is_coordinator)
|
|
|
|
local_plan = planner.create_local_plan()
|
|
local_plan = storage_reader.prepare_local_plan(local_plan)
|
|
return local_plan
|
|
|
|
@_dcp_method_logger(**ckpt_kwargs)
|
|
def global_step(all_local_plans):
|
|
assert planner is not None
|
|
all_local_plans = planner.create_global_plan(all_local_plans)
|
|
all_local_plans = storage_reader.prepare_global_plan(all_local_plans)
|
|
return all_local_plans
|
|
|
|
central_plan: LoadPlan = distW.reduce_scatter("plan", local_step, global_step)
|
|
if distW.is_coordinator:
|
|
# Compare central_plan items with storage_reader.storage_data keys
|
|
dest_fqns = set()
|
|
storage_fqns = set()
|
|
|
|
# Extract FQNs from central_plan items
|
|
for item in central_plan.items:
|
|
if hasattr(item, 'dest_index') and hasattr(item.dest_index, 'fqn'):
|
|
dest_fqns.add(item.dest_index.fqn)
|
|
if hasattr(item, 'storage_index') and hasattr(item.storage_index, 'fqn'):
|
|
storage_fqns.add(item.storage_index.fqn)
|
|
|
|
# Get storage data keys
|
|
storage_data_keys = set()
|
|
if hasattr(storage_reader, 'storage_data') and storage_reader.storage_data is not None:
|
|
storage_data_keys = set(item[0].fqn for item in storage_reader.storage_data.items())
|
|
state_dict_keys = set(state_dict.keys())
|
|
# Compare sets and log differences
|
|
# Remove any item that has "_extra_state" as substring in the sets
|
|
state_dict_keys = {fqn for fqn in state_dict_keys if "_extra_state" not in fqn}
|
|
dest_fqns = {fqn for fqn in dest_fqns if "_extra_state" not in fqn}
|
|
storage_fqns = {fqn for fqn in storage_fqns if "_extra_state" not in fqn}
|
|
storage_data_keys = {fqn for fqn in storage_data_keys if "_extra_state" not in fqn}
|
|
|
|
log.info("=== Load Plan FQN Analysis ===")
|
|
log.info(f"State Dict FQNs count: {len(state_dict_keys)}")
|
|
log.info(f"Destination FQNs count (without _extra_state): {len(dest_fqns)}")
|
|
log.info(f"Loaded FQNs count (without _extra_state): {len(storage_fqns)}")
|
|
log.info(f"In Storage keys count (without _extra_state): {len(storage_data_keys)}")
|
|
|
|
# Find missing keys in each direction
|
|
state_dict_missing_from_dest = state_dict_keys - dest_fqns
|
|
storage_data_missing_from_storage_fqns = storage_data_keys - storage_fqns
|
|
|
|
if state_dict_missing_from_dest:
|
|
log.info(f"State Dict FQNs missing from load plan ({len(state_dict_missing_from_dest)} items): {sorted(state_dict_missing_from_dest)}")
|
|
else:
|
|
log.info("✓ All State Dict FQNs found in storage_data")
|
|
|
|
if storage_data_missing_from_storage_fqns:
|
|
# If there are more than 100 "net_ema" keys in storage_data_missing_from_storage_fqns, summarize them
|
|
net_ema_keys = {k for k in storage_data_missing_from_storage_fqns if "net_ema" in k}
|
|
if len(net_ema_keys) > 100:
|
|
storage_data_missing_from_storage_fqns = storage_data_missing_from_storage_fqns - net_ema_keys
|
|
storage_data_missing_from_storage_fqns = set(storage_data_missing_from_storage_fqns) # ensure set type
|
|
storage_data_missing_from_storage_fqns.add("net_ema")
|
|
log.info(f"Summarized {len(net_ema_keys)} 'net_ema' keys as 'net_ema' in missing storage data keys.")
|
|
log.info(f"Storage data keys not loaded by load plan ({len(storage_data_missing_from_storage_fqns)} items): {sorted(storage_data_missing_from_storage_fqns)}")
|
|
else:
|
|
log.info("✓ All storage data keys found in Loaded FQNs")
|
|
|
|
log.info("=== End Load Plan FQN Analysis ===")
|
|
|
|
@_dcp_method_logger(**ckpt_kwargs)
|
|
def read_data():
|
|
assert planner is not None
|
|
final_local_plan = planner.finish_plan(central_plan)
|
|
all_reads = storage_reader.read_data(final_local_plan, planner)
|
|
|
|
all_reads.wait()
|
|
return None
|
|
|
|
_ = distW.all_gather("read", read_data)
|
|
|
|
except ImportError as e:
|
|
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
|
|
|
|
log.critical(f"{e}, using default planner")
|
|
|
|
StateDictItemPath = namedtuple("StateDictItemPath", ["state_dict", "save_path"])
|
|
|
|
class ModelWrapper(Stateful):
|
|
"""Wrapper for model state dict handling"""
|
|
|
|
def __init__(self, model: Union[nn.Module, List[nn.Module]], load_ema_to_reg: bool = False):
|
|
self.model = [model] if isinstance(model, nn.Module) else model
|
|
self.load_ema_to_reg = load_ema_to_reg
|
|
if self.load_ema_to_reg:
|
|
assert isinstance(model, DiffusionModel)
|
|
|
|
def state_dict(self) -> Dict[str, Any]:
|
|
_state_dict = {k: v for sd in map(get_model_state_dict, self.model) for k, v in sd.items()}
|
|
if self.load_ema_to_reg:
|
|
if not self.model[0].config.ema.enabled:
|
|
all_keys = list(_state_dict.keys())
|
|
assert all(k.startswith("net.") for k in all_keys), "All keys must start with net."
|
|
for k in all_keys:
|
|
_state_dict[k.replace("net.", "net_ema.")] = _state_dict.pop(k)
|
|
else:
|
|
log.warning("EMA is enabled, will only load EMA weights from checkpoint file.")
|
|
all_keys = list(_state_dict.keys())
|
|
for k in all_keys:
|
|
if k.startswith("net_ema."):
|
|
break
|
|
else:
|
|
raise ValueError("No EMA keys found in state_dict")
|
|
# do not load .net keys, since we do not need them anyway.
|
|
_state_dict = {k: _state_dict[k] for k in all_keys if not k.startswith("net.")}
|
|
return _state_dict
|
|
|
|
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
|
|
if self.load_ema_to_reg:
|
|
if not self.model[0].config.ema.enabled:
|
|
all_keys = list(state_dict.keys())
|
|
assert all(k.startswith("net_ema.") for k in all_keys), "All keys must start with net_ema."
|
|
for k in all_keys:
|
|
state_dict[k.replace("net_ema.", "net.")] = state_dict.pop(k)
|
|
else:
|
|
log.warning("EMA is enabled, will load EMA weights to regular model weights")
|
|
all_keys = list(state_dict.keys())
|
|
assert all(not k.startswith("net.") for k in all_keys), "No .net keys should be in state_dict"
|
|
for k in all_keys:
|
|
if k.startswith("net_ema."):
|
|
state_dict[k.replace("net_ema.", "net.")] = torch.clone(state_dict[k])
|
|
func = functools.partial(
|
|
set_model_state_dict,
|
|
model_state_dict=state_dict,
|
|
options=StateDictOptions(strict=False),
|
|
)
|
|
list(map(func, self.model))
|
|
|
|
|
|
class OptimizerWrapper(Stateful):
|
|
def __init__(
|
|
self,
|
|
model: Union[nn.Module, List[nn.Module]],
|
|
optim: Union[torch.optim.Optimizer, List[torch.optim.Optimizer]],
|
|
) -> None:
|
|
self.model = [model] if isinstance(model, nn.Module) else model
|
|
self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim
|
|
|
|
def state_dict(self) -> Dict[str, Any]:
|
|
func = functools.partial(
|
|
get_optimizer_state_dict,
|
|
options=StateDictOptions(flatten_optimizer_state_dict=True),
|
|
)
|
|
return {k: v for sd in map(func, self.model, self.optim) for k, v in sd.items()}
|
|
|
|
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
|
|
func = functools.partial(
|
|
set_optimizer_state_dict,
|
|
optim_state_dict=state_dict,
|
|
options=StateDictOptions(flatten_optimizer_state_dict=True),
|
|
)
|
|
list(map(func, self.model, self.optim))
|
|
|
|
|
|
class AsyncMode(str, enum.Enum):
|
|
DISABLED = "disabled"
|
|
ASYNC_WITH_PINNED_MEM = "async_with_pinned_mem"
|
|
|
|
|
|
class Terminate:
|
|
pass
|
|
|
|
|
|
class SaveDone:
|
|
def __init__(self, iteration: int, elapsed_time: float, succeeded: bool):
|
|
self.iteration = iteration
|
|
self.elapsed_time = elapsed_time
|
|
self.succeeded = succeeded
|
|
|
|
def __str__(self):
|
|
return f"SaveDone(iteration={self.iteration}, elapsed_time={self.elapsed_time}, succeeded={self.succeeded})"
|
|
|
|
|
|
def save_checkpoint_in_background(
|
|
receiver_queue: multiprocessing.Queue,
|
|
sender_queue: multiprocessing.Queue,
|
|
checkpoint_config: CheckpointConfig,
|
|
job_config: JobConfig,
|
|
) -> None:
|
|
"""
|
|
Handles model checkpoint saving in a separate background process using PyTorch's distributed functionality.
|
|
This function runs in a dedicated process to avoid blocking the main training loop.
|
|
|
|
Args:
|
|
receiver_queue: Queue to receive state dictionaries and commands from the main process
|
|
sender_queue: Queue to send completion signals back to the main process
|
|
checkpoint_config: Configuration settings for checkpoint saving behavior
|
|
job_config: Configuration settings for the training job
|
|
|
|
Flow:
|
|
1. Initializes distributed processing environment
|
|
2. Continuously waits for state dictionaries to save
|
|
3. Saves checkpoints asynchronously
|
|
4. Signals completion back to main process
|
|
5. Terminates when receiving a Terminate signal
|
|
|
|
Raises:
|
|
AssertionError: If received object is neither Terminate signal nor valid state dict tuple
|
|
|
|
Note:
|
|
- Uses a different port than the main process to avoid conflicts
|
|
- Disables TorchElastic agent store for checkpoint operations
|
|
- Automatically cleans up distributed process group on exit
|
|
"""
|
|
# Configure distributed environment
|
|
os.environ["MASTER_PORT"] = str(int(os.environ["MASTER_PORT"]) + 2)
|
|
os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False"
|
|
|
|
# Set up GPU device and distributed processing
|
|
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
|
|
distributed.init()
|
|
|
|
# Initialize checkpointing mechanism
|
|
checkpoint_handler = DistributedCheckpointer(checkpoint_config, job_config, None, disable_async=True)
|
|
|
|
try:
|
|
while True:
|
|
log.debug("Checkpoint background process is ready for next task, waiting for new state_dict")
|
|
received_data = receiver_queue.get()
|
|
log.debug("Received new state_dict")
|
|
|
|
if isinstance(received_data, Terminate):
|
|
log.info("Received termination signal in checkpoint background process, closing sender queue")
|
|
sender_queue.put(Terminate())
|
|
sender_queue.close()
|
|
return
|
|
|
|
assert isinstance(received_data, tuple), "Received data must be a tuple of (state_dict, checkpoint_path)"
|
|
state_dict, checkpoint_path = received_data
|
|
|
|
# Save checkpoint and measure time taken
|
|
start_time = time.monotonic()
|
|
iteration = state_dict["trainer"][0]["iteration"]
|
|
elapsed_time = 0
|
|
succeeded = False
|
|
try:
|
|
checkpoint_handler.save_state_dict_worker(state_dict, checkpoint_path)
|
|
elapsed_time = time.monotonic() - start_time
|
|
log.info(
|
|
f"Checkpoint saved successfully in background process. Time taken: {elapsed_time:.2f} seconds, iteration: {iteration}"
|
|
)
|
|
succeeded = True
|
|
except Exception as e:
|
|
log.error(f"Error saving checkpoint to {checkpoint_path}: {e}")
|
|
# continue because if the thread exits, the main thread keeps on adding to the queue
|
|
finally:
|
|
if elapsed_time == 0:
|
|
elapsed_time = time.monotonic() - start_time
|
|
sender_queue.put(SaveDone(iteration, elapsed_time, succeeded))
|
|
|
|
finally:
|
|
log.info("Cleaning up: destroying distributed process group")
|
|
torch.distributed.destroy_process_group()
|
|
|
|
|
|
class DistributedCheckpointer(AbstractCheckpointer):
|
|
KEYS_TO_SAVE = ["model", "optim", "scheduler", "trainer"]
|
|
|
|
def __init__(
|
|
self,
|
|
config_checkpoint: CheckpointConfig,
|
|
config_job: JobConfig,
|
|
callbacks: Optional[callback.CallBackGroup] = None,
|
|
disable_async: bool = False,
|
|
):
|
|
super().__init__(config_checkpoint, config_job, callbacks)
|
|
self.config_checkpoint = config_checkpoint
|
|
if config_checkpoint.dcp_async_mode_enabled:
|
|
self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
|
|
else:
|
|
self.async_mode = AsyncMode.DISABLED
|
|
|
|
if disable_async:
|
|
self.async_mode = AsyncMode.DISABLED
|
|
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
|
|
ctx = get_context("spawn")
|
|
self.mp_queue_send = ctx.Queue()
|
|
self.mp_queue_recv = ctx.Queue()
|
|
self.mp = ctx.Process(
|
|
target=save_checkpoint_in_background,
|
|
args=(
|
|
self.mp_queue_send,
|
|
self.mp_queue_recv,
|
|
config_checkpoint,
|
|
config_job,
|
|
),
|
|
daemon=True,
|
|
)
|
|
self.mp.start()
|
|
self.cpu_offload_state_dict = None
|
|
self.staging = False
|
|
self.staging_ckpt_file = None
|
|
self.staging_stream = torch.cuda.Stream()
|
|
|
|
def keys_to_resume_during_load(self) -> Tuple[Set, Union[str, None]]:
|
|
latest_checkpoint_file = self._read_latest_checkpoint_file()
|
|
|
|
resume_keys = []
|
|
|
|
if latest_checkpoint_file is not None:
|
|
# 1. Resume training from latest_checkpoint.txt under the same name.
|
|
checkpoint_path = os.path.join(self.load_dirname, latest_checkpoint_file)
|
|
resume_keys.extend(self.KEYS_TO_SAVE)
|
|
else:
|
|
if self.load_path and not self.load_path.endswith(".pth"):
|
|
# 2. Load the module weights specified by config_checkpoint.path.
|
|
checkpoint_path = self.load_path
|
|
if self.load_training_state:
|
|
resume_keys.extend(self.KEYS_TO_SAVE)
|
|
else:
|
|
resume_keys.append("model")
|
|
if self.only_load_scheduler_state:
|
|
resume_keys.append("scheduler")
|
|
elif self.load_path and self.load_path.endswith(".pth"):
|
|
checkpoint_path = self.load_path
|
|
else:
|
|
checkpoint_path = None
|
|
if len(self.keys_not_to_resume) > 0:
|
|
for key in self.keys_not_to_resume:
|
|
assert key in self.KEYS_TO_SAVE, f"Invalid key to resume: {key} not in {self.KEYS_TO_SAVE}"
|
|
resume_keys = [key for key in resume_keys if key not in self.keys_not_to_resume]
|
|
return set(resume_keys), checkpoint_path
|
|
|
|
@misc.timer("checkpoint loading")
|
|
def load(
|
|
self,
|
|
model: ImaginaireModel,
|
|
optimizer: torch.optim.Optimizer | None = None,
|
|
scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
|
|
grad_scaler: torch.amp.GradScaler | None = None,
|
|
) -> int:
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_load_checkpoint_start(model)
|
|
|
|
resume_keys, checkpoint_path = self.keys_to_resume_during_load()
|
|
resume_keys = sorted(resume_keys)
|
|
log.critical(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}")
|
|
|
|
iteration = 0
|
|
|
|
if checkpoint_path is not None and not checkpoint_path.endswith(".pth"):
|
|
self._check_checkpoint_exists(checkpoint_path)
|
|
for key in resume_keys:
|
|
load_planner = DefaultLoadPlanner(allow_partial_load=True)
|
|
if hasattr(load_planner, "set_partial_channel_weight"):
|
|
log.critical(f"set_partial_channel_weight: {self.config_checkpoint.dcp_allow_mismatched_size}")
|
|
load_planner.set_partial_channel_weight(self.config_checkpoint.dcp_allow_mismatched_size)
|
|
cur_key_ckpt_full_path = os.path.join(checkpoint_path, key)
|
|
log.critical(f"Start loading checkpoint from {checkpoint_path}")
|
|
storage_reader = self.get_storage_reader(cur_key_ckpt_full_path)
|
|
torch.distributed.barrier()
|
|
log.critical(f"starting {cur_key_ckpt_full_path}", rank0_only=False)
|
|
if key == "model":
|
|
log.info("- Loading the model...")
|
|
_model_wrapper = ModelWrapper(model)
|
|
_state_dict = _model_wrapper.state_dict()
|
|
load(
|
|
_state_dict,
|
|
storage_reader=storage_reader,
|
|
planner=load_planner
|
|
)
|
|
_model_wrapper.load_state_dict(_state_dict)
|
|
elif key == "optim":
|
|
log.info("- Loading the optimizer...")
|
|
_optim_wrapper = OptimizerWrapper(model, optimizer)
|
|
_state_dict = _optim_wrapper.state_dict()
|
|
dcp.load(
|
|
_state_dict,
|
|
storage_reader=storage_reader,
|
|
planner=load_planner,
|
|
)
|
|
_optim_wrapper.load_state_dict(_state_dict)
|
|
elif key == "scheduler":
|
|
log.info("- Loading the scheduler...")
|
|
_state_dict = scheduler.state_dict()
|
|
dcp.load(
|
|
_state_dict,
|
|
storage_reader=storage_reader,
|
|
planner=load_planner,
|
|
)
|
|
scheduler.load_state_dict(_state_dict)
|
|
elif key == "trainer":
|
|
log.info("- Loading the trainer...")
|
|
_state_dict = {
|
|
"grad_scaler": grad_scaler.state_dict(),
|
|
"iteration": iteration,
|
|
}
|
|
dcp.load(
|
|
_state_dict,
|
|
storage_reader=storage_reader,
|
|
planner=load_planner,
|
|
)
|
|
grad_scaler.load_state_dict(_state_dict["grad_scaler"])
|
|
iteration = _state_dict["iteration"]
|
|
else:
|
|
raise ValueError(f"Invalid key: {key}. not support to resume.")
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_load_checkpoint(model, state_dict=_state_dict)
|
|
log.critical(f"Loaded checkpoint from {checkpoint_path} in iteration {iteration}")
|
|
elif checkpoint_path is not None and checkpoint_path.endswith(".pth"):
|
|
|
|
state = easy_io.load(checkpoint_path)
|
|
model_state = model.net.state_dict()
|
|
|
|
for k, v in list(state.items()):
|
|
|
|
tgt = model_state.get(k, None)
|
|
if tgt is None:
|
|
continue
|
|
# If target param/buffer is a DTensor and checkpoint value is not, distribute it
|
|
if isinstance(tgt, DTensor) and not isinstance(v, DTensor):
|
|
# Match device, dtype, and placements from the target DTensor
|
|
v = v.to(tgt.device, dtype=tgt.dtype, copy=False)
|
|
v = distribute_tensor(v, tgt.device_mesh, tgt.placements)
|
|
state[k] = v
|
|
# If target is a plain Tensor but checkpoint is a DTensor, bring it local
|
|
if not isinstance(tgt, DTensor) and isinstance(v, DTensor):
|
|
state[k] = v.to_local().to(tgt.device, dtype=tgt.dtype, copy=False)
|
|
|
|
model.load_state_dict(state, strict=False, pretrain_copy=True)
|
|
# Clear unused reserved memory from fp32
|
|
torch.cuda.empty_cache()
|
|
log.critical(f"Loaded checkpoint from {checkpoint_path}. **** This only happen at iteration 0 **** ")
|
|
else:
|
|
log.info("Training from scratch.")
|
|
torch.cuda.empty_cache()
|
|
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path)
|
|
return iteration
|
|
|
|
def _async_with_pinned_memory(self, checkpoint_file: str, state_dict: Dict[str, Tuple[Any, str]]) -> None:
|
|
try:
|
|
from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Please install the latest PyTorch nightly to use async checkpointing with pinned memory."
|
|
) from e
|
|
if self.cpu_offload_state_dict is None:
|
|
log.debug(f"Preparing the CPU memory, {time.monotonic()=}.:.2f")
|
|
self.cpu_offload_state_dict = _create_cpu_state_dict(state_dict, pin_memory=True, share_memory=True)
|
|
|
|
log.debug(f"Staging the state_dict, {time.monotonic()=}.:.2f")
|
|
with torch.cuda.stream(self.staging_stream):
|
|
self.cpu_offload_state_dict = _copy_state_dict(
|
|
state_dict,
|
|
self.cpu_offload_state_dict,
|
|
non_blocking=True,
|
|
)
|
|
self.staging = True
|
|
self.staging_ckpt_file = checkpoint_file
|
|
|
|
self.maybe_wait_for_staging()
|
|
|
|
def maybe_wait_for_staging(self) -> None:
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM and self.staging:
|
|
if not self.staging_stream.query():
|
|
self.staging_stream.synchronize()
|
|
|
|
def sync_func():
|
|
self.mp_queue_send.put_nowait((self.cpu_offload_state_dict, self.staging_ckpt_file))
|
|
|
|
sync_func()
|
|
self.staging = False
|
|
|
|
def get_previous_checkpoint_results(self, wait_for: int = 0) -> None:
|
|
"""Get the results of previously submitted checkpoints and pass them to callbacks if checkpoint succeeded"""
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
|
|
try:
|
|
start_time = time.monotonic()
|
|
while not self.mp_queue_recv.empty() or wait_for > 0:
|
|
try:
|
|
ret = self.mp_queue_recv.get(timeout=1)
|
|
if isinstance(ret, Terminate):
|
|
log.info("Received termination event from checkpoint background process")
|
|
break
|
|
save_done: SaveDone = ret
|
|
log.logger.info(f"Received checkpoint save result: {save_done}")
|
|
if self.callbacks is not None and save_done.succeeded:
|
|
self.callbacks.on_save_checkpoint_success(
|
|
iteration=save_done.iteration, elapsed_time=save_done.elapsed_time
|
|
)
|
|
except queue.Empty:
|
|
elapsed_time = time.monotonic() - start_time
|
|
if elapsed_time > wait_for:
|
|
break
|
|
except (EOFError, BrokenPipeError):
|
|
log.info("Queue was closed by checkpoint background process")
|
|
|
|
def get_storage_writer(self, checkpoint_path: str) -> FileSystemWriter:
|
|
return FileSystemWriter(path=checkpoint_path)
|
|
|
|
def get_storage_reader(self, checkpoint_path: str) -> FileSystemReader:
|
|
return FileSystemReader(checkpoint_path)
|
|
|
|
def save_state_dict_worker(self, to_save_dict: Dict[str, Tuple[Any, str]], checkpoint_file: str) -> None:
|
|
for k, (v, full_checkpoint_path) in to_save_dict.items():
|
|
storage_writer = self.get_storage_writer(full_checkpoint_path)
|
|
dcp.save(
|
|
v,
|
|
storage_writer=storage_writer,
|
|
planner=DefaultSavePlanner(dedup_save_to_lowest_rank=True),
|
|
)
|
|
|
|
self._write_latest_checkpoint_file(checkpoint_file)
|
|
log.critical(f"Saved checkpoint to {os.path.join(self.save_dirname, checkpoint_file)}", rank0_only=True)
|
|
|
|
def save(
|
|
self,
|
|
model: ImaginaireModel,
|
|
optimizer: torch.optim.Optimizer,
|
|
scheduler: torch.optim.lr_scheduler.LRScheduler,
|
|
grad_scaler: torch.amp.GradScaler,
|
|
iteration: int,
|
|
) -> None:
|
|
"""Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
|
|
|
|
Args:
|
|
model (ImaginaireModel): The PyTorch model.
|
|
optimizer (torch.optim.Optimizer): The model optimizer.
|
|
scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
|
|
grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
|
|
iteration (int): Current iteration number.
|
|
"""
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
|
|
self.get_previous_checkpoint_results(wait_for=0)
|
|
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_save_checkpoint_start(model, iteration)
|
|
|
|
checkpoint_file = f"iter_{iteration:09}"
|
|
to_save_dict = {
|
|
"model": ModelWrapper(model).state_dict(),
|
|
"optim": OptimizerWrapper(model, optimizer).state_dict(),
|
|
"scheduler": scheduler.state_dict(),
|
|
"trainer": {
|
|
"grad_scaler": grad_scaler.state_dict(),
|
|
"iteration": iteration,
|
|
},
|
|
}
|
|
for k in to_save_dict.keys():
|
|
output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{k}")
|
|
to_save_dict[k] = (to_save_dict[k], output_dirname)
|
|
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict)
|
|
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
|
|
self._async_with_pinned_memory(checkpoint_file, to_save_dict)
|
|
else:
|
|
start_time = time.monotonic()
|
|
try:
|
|
self.save_state_dict_worker(to_save_dict, checkpoint_file)
|
|
finally:
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_save_checkpoint_success(
|
|
iteration=iteration, elapsed_time=time.monotonic() - start_time
|
|
)
|
|
|
|
# This measures exposed (synchronous) checkpoint time, on_save_checkpoint_success()
|
|
# is instead called to measure the entire duration for asynchronous checkpoint for the async case too.
|
|
if self.callbacks is not None:
|
|
self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)
|
|
|
|
def finalize(self) -> None:
|
|
super().finalize()
|
|
if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
|
|
if self.mp and self.mp.is_alive():
|
|
self.mp_queue_send.put(Terminate())
|
|
self.get_previous_checkpoint_results(wait_for=60)
|
|
self.mp.join()
|