diff --git a/.github/scripts/ci/run_colossalai_batch.sh b/.github/scripts/ci/run_colossalai_batch.sh index 01db482cca27..a44385f993f9 100644 --- a/.github/scripts/ci/run_colossalai_batch.sh +++ b/.github/scripts/ci/run_colossalai_batch.sh @@ -72,6 +72,7 @@ case "${BATCH}" in DEFAULT_TIMEOUT_MIN=15 TESTS=( docker/torch213/test_verify_image.py + tests/test_checkpoint_io/test_optimizer_safetensors_v2.py tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py tests/test_lazy/test_chatglm_cpu.py diff --git a/.github/scripts/ci/run_cpu_unit_tests.sh b/.github/scripts/ci/run_cpu_unit_tests.sh index 2677eccf2d84..ae22911eb0ec 100644 --- a/.github/scripts/ci/run_cpu_unit_tests.sh +++ b/.github/scripts/ci/run_cpu_unit_tests.sh @@ -10,6 +10,7 @@ PYTHON="${COLOSSAL_PYTHON:-python}" TESTS=( .github/workflows/scripts/example_checks/test_run_test_ci.py docker/torch213/test_verify_image.py + tests/test_checkpoint_io/test_optimizer_safetensors_v2.py tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py tests/test_lazy/test_chatglm_cpu.py diff --git a/colossalai/booster/plugin/gemini_plugin.py b/colossalai/booster/plugin/gemini_plugin.py index a81f9b05d7d7..0d8846e07b30 100644 --- a/colossalai/booster/plugin/gemini_plugin.py +++ b/colossalai/booster/plugin/gemini_plugin.py @@ -20,8 +20,11 @@ create_pinned_state_dict, get_model_base_filenames, get_optimizer_base_filenames, + load_param_groups, load_state_dict_shards, + prepare_pinned_state_dict, save_config_file, + save_param_groups, save_state_dict, save_state_dict_shards, ) @@ -128,8 +131,9 @@ def save_unsharded_optimizer( from colossalai.utils.safetensors import _flatten_optim_state_dict, save flatten_state_dict, metadata = _flatten_optim_state_dict(state_dict) - if id(optimizer) not in self.pinned_state_dicts: - self.pinned_state_dicts[id(optimizer)] = create_pinned_state_dict(flatten_state_dict) + self.pinned_state_dicts[id(optimizer)] = prepare_pinned_state_dict( + flatten_state_dict, self.pinned_state_dicts.get(id(optimizer)) + ) for k, v in flatten_state_dict.items(): self.pinned_state_dicts[id(optimizer)][k].copy_(v) flatten_state_dict[k] = self.pinned_state_dicts[id(optimizer)][k] @@ -264,12 +268,13 @@ def save_sharded_optimizer( states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix, use_safetensors=use_async) index_file = CheckpointIndexFile(checkpoint) index_file.append_meta_data("param_groups", param_group_file) + index_file.append_meta_data("param_groups_format", 2) # Store the information of param groups to param_group_file. if self.coordinator.is_master(): group_file_path = os.path.join(checkpoint, param_group_file) param_groups = optimizer.get_param_groups_for_saving() - torch.save(param_groups, group_file_path) + save_param_groups({"param_groups": param_groups}, group_file_path) # States are broken into shards within max_shard_size. if use_async and self.coordinator.is_master(): @@ -333,7 +338,7 @@ def load_sharded_optimizer( assert isinstance(optimizer, GeminiOptimizer) # Read checkpoint index file. - ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) + ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file, is_optimizer=True) # Load param_groups. param_group_path = ckpt_index_file.get_param_group_filename() @@ -342,7 +347,7 @@ def load_sharded_optimizer( f"Invalid index file path {checkpoint_index_file} for an optimizer. \ Lacking param group file under current directory." ) - saved_param_groups = torch.load(param_group_path) + saved_param_groups = load_param_groups(param_group_path) optimizer.load_param_groups(saved_param_groups) checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames() diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py index a063f3eaee4f..5aeb928e189e 100644 --- a/colossalai/booster/plugin/low_level_zero_plugin.py +++ b/colossalai/booster/plugin/low_level_zero_plugin.py @@ -23,8 +23,8 @@ create_pinned_state_dict, get_optimizer_base_filenames, get_shard_filename, + load_optimizer_state, load_param_groups_into_optimizer, - load_state_dict, load_state_dict_shards, load_states_into_optimizer, save_param_groups, @@ -153,9 +153,9 @@ def load_unsharded_optimizer( if use_async: from colossalai.utils.safetensors import load_flat - checkpoint = load_flat(checkpoint) + checkpoint = load_flat(checkpoint, expected_kind="full") else: - checkpoint = load_state_dict(checkpoint) + checkpoint = load_optimizer_state(checkpoint) if not low_cpu_mem_mode: checkpoint = create_pinned_state_dict(checkpoint, empty=False, num_threads=num_threads) optimizer.load_state_dict(checkpoint) @@ -207,6 +207,7 @@ def save_sharded_optimizer( states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix, use_safetensors=use_async) index_file = CheckpointIndexFile(checkpoint) index_file.append_meta_data("param_groups", param_group_file) + index_file.append_meta_data("param_groups_format", 2) # Store the information of param groups to param_group_file. if self.coordinator.is_master(): @@ -263,7 +264,7 @@ def load_sharded_optimizer( optimizer = optimizer.unwrap() # Read checkpoint index file. - ckpt_index_file = CheckpointIndexFile.from_file(index_file_path) + ckpt_index_file = CheckpointIndexFile.from_file(index_file_path, is_optimizer=True) # Load param_groups param_group_path = ckpt_index_file.get_param_group_filename() diff --git a/colossalai/booster/plugin/torch_fsdp_plugin.py b/colossalai/booster/plugin/torch_fsdp_plugin.py index 6e652e549b5e..05cc432a673d 100644 --- a/colossalai/booster/plugin/torch_fsdp_plugin.py +++ b/colossalai/booster/plugin/torch_fsdp_plugin.py @@ -26,7 +26,11 @@ from torch.utils.data import DataLoader from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO, utils -from colossalai.checkpoint_io.utils import async_save_state_dict_shards, create_pinned_state_dict +from colossalai.checkpoint_io.utils import ( + async_save_state_dict_shards, + create_pinned_state_dict, + prepare_pinned_state_dict, +) from colossalai.cluster import DistCoordinator from colossalai.interface import ModelWrapper, OptimizerWrapper from colossalai.logging import get_dist_logger @@ -56,9 +60,9 @@ def load_unsharded_optimizer( ): assert isinstance(optimizer, FSDPOptimizerWrapper), "Please boost the optimizer before loading!" if checkpoint.endswith(".safetensors"): - checkpoint = load_flat(checkpoint, seperator=".") + checkpoint = load_flat(checkpoint, seperator=".", expected_kind="full") else: - checkpoint = utils.load_state_dict(checkpoint) + checkpoint = utils.load_optimizer_state(checkpoint) fsdp_model = optimizer.unwrap_model() full_optimizer_state = FSDP.full_optim_state_dict(fsdp_model, optim=optimizer, rank0_only=False) @@ -151,8 +155,9 @@ def pack_group(group: Dict[str, Any]) -> Dict[str, Any]: from colossalai.utils.safetensors import _flatten_optim_state_dict, save flatten_state_dict, metadata = _flatten_optim_state_dict(full_optimizer_state, seperator=".") - if id(optimizer) not in self.pinned_state_dicts: - self.pinned_state_dicts[id(optimizer)] = create_pinned_state_dict(flatten_state_dict) + self.pinned_state_dicts[id(optimizer)] = prepare_pinned_state_dict( + flatten_state_dict, self.pinned_state_dicts.get(id(optimizer)) + ) for k, v in flatten_state_dict.items(): self.pinned_state_dicts[id(optimizer)][k].copy_(v) flatten_state_dict[k] = self.pinned_state_dicts[id(optimizer)][k] @@ -318,6 +323,8 @@ def pack_group(group: Dict[str, Any]) -> Dict[str, Any]: index_file = CheckpointIndexFile(checkpoint) index_file.append_meta_data("param_groups", param_group_file) + + index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) utils.save_param_groups(fsdp_optim_state, group_file_path) @@ -373,7 +380,7 @@ def load_sharded_optimizer( """ assert isinstance(optimizer, FSDPOptimizerWrapper), "Please boost the optimizer before saving!" - ckpt_index_file = CheckpointIndexFile.from_file(index_file_path) + ckpt_index_file = CheckpointIndexFile.from_file(index_file_path, is_optimizer=True) # Load param_groups param_group_path = ckpt_index_file.get_param_group_filename() @@ -383,7 +390,7 @@ def load_sharded_optimizer( "Looking param group file under current directory." ) - saved_param_groups = torch.load(param_group_path) + saved_param_groups = utils.load_param_groups(param_group_path) # Load param fsdp_optim_state = {} diff --git a/colossalai/checkpoint_io/general_checkpoint_io.py b/colossalai/checkpoint_io/general_checkpoint_io.py index bda06a338a88..69117646deb1 100644 --- a/colossalai/checkpoint_io/general_checkpoint_io.py +++ b/colossalai/checkpoint_io/general_checkpoint_io.py @@ -17,12 +17,14 @@ get_model_base_filenames, get_optimizer_base_filenames, is_safetensors_available, + load_optimizer_state, load_param_groups_into_optimizer, load_state_dict, load_state_dict_into_model, load_state_dict_shards, load_states_into_optimizer, merge_optimizer_param_group_defaults, + prepare_pinned_state_dict, save_config_file, save_param_groups, save_state_dict, @@ -82,7 +84,7 @@ def load_sharded_optimizer( """ # Read checkpoint index file. - ckpt_index_file = CheckpointIndexFile.from_file(index_file_path) + ckpt_index_file = CheckpointIndexFile.from_file(index_file_path, is_optimizer=True) # Load param_groups param_group_path = ckpt_index_file.get_param_group_filename() @@ -135,6 +137,7 @@ def save_sharded_optimizer( # Store the information of param groups to param_group_file. index_file.append_meta_data("param_groups", param_group_file) + index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) save_param_groups(state_dict, group_file_path) @@ -176,9 +179,9 @@ def load_unsharded_optimizer( self, optimizer: Optimizer, checkpoint: Path, low_cpu_mem_mode: bool = True, num_threads: int = 1 ): if checkpoint.endswith(".safetensors"): - checkpoint = load_flat(checkpoint) + checkpoint = load_flat(checkpoint, expected_kind="full") else: - checkpoint = load_state_dict(checkpoint) + checkpoint = load_optimizer_state(checkpoint) if not low_cpu_mem_mode: checkpoint = create_pinned_state_dict(checkpoint, empty=False, num_threads=num_threads) merge_optimizer_param_group_defaults(checkpoint["param_groups"], optimizer.param_groups) @@ -197,8 +200,9 @@ def save_unsharded_optimizer( from colossalai.utils.safetensors import _flatten_optim_state_dict, move_and_save flatten_state_dict, metadata = _flatten_optim_state_dict(state_dict) - if id(optimizer) not in self.pinned_state_dicts: - self.pinned_state_dicts[id(optimizer)] = create_pinned_state_dict(flatten_state_dict) + self.pinned_state_dicts[id(optimizer)] = prepare_pinned_state_dict( + flatten_state_dict, self.pinned_state_dicts.get(id(optimizer)) + ) writer = move_and_save( path=checkpoint, state_dict=flatten_state_dict, diff --git a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py index 4749a9ab479c..b42fbd3a7533 100644 --- a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py +++ b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py @@ -38,9 +38,13 @@ get_model_base_filenames, get_optimizer_base_filenames, is_safetensors_available, + load_optim_or_model_shard, + load_optimizer_state, + load_param_groups, load_shard_state_dict, load_state_dict, load_state_dict_into_model, + prepare_pinned_state_dict, save_config_file, save_param_groups, save_state_dict, @@ -596,6 +600,7 @@ def save_sharded_optimizer( if control_saving: # Store param groups. index_file.append_meta_data("param_groups", param_group_file) + index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) param_groups = [ {**group, "params": group_info["params"]} @@ -665,15 +670,20 @@ def save_sharded_optimizer( if self.pp_rank == 0: final_index_file = CheckpointIndexFile(checkpoint) final_index_file.append_meta_data("total_size", 0) + if use_async: + final_index_file.append_meta_data("optimizer_format", 2) for filename in os.listdir(tmp_index_file_folder): - stage_index_file = CheckpointIndexFile.from_file(os.path.join(tmp_index_file_folder, filename)) + stage_index_file = CheckpointIndexFile.from_file( + os.path.join(tmp_index_file_folder, filename), is_optimizer=True + ) final_index_file.metadata["total_size"] += stage_index_file.metadata["total_size"] for param_id, state_filename in stage_index_file.weight_map.items(): final_index_file.append_weight_map(param_id, state_filename) # Store param groups. final_index_file.append_meta_data("param_groups", param_group_file) + final_index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) param_groups = [ {**group, "params": group_info["params"]} @@ -729,7 +739,7 @@ def _get_param_id_from_optimizer_param( id_map[param_id] = param # Read checkpoint index file. - ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) + ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file, is_optimizer=True) ckpt_root_path = ckpt_index_file.root_path weight_map = ckpt_index_file.weight_map weight_map = {int(k): v for k, v in weight_map.items()} # convert saved id from str to int @@ -741,7 +751,7 @@ def _get_param_id_from_optimizer_param( f"Invalid index file path {checkpoint_index_file} for an optimizer. \ Lacking param group file under current directory." ) - saved_groups = torch.load(param_group_path) + saved_groups = load_param_groups(param_group_path) updated_groups = [] for old_pg, saved_pg in zip(optimizer.optim.param_groups, saved_groups): @@ -768,10 +778,7 @@ def _get_param_id_from_optimizer_param( continue file_path = os.path.join(ckpt_root_path, filename) - if file_path.endswith(".safetensors"): - state_dict = load_flat(file_path) - else: - state_dict = load_shard_state_dict(Path(file_path), use_safetensors=False) + state_dict = load_optim_or_model_shard(file_path, is_optim=True, use_safetensors=False) if not low_cpu_mem_mode: state_dict = create_pinned_state_dict(state_dict, empty=False, num_threads=num_threads) self.load_states_into_optimizer(optimizer, state_dict, id_map) @@ -963,11 +970,11 @@ def save_unsharded_optimizer( from colossalai.utils.safetensors import save flatten_state_dict, metadata = _flatten_optim_state_dict(state_dict) - if id(optimizer) not in self.pinned_state_dicts: - self.pinned_state_dicts = create_pinned_state_dict(flatten_state_dict) + buffers = prepare_pinned_state_dict(flatten_state_dict, self.pinned_state_dicts.get(id(optimizer))) + self.pinned_state_dicts[id(optimizer)] = buffers for k, v in flatten_state_dict.items(): - self.pinned_state_dicts[k].copy_(v) - flatten_state_dict[k] = self.pinned_state_dicts[k] + buffers[k].copy_(v) + flatten_state_dict[k] = buffers[k] writer = save(path=checkpoint, state_dict=flatten_state_dict, metadata=metadata) self.async_writers.append(writer) else: @@ -991,11 +998,11 @@ def save_unsharded_optimizer( from colossalai.utils.safetensors import save flatten_state_dict, metadata = _flatten_optim_state_dict(state_dict) - if id(optimizer) not in self.pinned_state_dicts: - self.pinned_state_dicts = create_pinned_state_dict(flatten_state_dict) + buffers = prepare_pinned_state_dict(flatten_state_dict, self.pinned_state_dicts.get(id(optimizer))) + self.pinned_state_dicts[id(optimizer)] = buffers for k, v in flatten_state_dict.items(): - self.pinned_state_dicts[k].copy_(v) - flatten_state_dict[k] = self.pinned_state_dicts[k] + buffers[k].copy_(v) + flatten_state_dict[k] = buffers[k] writer = save(path=checkpoint, state_dict=flatten_state_dict, metadata=metadata) self.async_writers.append(writer) else: @@ -1028,9 +1035,9 @@ def _get_param_id_from_optimizer_param( # Complete optimizer state_dict loaded from checkpoint, need to be processed later. if checkpoint.endswith(".safetensors"): - state_dict = load_flat(checkpoint) + state_dict = load_flat(checkpoint, expected_kind="full") else: - state_dict = load_state_dict(checkpoint) + state_dict = load_optimizer_state(checkpoint) if not low_cpu_mem_mode: state_dict = create_pinned_state_dict(state_dict, empty=False, num_threads=num_threads) diff --git a/colossalai/checkpoint_io/index_file.py b/colossalai/checkpoint_io/index_file.py index da12c146f2c3..173f44bdb979 100644 --- a/colossalai/checkpoint_io/index_file.py +++ b/colossalai/checkpoint_io/index_file.py @@ -20,25 +20,27 @@ class CheckpointIndexFile: >>> index.export('new_index.json') """ - def __init__(self, root_path=None) -> None: + def __init__(self, root_path=None, *, is_optimizer: bool = False) -> None: self.root_path = root_path + self.is_optimizer = is_optimizer # use ordered dict to preserve the tensor checkpoint order self.metadata: Dict = OrderedDict() self.weight_map: Dict = OrderedDict() @staticmethod - def from_file(index_path: Union[str, Path]): + def from_file(index_path: Union[str, Path], *, is_optimizer: bool = False): """ Create a CheckpointIndexFile object from a json file. Args: index_path (str): path to the json file. + is_optimizer (bool): enforce optimizer metadata and path limits. Returns: CheckpointIndexFile: CheckpointIndexFile object. """ - index = CheckpointIndexFile() + index = CheckpointIndexFile(is_optimizer=is_optimizer) index.load(index_path) return index @@ -49,9 +51,18 @@ def load(self, json_path: str): Args: json_path (str): path to the json file. """ - # load the json file - with open(json_path, "r") as f: - index = json.load(f) + with open(json_path, "r", encoding="utf-8") as f: + if self.is_optimizer: + from colossalai.utils._optimizer_safetensors import MAX_SCHEMA_BYTES, _pairs + + raw = f.read(MAX_SCHEMA_BYTES + 1) + if len(raw.encode("utf-8")) > MAX_SCHEMA_BYTES: + raise ValueError("Checkpoint index exceeds the metadata byte limit") + index = json.loads(raw, object_pairs_hook=_pairs) + else: + index = json.load(f) + if type(index) is not dict: + raise ValueError("Checkpoint index must be a dictionary") # assign attributes if exists if "metadata" in index: @@ -61,6 +72,26 @@ def load(self, json_path: str): # assign the root directory for the index file self.root_path = Path(json_path).absolute().parent + if not isinstance(self.metadata, dict) or not isinstance(self.weight_map, dict): + raise ValueError("Invalid checkpoint index metadata or weight map") + for filename in self.weight_map.values(): + self._resolve_file(filename) + if "param_groups" in self.metadata: + self._resolve_file(self.metadata["param_groups"]) + + def _resolve_file(self, filename): + if type(filename) is not str or not filename: + raise ValueError("Checkpoint index filenames must be nonempty strings") + path = Path(self.root_path) / filename + if self.is_optimizer: + if Path(filename).is_absolute(): + raise ValueError("Checkpoint index filenames must be relative paths") + root = Path(self.root_path).resolve() + if not path.resolve().is_relative_to(root): + raise ValueError("Checkpoint index path escapes its root directory") + # Validate real paths, but retain the named path for format dispatch. + # Model snapshots may link to a shared blob directory outside the root. + return str(path) def export(self, json_path: str): """ @@ -122,7 +153,7 @@ def get_checkpoint_filenames(self) -> List[str]: checkpoint_files = sorted(list(set(self.weight_map.values()))) # get the absolute paths for all checkpoint files - checkpoint_files = [str(self.root_path.joinpath(f)) for f in checkpoint_files] + checkpoint_files = [self._resolve_file(f) for f in checkpoint_files] dtensor_list = [] checkpoint_list = [] @@ -151,6 +182,7 @@ def get_checkpoint_file(self, param_name: str) -> str: str: checkpoint file name. """ ckpt_path = self.weight_map[param_name] + self._resolve_file(ckpt_path) return ckpt_path def get_all_param_names(self): @@ -167,7 +199,7 @@ def get_param_group_filename(self) -> Union[str, None]: """ filename = self.metadata.get("param_groups", None) if filename: - return str(self.root_path.joinpath(filename)) + return self._resolve_file(filename) else: return None diff --git a/colossalai/checkpoint_io/moe_checkpoint.py b/colossalai/checkpoint_io/moe_checkpoint.py index a24cd4d638d8..7813b9133dca 100644 --- a/colossalai/checkpoint_io/moe_checkpoint.py +++ b/colossalai/checkpoint_io/moe_checkpoint.py @@ -23,8 +23,9 @@ get_lora_state_dict, get_model_base_filenames, get_optimizer_base_filenames, - load_shard_state_dict, - load_state_dict, + load_optim_or_model_shard, + load_optimizer_state, + load_param_groups, load_states_into_optimizer, save_config_file, save_param_groups, @@ -496,6 +497,7 @@ def save_sharded_optimizer( if control_saving: # Store param groups. index_file.append_meta_data("param_groups", param_group_file) + index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) param_groups = [ {**group, "params": group_info["params"]} @@ -551,13 +553,16 @@ def save_sharded_optimizer( final_index_file.append_meta_data("total_size", 0) for filename in os.listdir(tmp_index_file_folder): - stage_index_file = CheckpointIndexFile.from_file(os.path.join(tmp_index_file_folder, filename)) + stage_index_file = CheckpointIndexFile.from_file( + os.path.join(tmp_index_file_folder, filename), is_optimizer=True + ) final_index_file.metadata["total_size"] += stage_index_file.metadata["total_size"] for param_id, state_filename in stage_index_file.weight_map.items(): final_index_file.append_weight_map(param_id, state_filename) # Store param groups. final_index_file.append_meta_data("param_groups", param_group_file) + final_index_file.append_meta_data("param_groups_format", 2) group_file_path = os.path.join(checkpoint, param_group_file) param_groups = [ {**group, "params": group_info["params"]} @@ -613,7 +618,7 @@ def _get_param_id_from_optimizer_param( id_map[param_id] = param # Read checkpoint index file. - ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) + ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file, is_optimizer=True) ckpt_root_path = ckpt_index_file.root_path weight_map = ckpt_index_file.weight_map weight_map = {int(k): v for k, v in weight_map.items()} # convert saved id from str to int @@ -625,7 +630,7 @@ def _get_param_id_from_optimizer_param( f"Invalid index file path {checkpoint_index_file} for an optimizer. \ Lacking param group file under current directory." ) - saved_groups = torch.load(param_group_path) + saved_groups = load_param_groups(param_group_path) updated_groups = [] for old_pg, saved_pg in zip(optimizer.optim.param_groups, saved_groups): @@ -657,7 +662,7 @@ def _get_param_id_from_optimizer_param( continue file_path = os.path.join(ckpt_root_path, filename) - state_dict = load_shard_state_dict(Path(file_path), use_safetensors=False) + state_dict = load_optim_or_model_shard(file_path, is_optim=True, use_safetensors=False) load_states_into_optimizer(optimizer.optim, state_dict, id_map, strict=True) loaded_file.add(filename) @@ -961,7 +966,7 @@ def _get_param_id_from_optimizer_param( assert isinstance(optimizer, OptimizerWrapper), "Please boost the optimizer before loading!" # Complete optimizer state_dict loaded from checkpoint, need to be processed later. - state_dict = load_state_dict(checkpoint) + state_dict = load_optimizer_state(checkpoint) # Load param_groups. updated_groups = [] diff --git a/colossalai/checkpoint_io/optimizer_migration.py b/colossalai/checkpoint_io/optimizer_migration.py new file mode 100644 index 000000000000..30f0cd893d26 --- /dev/null +++ b/colossalai/checkpoint_io/optimizer_migration.py @@ -0,0 +1,131 @@ +"""Explicit offline migration of trusted legacy optimizer checkpoints. + +This module is not used by training loaders. Legacy decoding can execute code. +""" + +import hashlib +import json +import os +import tempfile +from pathlib import Path + +import torch +from safetensors import safe_open + +from colossalai.utils.safetensors import load_flat, save_optimizer_file + +from .index_file import CheckpointIndexFile + + +def _sha256(path): + digest = hashlib.sha256() + with open(path, "rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_trusted_legacy(path): + if path.suffix != ".safetensors": + return torch.load(path, map_location="cpu", weights_only=False) + with safe_open(path, framework="pt", device="cpu") as reader: + metadata = reader.metadata() or {} + if "colossalai_optimizer_format" in metadata: + return load_flat(path) + objects = json.loads(metadata.get("non_tensor_keys", "[]")) + if type(objects) is not list or any(type(key) is not str for key in objects): + raise ValueError("Invalid legacy optimizer object list") + if len(set(objects)) != len(objects) or not set(objects).issubset(reader.keys()): + raise ValueError("Invalid legacy optimizer object keys") + flat = {} + for key in reader.keys(): + value = reader.get_tensor(key) + if key in objects: + import pickle + + if value.dtype != torch.uint8 or value.dim() != 1: + raise ValueError("Legacy optimizer object must be a byte tensor") + # Deliberately isolated behind the caller's explicit trust opt-in. + value = pickle.loads(value.numpy().tobytes()) + flat[key] = value + states = {} + groups = flat.pop("param_groups", None) + for key, value in flat.items(): + parts = key.split(".") + if len(parts) != 3 or parts[0] != "state" or not parts[1].isdigit(): + raise ValueError("Invalid legacy optimizer state key") + states.setdefault(int(parts[1]), {})[parts[2]] = value + return states if groups is None else {"state": states, "param_groups": groups} + + +def _write_exclusive(path, state): + with tempfile.TemporaryDirectory(dir=path.parent, prefix=".optimizer-migration-") as staging: + temporary = Path(staging) / "optimizer.safetensors" + save_optimizer_file(temporary, state) + # Unlike replace(), link() cannot overwrite a concurrently created file. + os.link(temporary, path) + + +def migrate_optimizer_checkpoint(source, destination, *, trust_legacy_pickle=False): + """Convert a trusted file or optimizer index/directory without overwriting it. + + ``trust_legacy_pickle`` permits arbitrary code execution while reading the + legacy input. Use an isolated environment without network or credentials. + """ + if trust_legacy_pickle is not True: + raise ValueError("Explicit trust_legacy_pickle=True is required; legacy inputs can execute code") + source = Path(source).resolve(strict=True) + destination = Path(destination).absolute() + if destination.exists() or destination.is_symlink(): + raise FileExistsError(destination) + if source.is_dir(): + indices = list(source.glob("*optim*.index.json")) + if len(indices) != 1: + raise ValueError("Specify an optimizer index file; exactly one optimizer index is required") + source = indices[0] + if not source.name.endswith(".index.json") and destination.suffix != ".safetensors": + raise ValueError("A single-file migration destination must end with .safetensors") + inputs = {source.name: _sha256(source)} + + def read_input(path): + inputs[str(path.relative_to(source.parent))] = _sha256(path) + return _load_trusted_legacy(path) + + if not source.name.endswith(".index.json"): + state = _load_trusted_legacy(source) + if type(state) is list: + state = {"param_groups": state} + _write_exclusive(destination, state) + return {"format_version": 2, "inputs": inputs, "outputs": {destination.name: _sha256(destination)}} + index = CheckpointIndexFile.from_file(source, is_optimizer=True) + groups_path = index.get_param_group_filename() + if groups_path is None: + raise ValueError("Optimizer index has no parameter group file") + groups = read_input(Path(groups_path)) + if type(groups) is dict and "param_groups" in groups: + groups = groups["param_groups"] + # Reserve a new output directory. The index is published only after every + # file has been converted; failed conversion leaves no usable index. + destination.mkdir() + save_optimizer_file(destination / "optimizer_group.safetensors", {"param_groups": groups}) + output = CheckpointIndexFile(destination) + output.append_meta_data("param_groups", "optimizer_group.safetensors") + output.append_meta_data("param_groups_format", 2) + output.append_meta_data("optimizer_format", 2) + converted = {} + total_size = 0 + for filename in dict.fromkeys(index.weight_map.values()): + state = read_input(Path(index._resolve_file(filename))) + new_name = f"optimizer-{len(converted) + 1:05d}.safetensors" + save_optimizer_file(destination / new_name, state) + converted[filename] = new_name + total_size += (destination / new_name).stat().st_size + for param_id, filename in index.weight_map.items(): + output.append_weight_map(param_id, converted[filename]) + output.append_meta_data("total_size", total_size) + output.write_index_file("optimizer.safetensors.index.json") + return { + "format_version": 2, + "inputs": inputs, + "outputs": {path.name: _sha256(path) for path in destination.iterdir()}, + } diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py index 58aba87537dc..f259f9b31053 100644 --- a/colossalai/checkpoint_io/utils.py +++ b/colossalai/checkpoint_io/utils.py @@ -28,7 +28,7 @@ to_global_for_customized_distributed_tensor, ) from colossalai.utils import get_current_device -from colossalai.utils.safetensors import _flatten_optim_state_dict, load_flat +from colossalai.utils.safetensors import _flatten_optim_state_dict, load_flat, save_optimizer_file SAFE_WEIGHTS_NAME = "model.safetensors" WEIGHTS_NAME = "pytorch_model.bin" @@ -38,7 +38,7 @@ WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json" STATES_INDEX_NAME = "pytorch_optim.bin.index.json" SAFE_STATES_INDEX_NAME = "optimizer.safetensors.index.json" -GROUP_FILE_NAME = "pytorch_optim_group.bin" +GROUP_FILE_NAME = "optimizer_group.safetensors" # ====================================== # General helper functions @@ -300,6 +300,9 @@ def async_save_state_dict_shards( """ from colossalai.utils.safetensors import save + if state_preprocess: + index_file.append_meta_data("optimizer_format", 2) + total_size = 0 shard_filenames = [] writers = [] @@ -359,6 +362,9 @@ def async_move_save_state_dict_shards( """ from colossalai.utils.safetensors import move_and_save + if state_preprocess: + index_file.append_meta_data("optimizer_format", 2) + total_size = 0 shard_filenames = [] if pinned_state_dict is None: @@ -384,11 +390,7 @@ def async_move_save_state_dict_shards( state_dict = shard metadata = None - if pinned_state_dict is not None: - sub_pinned_state_dict = {k: pinned_state_dict[k] for k in state_dict.keys()} - else: - sub_pinned_state_dict = create_pinned_state_dict(state_dict) - returned_state_dict.update(sub_pinned_state_dict) + sub_pinned_state_dict = prepare_pinned_state_dict(state_dict, returned_state_dict) # Only save on master rank. writer = move_and_save(checkpoint_file_path, state_dict, sub_pinned_state_dict, metadata) @@ -501,7 +503,30 @@ def save_param_groups(state_dict: dict, group_file_path: str) -> None: group_file_path (str): path to the group file. """ param_groups = state_dict["param_groups"] - torch.save(param_groups, group_file_path) + save_optimizer_file(group_file_path, {"param_groups": param_groups}) + + +def _load_legacy_optimizer_file(path): + """Restrict PyTorch optimizer files to versions fixing known weights-only issues.""" + if Version(torch.__version__) < Version("2.10.0"): + raise RuntimeError( + "Loading legacy optimizer checkpoints requires torch >= 2.10.0. " + "Upgrade torch, or explicitly migrate a trusted checkpoint with " + "colossalai checkpoint migrate-optimizer --trust-legacy-pickle." + ) + return torch.load(path, map_location="cpu", weights_only=True) + + +def load_param_groups(group_file_path): + if str(group_file_path).endswith(".safetensors"): + return load_flat(group_file_path, expected_kind="param_groups")["param_groups"] + return _load_legacy_optimizer_file(group_file_path) + + +def load_optimizer_state(checkpoint): + if str(checkpoint).endswith(".safetensors"): + return load_flat(checkpoint, expected_kind="full") + return _load_legacy_optimizer_file(checkpoint) def merge_optimizer_param_group_defaults(saved_groups: List[dict], current_groups: List[dict]) -> None: @@ -706,7 +731,7 @@ def load_shard_state_dict(checkpoint_file: Path, use_safetensors: bool = False): return safe_load_file(checkpoint_file) else: - return torch.load(checkpoint_file, map_location=torch.device("cpu")) + return torch.load(checkpoint_file, map_location=torch.device("cpu"), weights_only=True) def load_state_dict_into_model( @@ -769,7 +794,7 @@ def load_param_groups_into_optimizer(optimizer: Optimizer, param_group_path: str # Load list of param_groups from given file path. # The params in saved_groups are in the form of integer indices. - saved_groups = torch.load(param_group_path, map_location=torch.device("cpu")) + saved_groups = load_param_groups(param_group_path) if not isinstance(saved_groups, List): raise ValueError(f"The param_groups saved at {param_group_path} is not of List type") @@ -937,7 +962,7 @@ def load_state_dict(checkpoint_file_path: Path): else: # load with torch - return torch.load(checkpoint_file_path, map_location=torch.device("cpu")) + return torch.load(checkpoint_file_path, map_location=torch.device("cpu"), weights_only=True) def add_prefix(weights_name: str, prefix: Optional[str] = None) -> str: @@ -1013,12 +1038,26 @@ def create_pinned_state_dict( return tree_unflatten(elems, spec) +def prepare_pinned_state_dict(state_dict, cache=None): + """Reuse compatible buffers, including newly initialized optimizer state.""" + if cache is None: + cache = {} + result = {} + for key, tensor in state_dict.items(): + previous = cache.get(key) + if not isinstance(previous, torch.Tensor) or previous.shape != tensor.shape or previous.dtype != tensor.dtype: + previous = _pin_tensor(tensor) + cache[key] = previous + result[key] = previous + return result + + def load_optim_or_model_shard(path: str, is_optim: bool, use_safetensors: bool) -> dict: if is_optim: if path.endswith(".safetensors"): - state_dict = load_flat(path) + state_dict = load_flat(path, expected_kind="state_shard") else: - state_dict = load_shard_state_dict(Path(path), use_safetensors=False) + state_dict = _load_legacy_optimizer_file(path) else: state_dict = load_shard_state_dict(Path(path), use_safetensors) return state_dict diff --git a/colossalai/cli/checkpoint.py b/colossalai/cli/checkpoint.py new file mode 100644 index 000000000000..94565828ae65 --- /dev/null +++ b/colossalai/cli/checkpoint.py @@ -0,0 +1,29 @@ +import json +from pathlib import Path + +import click + + +@click.group() +def checkpoint(): + """Checkpoint format tools.""" + + +@checkpoint.command("migrate-optimizer") +@click.argument("source", type=click.Path(exists=True, path_type=Path)) +@click.argument("destination", type=click.Path(path_type=Path)) +@click.option( + "--trust-legacy-pickle", + is_flag=True, + required=True, + help="Allow legacy decoding, which can execute code. Use only trusted input in an isolated environment.", +) +def migrate_optimizer(source, destination, trust_legacy_pickle): + """Convert a trusted legacy optimizer file or directory to data-only v2.""" + from colossalai.checkpoint_io.optimizer_migration import migrate_optimizer_checkpoint + + try: + result = migrate_optimizer_checkpoint(source, destination, trust_legacy_pickle=trust_legacy_pickle) + click.echo(json.dumps(result, indent=2)) + except (OSError, ValueError, TypeError) as exc: + raise click.ClickException(str(exc)) from exc diff --git a/colossalai/cli/cli.py b/colossalai/cli/cli.py index 0d94fe59f8ae..0b5e562ee30f 100644 --- a/colossalai/cli/cli.py +++ b/colossalai/cli/cli.py @@ -1,6 +1,7 @@ import click from .check import check +from .checkpoint import checkpoint from .launcher import run @@ -17,6 +18,7 @@ def cli(): cli.add_command(run) cli.add_command(check) +cli.add_command(checkpoint) if __name__ == "__main__": cli() diff --git a/colossalai/utils/_optimizer_safetensors.py b/colossalai/utils/_optimizer_safetensors.py new file mode 100644 index 000000000000..d7e0f4b6ccab --- /dev/null +++ b/colossalai/utils/_optimizer_safetensors.py @@ -0,0 +1,251 @@ +"""Data-only optimizer checkpoint schema. No Python object reconstruction.""" + +import json +import math +from collections import OrderedDict + +import torch + +FORMAT_KEY = "colossalai_optimizer_format" +FORMAT_VERSION = "2" +SCHEMA_KEY = "optimizer_schema" +KINDS = ("full", "state_shard", "param_groups") +MAX_SCHEMA_BYTES = 16 * 1024 * 1024 +MAX_NODES = 1_000_000 +MAX_DEPTH = 64 + + +class LegacyOptimizerCheckpointError(ValueError): + """An optimizer checkpoint requires legacy pickle object decoding.""" + + +def _pairs(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate optimizer schema key: {key}") + result[key] = value + return result + + +def _integer(value): + if len(value) > 4096: + raise ValueError("Optimizer schema integer is too large") + return int(value) + + +def _constant(value): + raise ValueError(f"Non-standard JSON constant in optimizer schema: {value}") + + +def _budget(depth, count): + count[0] += 1 + if depth > MAX_DEPTH or count[0] > MAX_NODES: + raise ValueError("Optimizer schema exceeds the depth or node limit") + + +def _fields(node, expected): + if type(node) is not dict or set(node) != set(expected): + raise ValueError("Invalid optimizer schema fields") + + +def _validate_payload(payload, kind): + if type(payload) is not dict: + raise ValueError("Optimizer payload must be a dictionary") + expected = {"state", "param_groups"} if kind == "full" else {"state" if kind == "state_shard" else "param_groups"} + if set(payload) != expected: + raise ValueError(f"Invalid fields for optimizer checkpoint kind {kind}") + if "state" in payload: + states = payload["state"] + if type(states) not in (dict, OrderedDict): + raise ValueError("Optimizer state must be a dictionary") + for param_id, state in states.items(): + if type(param_id) is not int or param_id < 0 or type(state) not in (dict, OrderedDict): + raise ValueError("Optimizer state requires non-negative integer IDs and dictionaries") + if any(type(key) is not str for key in state): + raise ValueError("Optimizer state field names must be strings") + if "param_groups" in payload: + groups = payload["param_groups"] + if type(groups) is not list: + raise ValueError("Optimizer param_groups must be a list") + seen = set() + for group in groups: + if type(group) not in (dict, OrderedDict) or "params" not in group: + raise ValueError("Each optimizer parameter group requires params") + if any(type(key) is not str for key in group): + raise ValueError("Optimizer parameter group keys must be strings") + if type(group["params"]) is not list: + raise ValueError("Optimizer params must be a list of integer IDs") + for param_id in group["params"]: + if type(param_id) is not int or param_id < 0 or param_id in seen: + raise ValueError("Optimizer parameter IDs must be non-negative and unique") + seen.add(param_id) + if kind == "full" and not set(payload["state"]).issubset(seen): + raise ValueError("Optimizer state references an ID absent from param_groups") + + +def flatten_optimizer(state_dict, kind=None): + if type(state_dict) not in (dict, OrderedDict): + raise TypeError("Optimizer state_dict must be a dictionary") + if kind is None: + if "param_groups" in state_dict: + kind = "full" if "state" in state_dict else "param_groups" + else: + kind = "state_shard" + if kind not in KINDS: + raise ValueError(f"Unsupported optimizer checkpoint kind: {kind}") + payload = dict(state_dict) if kind != "state_shard" else {"state": state_dict} + _validate_payload(payload, kind) + tensors = {} + active = set() + count = [0] + + def encode(value, path, depth=0): + _budget(depth, count) + if isinstance(value, torch.Tensor): + if value.layout != torch.strided or value.is_quantized: + raise TypeError(f"Only dense strided tensors are supported at {path}") + # Paths, rather than per-shard counters, keep pinned-cache names stable. + key = "optimizer.tensor." + json.dumps(path, ensure_ascii=True, separators=(",", ":")) + tensor = value.detach().contiguous() + tensors[key] = tensor + return {"type": "tensor_ref", "key": key, "dtype": str(tensor.dtype), "shape": list(tensor.shape)} + value_type = type(value) + if value is None: + return {"type": "null"} + if value_type in (bool, int, float, str): + if value_type is int and value.bit_length() > 4096: + raise ValueError(f"Optimizer integer is too large at {path}") + if value_type is float and not math.isfinite(value): + value = "nan" if math.isnan(value) else ("+inf" if value > 0 else "-inf") + return {"type": value_type.__name__, "value": value} + if value_type not in (list, tuple, dict, OrderedDict): + raise TypeError(f"Unsupported optimizer value {value_type.__name__} at {path}") + if id(value) in active: + raise ValueError(f"Cyclic optimizer data at {path}") + active.add(id(value)) + try: + if value_type in (list, tuple): + items = [encode(item, path + [("item", i)], depth + 1) for i, item in enumerate(value)] + return {"type": value_type.__name__, "items": items} + items = [] + for key, item in value.items(): + if type(key) not in (str, int): + raise TypeError(f"Unsupported optimizer dictionary key at {path}") + items.append( + [encode(key, path + [("key", key)], depth + 1), encode(item, path + [("value", key)], depth + 1)] + ) + return {"type": "dict", "items": items} + finally: + active.remove(id(value)) + + schema = json.dumps(encode(payload, []), ensure_ascii=True, allow_nan=False, separators=(",", ":")) + if len(schema.encode("utf-8")) > MAX_SCHEMA_BYTES: + raise ValueError("Optimizer schema exceeds the byte limit") + return tensors, {FORMAT_KEY: FORMAT_VERSION, "kind": kind, SCHEMA_KEY: schema} + + +def parse_optimizer_metadata(metadata, expected_kind=None): + metadata = metadata or {} + if "non_tensor_keys" in metadata: + raise LegacyOptimizerCheckpointError( + "Legacy pickle optimizer checkpoints are not loaded. Migrate a trusted checkpoint with " + "colossalai checkpoint migrate-optimizer --trust-legacy-pickle." + ) + version = metadata.get(FORMAT_KEY) + if version is None: + if SCHEMA_KEY in metadata or "kind" in metadata: + raise ValueError("Optimizer format version is missing") + if expected_kind in ("full", "param_groups"): + raise LegacyOptimizerCheckpointError("Legacy optimizer metadata requires explicit migration") + return None, None + if version != FORMAT_VERSION: + raise ValueError(f"Unsupported optimizer checkpoint version: {version}") + kind = metadata.get("kind") + if kind not in KINDS or (expected_kind is not None and kind != expected_kind): + raise ValueError(f"Unexpected optimizer checkpoint kind: {kind}") + raw_schema = metadata.get(SCHEMA_KEY) + if type(raw_schema) is not str or len(raw_schema.encode("utf-8")) > MAX_SCHEMA_BYTES: + raise ValueError("Missing or oversized optimizer schema") + try: + schema = json.loads(raw_schema, object_pairs_hook=_pairs, parse_int=_integer, parse_constant=_constant) + except (RecursionError, OverflowError) as exc: + raise ValueError("Optimizer schema exceeds JSON nesting or numeric limits") from exc + return kind, schema + + +def unflatten_optimizer(tensors, metadata, expected_kind=None, separator="."): + kind, schema = parse_optimizer_metadata(metadata, expected_kind) + if kind is None: + # Legacy tensor-only shards contain no serialized Python objects. + states = {} + for key in tensors: + parts = key.split(separator) + if len(parts) != 3 or parts[0] != "state" or len(parts[1]) > 4096 or not parts[1].isdigit() or not parts[2]: + raise LegacyOptimizerCheckpointError("Legacy optimizer layout requires explicit migration") + param_id = int(parts[1]) + if parts[2] in states.setdefault(param_id, {}): + raise ValueError("Duplicate legacy optimizer state field") + states[param_id][parts[2]] = tensors[key] + return states + count = [0] + referenced = set() + cache = {} + + def decode(node, depth=0): + _budget(depth, count) + if type(node) is not dict or type(node.get("type")) is not str: + raise ValueError("Optimizer schema node requires a type") + tag = node["type"] + if tag == "null": + _fields(node, ("type",)) + return None + if tag in ("bool", "int", "float", "str"): + _fields(node, ("type", "value")) + value = node["value"] + if tag == "float" and type(value) is str and value in ("nan", "+inf", "-inf"): + return {"nan": float("nan"), "+inf": float("inf"), "-inf": float("-inf")}[value] + expected = {"bool": bool, "int": int, "float": float, "str": str}[tag] + if type(value) is not expected or (tag == "float" and not math.isfinite(value)): + raise ValueError(f"Invalid optimizer {tag} value") + if tag == "int" and value.bit_length() > 4096: + raise ValueError("Optimizer schema integer is too large") + return value + if tag == "tensor_ref": + _fields(node, ("type", "key", "dtype", "shape")) + key = node["key"] + if type(key) is not str or not key.startswith("optimizer.tensor.") or key not in tensors: + raise ValueError("Invalid optimizer tensor reference") + shape = node["shape"] + if type(shape) is not list or any(type(dim) is not int or dim < 0 for dim in shape): + raise ValueError("Invalid optimizer tensor shape") + if key not in cache: + cache[key] = tensors[key] + tensor = cache[key] + if list(tensor.shape) != shape or str(tensor.dtype) != node["dtype"]: + raise ValueError("Optimizer tensor reference has the wrong shape or dtype") + referenced.add(key) + return tensor + if tag in ("list", "tuple", "dict"): + _fields(node, ("type", "items")) + if type(node["items"]) is not list: + raise ValueError("Optimizer container requires a list of items") + if tag != "dict": + items = [decode(item, depth + 1) for item in node["items"]] + return items if tag == "list" else tuple(items) + result = {} + for pair in node["items"]: + if type(pair) is not list or len(pair) != 2: + raise ValueError("Invalid optimizer dictionary entry") + key = decode(pair[0], depth + 1) + if type(key) not in (str, int) or key in result: + raise ValueError("Invalid or duplicate optimizer dictionary key") + result[key] = decode(pair[1], depth + 1) + return result + raise ValueError(f"Unsupported optimizer schema type: {tag}") + + payload = decode(schema) + if referenced != set(tensors): + raise ValueError("Optimizer checkpoint contains unreferenced tensors") + _validate_payload(payload, kind) + return payload["state"] if kind == "state_shard" else payload diff --git a/colossalai/utils/safetensors.py b/colossalai/utils/safetensors.py index b15b92fcc861..5f5a73f7beea 100644 --- a/colossalai/utils/safetensors.py +++ b/colossalai/utils/safetensors.py @@ -1,11 +1,14 @@ # a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214 import json import warnings +from collections.abc import Mapping from dataclasses import asdict, dataclass from typing import Dict, List, Optional, Tuple import torch -from safetensors.torch import _TYPES, load_file, safe_open +from safetensors.torch import _TYPES, safe_open, save_file + +from ._optimizer_safetensors import flatten_optimizer, unflatten_optimizer try: from tensornvme.async_file_io import AsyncFileWriter @@ -18,10 +21,6 @@ "Please install the latest tensornvme to use async save. pip install git+https://github.com/hpcaitech/TensorNVMe.git" ) _TYPES_INV = {v: k for k, v in _TYPES.items()} -import io - -from torch.distributed.distributed_c10d import _pickler, _unpickler - ASYNC_WRITE_ENTRIES = 32 @@ -31,23 +30,6 @@ def _require_async_file_writer(): return AsyncFileWriter -def _object_to_tensor(obj, device): - f = io.BytesIO() - _pickler(f).dump(obj) - byte_storage = torch.ByteStorage._from_buffer(f.getvalue()) # type: ignore[attr-defined] - # Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor and specifying dtype. - # Otherwise, it will casue 100X slowdown. - # See: https://github.com/pytorch/pytorch/issues/65696 - byte_tensor = torch.ByteTensor(byte_storage).to(device) - return byte_tensor - - -def _tensor_to_object(tensor, tensor_size): - tensor = tensor.cpu() - buf = tensor.numpy().tobytes()[:tensor_size] - return _unpickler(io.BytesIO(buf)).load() - - @dataclass class TensorInfo: dtype: str @@ -62,82 +44,25 @@ class PreparedData: offset: int -def _cast_to_tensor(obj): - if isinstance(obj, torch.Tensor): - return obj - return _object_to_tensor(obj, "cpu") - - -def _cast_to_object(tensor: torch.Tensor): - return _tensor_to_object(tensor, tensor.numel() * tensor.element_size()) - - -def _flatten_optim_state_dict(state_dict: dict, seperator: str = ".") -> Tuple[dict, Optional[dict]]: - flat_dict = {} - non_tensor_keys = [] - if "state" in state_dict: - # 3-level dict - states = state_dict["state"] - else: - # 2-level dict, usually for optimizer state dict shard - states = state_dict - - for idx, d in states.items(): - for k, v in d.items(): - if v is None: - continue - nested_key = f"state{seperator}{idx}{seperator}{k}" - if not isinstance(v, torch.Tensor): - non_tensor_keys.append(nested_key) - flat_dict[nested_key] = _cast_to_tensor(v) - if "param_groups" in state_dict: - flat_dict["param_groups"] = _cast_to_tensor(state_dict["param_groups"]) - non_tensor_keys.append("param_groups") - if len(non_tensor_keys) > 0: - metadata = {"non_tensor_keys": non_tensor_keys} - else: - metadata = None - return flat_dict, metadata - - -def _unflatten_optim_state_dict(flat_dict: dict, metadata: Optional[dict] = None, seperator: str = "."): - state_dict = {} - - if metadata is not None and "non_tensor_keys" in metadata: - non_tensor_keys = json.loads(metadata["non_tensor_keys"]) - else: - non_tensor_keys = [] - flat_dict = {k: _cast_to_object(v) if k in non_tensor_keys else v for k, v in flat_dict.items()} - if "param_groups" in flat_dict: - # 3-level dict - state_dict["param_groups"] = flat_dict.pop("param_groups") - state_dict["state"] = {} - states = state_dict["state"] - else: - # 2-level dict, usually for optimizer state dict shard - states = state_dict - - for k, v in flat_dict.items(): - parts = k.split(seperator) - assert len(parts) == 3 and parts[0] == "state" - idx = int(parts[1]) - key = parts[2] - if idx not in states: - states[idx] = {} - states[idx][key] = v - - return state_dict +def _flatten_optim_state_dict(state_dict: dict, seperator: str = ".") -> Tuple[dict, dict]: + """Encode optimizer data with the versioned, non-executable schema.""" + return flatten_optimizer(state_dict) + + +def _unflatten_optim_state_dict( + flat_dict: dict, metadata: Optional[dict] = None, seperator: str = ".", *, expected_kind=None +): + return unflatten_optimizer(flat_dict, metadata, expected_kind, seperator) def prepare( data: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None ) -> Tuple[PreparedData, List[torch.Tensor], List[str]]: - if metadata is not None: - assert isinstance(metadata, dict) - for k, v in metadata.items(): - metadata[k] = json.dumps(v) - assert isinstance(k, str) - assert isinstance(metadata[k], str) + # Safetensors metadata values are strings. Do not mutate cached metadata or + # JSON-encode already encoded optimizer schemas a second time. + metadata = {k: v if isinstance(v, str) else json.dumps(v) for k, v in (metadata or {}).items()} + if any(type(k) is not str for k in metadata): + raise TypeError("Safetensors metadata keys must be strings") tensors = [] tensor_keys = [] @@ -200,7 +125,8 @@ def move_and_save( f_writer.write(n.to_bytes(8, byteorder="little")) f_writer.write(header_bytes) - f_writer.register_h2d(len(tensor_keys)) + if tensor_keys: + f_writer.register_h2d(len(tensor_keys)) for name in tensor_keys: if state_dict_pinned: f_writer.write_tensor(state_dict[name], state_dict_pinned[name]) @@ -209,9 +135,50 @@ def move_and_save( return f_writer -def load_flat(checkpoint_path, seperator: str = "."): - with safe_open(checkpoint_path, framework="pt") as f: - metadata = f.metadata() - state_dict_load = load_file(checkpoint_path) - state_dict = _unflatten_optim_state_dict(state_dict_load, metadata, seperator) - return state_dict +class _FileTensors(Mapping): + def __init__(self, reader): + self.reader = reader + self.names = set(reader.keys()) + + def __getitem__(self, key): + if key not in self.names: + raise KeyError(key) + return self.reader.get_tensor(key) + + def __contains__(self, key): + return key in self.names + + def __iter__(self): + return iter(self.names) + + def __len__(self): + return len(self.names) + + +def load_flat(checkpoint_path, seperator: str = ".", *, expected_kind=None): + """Load data-only optimizer state; legacy embedded pickle is rejected.""" + with safe_open(checkpoint_path, framework="pt", device="cpu") as reader: + metadata = reader.metadata() + # The lazy mapping validates metadata before fetching tensor values; + # both metadata and values come from this one open file. + return _unflatten_optim_state_dict(_FileTensors(reader), metadata, seperator, expected_kind=expected_kind) + + +def save_optimizer_file(path, state_dict): + """Synchronously write the same data-only schema used by async optimizers.""" + import os + import tempfile + from pathlib import Path + + tensors, metadata = flatten_optimizer(state_dict) + path = Path(path) + # Private storage also supports aliased tensors with safetensors.save_file. + tensors = {key: tensor.cpu().clone().contiguous() for key, tensor in tensors.items()} + with tempfile.NamedTemporaryFile(dir=path.parent, prefix=".optimizer-", delete=False) as temporary: + temporary_path = temporary.name + try: + save_file(tensors, temporary_path, metadata=metadata) + os.replace(temporary_path, path) + finally: + if os.path.exists(temporary_path): + os.unlink(temporary_path) diff --git a/docs/source/en/basics/booster_checkpoint.md b/docs/source/en/basics/booster_checkpoint.md index ea6c11ae2cdc..89c321dd8d6e 100644 --- a/docs/source/en/basics/booster_checkpoint.md +++ b/docs/source/en/basics/booster_checkpoint.md @@ -53,6 +53,25 @@ Optimizer must be boosted by `colossalai.booster.Booster` before saving. Optimizer must be boosted by `colossalai.booster.Booster` before loading. +### Optimizer safetensors format v2 + +Async optimizer checkpoints use a versioned, data-only safetensors format. Tensor values remain in the tensor data area; parameter groups and nested values use a typed JSON schema. The format preserves lists, tuples, integers, explicit `None`, tensor learning rates and nested tensors. New sharded saves write parameter groups to `optimizer_group.safetensors`; index metadata identifies the group and shard formats. Tensor keys are stable across shards and repeated saves. + +The normal loader rejects legacy safetensors files with `non_tensor_keys`, including files produced by older ColossalAI versions. It never falls back to pickle. Tensor-only legacy shards can still be read, but they do not provide the parameter groups required for a complete resume. Existing synchronous PyTorch `.bin` checkpoints retain their format and are read with explicit `weights_only=True`; unsupported Python objects are rejected. + +Normal loading of PyTorch optimizer files requires torch 2.10.0 or later because earlier versions have known `weights_only` vulnerabilities. This check applies to full optimizer files, state shards and legacy parameter-group companions; it does not change the project's overall torch dependency or impose this legacy-file check on v2 safetensors. A version check does not certify checkpoint provenance. The explicit trusted migration tool remains available on older torch versions. + +To migrate a **trusted** legacy optimizer file or optimizer checkpoint directory, use a new destination: + +```bash +colossalai checkpoint migrate-optimizer --trust-legacy-pickle old_optimizer.safetensors new_optimizer.safetensors +colossalai checkpoint migrate-optimizer --trust-legacy-pickle old_optimizer_directory new_optimizer_directory +``` + +Legacy decoding can execute code. Run this opt-in tool only for trusted inputs in an isolated environment without network access or credentials. It is separate from training loaders, never overwrites an existing destination, and prints input/output SHA-256 hashes as JSON. Conversion is not a way to certify untrusted input. Arbitrary custom Python classes are not supported by v2. Old ColossalAI versions cannot read v2 checkpoints; retain the originals when planning a rollback. + +The schema has limits of 16 MiB of JSON, depth 64 and 1,000,000 nodes; integer values are limited to 4096 bits. Unsupported types and oversized schemas raise errors rather than silently dropping optimizer state. Plugin-specific optimizer and sharding constraints still apply. + ## LR Scheduler Checkpoint {{ autodoc:colossalai.booster.Booster.save_lr_scheduler }} diff --git a/docs/source/zh-Hans/basics/booster_checkpoint.md b/docs/source/zh-Hans/basics/booster_checkpoint.md index 1ff2e330521c..4f8893ed1d2c 100644 --- a/docs/source/zh-Hans/basics/booster_checkpoint.md +++ b/docs/source/zh-Hans/basics/booster_checkpoint.md @@ -54,6 +54,25 @@ booster.load(model, model_dir) 优化器在加载前必须被 `colossalai.booster.Booster` 封装。 +### 优化器 safetensors v2 格式 + +异步优化器 checkpoint 使用带版本的纯数据 safetensors 格式：tensor 保存在数据区，参数组及嵌套结构使用有类型 JSON 描述。它保留列表、元组、整数、显式 None、tensor 学习率和嵌套 tensor。新分片保存将参数组写到 `optimizer_group.safetensors`，索引记录参数组和 shard 的格式；tensor key 在分片与重复保存之间保持稳定。 + +默认加载器拒绝包含 `non_tensor_keys` 的旧 safetensors 优化器文件，包括旧版 ColossalAI 正常写出的文件，不会回退到 pickle。旧纯 tensor shard 仍可读取，但它本身没有完整恢复需要的参数组。已有同步 PyTorch `.bin` checkpoint 保留原格式，读取显式使用 `weights_only=True`；不支持的 Python 对象会被拒绝。 + +普通读取 PyTorch 优化器文件要求 torch 2.10.0 或更高版本，因为更早版本存在已知的 `weights_only` 漏洞。这项检查覆盖完整优化器、状态分片和旧参数组伴随文件，不提升整个项目的 torch 最低依赖，也不对 v2 safetensors 应用旧格式门槛。版本检查不能认证文件来源。旧 torch 环境仍可显式迁移可信旧文件。 + +来源可信的旧优化器文件或目录，可以显式迁移到一个不存在的新目标： + +```bash +colossalai checkpoint migrate-optimizer --trust-legacy-pickle old_optimizer.safetensors new_optimizer.safetensors +colossalai checkpoint migrate-optimizer --trust-legacy-pickle old_optimizer_directory new_optimizer_directory +``` + +旧格式解码仍可能执行代码。此独立工具只用于可信输入，应在无网络、无凭据的隔离环境运行；训练加载器不自动调用它。迁移不覆盖已有目标，完成后输出含输入 / 输出 SHA-256 的 JSON 摘要。迁移不能为不可信输入作安全认证。v2 不支持任意 Python 自定义类；旧版本程序无法读取 v2，回滚计划应保留原文件。 + +JSON schema 限制为 16 MiB、深度 64、1,000,000 个节点，整数值限制 4096 bit。不支持的类型或超限数据会明确报错，不静默丢弃优化器状态。各插件原有的优化器与分片约束仍然适用。 + ## 学习率调度器 Checkpoint {{ autodoc:colossalai.booster.Booster.save_lr_scheduler }} diff --git a/tests/test_checkpoint_io/test_optimizer_safetensors_plugins.py b/tests/test_checkpoint_io/test_optimizer_safetensors_plugins.py new file mode 100644 index 000000000000..c94285e70fd3 --- /dev/null +++ b/tests/test_checkpoint_io/test_optimizer_safetensors_plugins.py @@ -0,0 +1,200 @@ +"""Small real-plugin checkpoint regression, requiring two guarded CUDA devices.""" + +import copy +import tempfile +from contextlib import nullcontext + +import pytest +import torch +import torch.distributed as dist + +import colossalai +from colossalai._compat import is_transformers_v5 +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchFSDPPlugin +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.safetensors import HAS_TENSORNVME + + +def _equal(left, right): + if isinstance(left, torch.Tensor): + assert torch.equal(left, right.to(left.device)) + elif isinstance(left, dict): + assert left.keys() == right.keys(), (left.keys(), right.keys()) + for key in left: + _equal(left[key], right[key]) + elif isinstance(left, (list, tuple)): + assert type(left) is type(right) and len(left) == len(right) + for a, b in zip(left, right): + _equal(a, b) + else: + assert left == right + + +def _run(rank, world_size, port, plugin_name): + colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost", backend="nccl") + torch.cuda.set_device(rank) + torch.manual_seed(1234) + if plugin_name == "fsdp": + plugin = TorchFSDPPlugin() + elif plugin_name == "zero": + plugin = LowLevelZeroPlugin(stage=2, precision="bf16", initial_scale=1) + elif plugin_name == "gemini": + plugin = GeminiPlugin(precision="bf16", placement_policy="static", initial_scale=1) + else: + plugin = HybridParallelPlugin(tp_size=1, pp_size=1, zero_stage=1, precision="bf16", initial_scale=1) + booster = Booster(plugin=plugin) + if plugin_name == "hybrid": + from transformers import LlamaConfig, LlamaModel + + model = LlamaModel( + LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + attn_implementation="eager", + ) + ).cuda() + else: + model = torch.nn.Sequential(torch.nn.Linear(8, 16), torch.nn.GELU(), torch.nn.Linear(16, 4)).cuda() + optimizer = ( + HybridAdam(model.parameters(), lr=0.001) if plugin_name == "gemini" else torch.optim.Adam(model.parameters()) + ) + model, optimizer, *_ = booster.boost(model, optimizer) + if plugin_name == "hybrid": + output = model(input_ids=torch.randint(0, 64, (2, 8), device="cuda")).last_hidden_state + else: + dtype = torch.float32 if plugin_name == "fsdp" else torch.bfloat16 + output = model(torch.rand(2, 8, device="cuda", dtype=dtype)) + booster.backward(output.float().square().mean(), optimizer) + optimizer.step() + optimizer.zero_grad() + inner = optimizer.unwrap() if hasattr(optimizer, "unwrap") else optimizer.optim + + def snapshot(): + # Gemini's public checkpoint state adds PyTorch Adam defaults and + # gathers partitioned states; local HybridAdam internals differ. + state = optimizer.state_dict(only_rank_0=False) if plugin_name == "gemini" else inner.state_dict() + return copy.deepcopy(state) + + with tempfile.TemporaryDirectory() if rank == 0 else nullcontext(None) as directory: + directories = [directory] + dist.broadcast_object_list(directories, src=0) + for use_async in (False, True): + for shard in (False, True): + for low_mem in (True, False): + # Updating this changes schema size between saves and also + # verifies that parameter-group options survive each loader. + inner.param_groups[0]["checkpoint_note"] = "data" * (100 if low_mem else 1) + expected = snapshot() + assert expected["state"], "optimizer state must be nonempty" + path = directories[0] + f"/{use_async}_{shard}_{low_mem}" + if use_async and not shard: + path += ".safetensors" + booster.save_optimizer(optimizer, path, shard=shard, use_async=use_async, size_per_shard=1) + booster.checkpoint_io._sync_d2h() + booster.checkpoint_io._sync_io() + dist.barrier() + for state in inner.state.values(): + for key, value in state.items(): + if isinstance(value, torch.Tensor): + value.fill_(-12) + elif type(value) is int: + state[key] = -12 + booster.load_optimizer(optimizer, path, low_cpu_mem_mode=low_mem) + _equal(expected, snapshot()) + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.skipif(not HAS_TENSORNVME or torch.cuda.device_count() < 2, reason="Requires two guarded CUDA devices") +@pytest.mark.parametrize("plugin_name", ["fsdp", "zero", "gemini", "hybrid"]) +@rerun_if_address_is_in_use() +def test_optimizer_format_plugins(plugin_name): + spawn(_run, 2, plugin_name=plugin_name) + + +def _run_moe_migration(rank, world_size, port, config): + from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin + from colossalai.checkpoint_io.optimizer_migration import migrate_optimizer_checkpoint + from colossalai.checkpoint_io.utils import load_param_groups + from colossalai.testing.random import seed_all + from tests.test_shardformer.test_model.test_moe_optimizer_checkpoint import _state_snapshot + from tests.test_shardformer.test_model.test_moe_optimizer_coverage import NUM_TOKENS, _build + + colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost", backend="nccl") + torch.cuda.set_device(rank) + model_name, ep_size = config + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=1, + ep_size=ep_size, + zero_stage=1, + precision="bf16", + initial_scale=1, + overlap_communication=False, + find_unused_parameters=True, + ) + booster = Booster(plugin=plugin) + seed_all(1234) + model = _build(model_name).to(torch.bfloat16).cuda() + hidden = model.config.hidden_size + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + model, optimizer, *_ = booster.boost(model, optimizer) + seed_all(4321 + dist.get_rank(plugin.dp_group)) + inputs = torch.rand(2, NUM_TOKENS // 2, hidden, device="cuda", dtype=torch.bfloat16) + + def step(): + loss = model(inputs_embeds=inputs).last_hidden_state.float().square().mean() + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + step() + expected = _state_snapshot(optimizer) + assert expected, "the optimizer must have real state before migration" + with tempfile.TemporaryDirectory() if rank == 0 else nullcontext(None) as directory: + directories = [directory] + dist.broadcast_object_list(directories, src=0) + root = directories[0] + booster.save_optimizer(optimizer, root + "/old", shard=True) + dist.barrier() + if rank == 0: + import json + from pathlib import Path + + old = Path(root) / "old" + index_path = next(old.glob("*.index.json")) + index = json.loads(index_path.read_text()) + # The state shards already use the legacy PyTorch layout. Recreate + # the old groups companion so both parts exercise real migration. + groups = load_param_groups(old / index["metadata"]["param_groups"]) + torch.save(groups, old / "legacy_groups.bin") + index["metadata"]["param_groups"] = "legacy_groups.bin" + index["metadata"].pop("param_groups_format", None) + index_path.write_text(json.dumps(index)) + migrate_optimizer_checkpoint(old, Path(root) / "new", trust_legacy_pickle=True) + dist.barrier() + inner = optimizer.unwrap() if hasattr(optimizer, "unwrap") else optimizer.optim + for state in inner.state.values(): + for value in state.values(): + if isinstance(value, torch.Tensor): + value.fill_(-12) + booster.load_optimizer(optimizer, root + "/new") + _equal(expected, _state_snapshot(optimizer)) + dist.barrier() + # Exercise a real resumed training step after the exact-state check. + step() + dist.destroy_process_group() + + +@pytest.mark.skipif(not is_transformers_v5(), reason="Fused MoE experts only exist in transformers v5") +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two guarded CUDA devices") +@pytest.mark.parametrize("config", [("mixtral", 2), ("deepseek_v3", 2), ("mixtral_odd", 1)]) +@rerun_if_address_is_in_use() +def test_moe_migrated_optimizer(config): + spawn(_run_moe_migration, 2, config=config) diff --git a/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py b/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py new file mode 100644 index 000000000000..c5a11791e1b9 --- /dev/null +++ b/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py @@ -0,0 +1,756 @@ +import json +import math +import pickle + +import pytest +import torch +from packaging.version import Version +from safetensors.torch import save_file + +from colossalai.checkpoint_io import GeneralCheckpointIO +from colossalai.checkpoint_io.index_file import CheckpointIndexFile +from colossalai.checkpoint_io.optimizer_migration import migrate_optimizer_checkpoint +from colossalai.checkpoint_io.utils import ( + load_optimizer_state, + load_param_groups, + load_state_dict_shards, + prepare_pinned_state_dict, +) +from colossalai.utils import _optimizer_safetensors as schema +from colossalai.utils._optimizer_safetensors import LegacyOptimizerCheckpointError +from colossalai.utils.safetensors import ( + HAS_TENSORNVME, + _flatten_optim_state_dict, + load_flat, + prepare, + save_nested, + save_optimizer_file, +) + + +def assert_equal(a, b): + assert type(a) is type(b) + if isinstance(a, torch.Tensor): + assert torch.equal(a, b) + elif isinstance(a, dict): + assert a.keys() == b.keys() + for key in a: + assert_equal(a[key], b[key]) + elif isinstance(a, (tuple, list)): + assert len(a) == len(b) + for x, y in zip(a, b): + assert_equal(x, y) + elif isinstance(a, float) and math.isnan(a): + assert math.isnan(b) + else: + assert a == b + + +def adam(): + params = [torch.nn.Parameter(torch.ones(2, 2)), torch.nn.Parameter(torch.ones(2))] + opt = torch.optim.Adam(params) + for param in params: + param.grad = torch.ones_like(param) * 0.1 + opt.step() + return opt + + +def write_v2(path, value): + flat, metadata = _flatten_optim_state_dict(value) + save_file(flat, str(path), metadata=metadata) + + +@pytest.mark.parametrize("kind", ["full", "state_shard", "param_groups"]) +def test_typed_roundtrip(tmp_path, kind): + state = adam().state_dict() + state["state"][0]["nested"] = [ + None, + True, + 4, + "text", + (0.5, torch.arange(3)), + {1: -0.0}, + ] + state["state"][0]["empty"] = torch.empty(0, dtype=torch.bfloat16) + state["param_groups"][0]["lr"] = torch.tensor(0.001) + state["param_groups"][0]["special"] = (float("inf"), float("-inf"), float("nan")) + value = ( + state + if kind == "full" + else (state["state"] if kind == "state_shard" else {"param_groups": state["param_groups"]}) + ) + path = tmp_path / "optimizer.safetensors" + write_v2(path, value) + assert_equal(value, load_flat(path, expected_kind=kind)) + + +@pytest.mark.parametrize( + "name", + [ + "ASGD", + "Adadelta", + "Adafactor", + "Adagrad", + "Adam", + "AdamW", + "Adamax", + "LBFGS", + "Muon", + "NAdam", + "RAdam", + "RMSprop", + "Rprop", + "SGD", + "SparseAdam", + "colossalai.CAME", + "colossalai.Lamb", + "colossalai.Lars", + "colossalai.Adafactor", + ], +) +def test_optimizer_next_step(tmp_path, name): + p = torch.nn.Parameter(torch.ones(2, 2)) + if name.startswith("colossalai."): + import colossalai.nn.optimizer as optimizers + + cls = getattr(optimizers, name.split(".")[1]) + else: + cls = getattr(torch.optim, name, None) + if cls is None: + pytest.skip(f"{name} is unavailable in this torch version") + kwargs = {"lr": 0.01} if name in ("Muon", "colossalai.CAME") else {} + opt = cls([p], **kwargs) + + def step(optimizer, param): + if name == "LBFGS": + + def closure(): + optimizer.zero_grad() + loss = param.square().sum() + loss.backward() + return loss + + optimizer.step(closure) + else: + param.grad = torch.full_like(param, 0.1) + if name == "SparseAdam": + param.grad = param.grad.to_sparse() + optimizer.step() + + step(opt, p) + q = torch.nn.Parameter(p.detach().clone()) + restored = cls([q], **kwargs) + path = tmp_path / "optimizer.safetensors" + save_optimizer_file(path, opt.state_dict()) + GeneralCheckpointIO().load_optimizer(restored, str(path)) + assert_equal(opt.state_dict(), restored.state_dict()) + step(opt, p) + step(restored, q) + assert torch.equal(p, q) + assert_equal(opt.state_dict(), restored.state_dict()) + + +def test_prepare_does_not_double_encode_or_mutate(): + tensors, metadata = _flatten_optim_state_dict(adam().state_dict()) + original = dict(metadata) + prepared, _, _ = prepare(tensors, metadata) + assert metadata == original + header = json.loads(prepared.header_bytes) + assert header["__metadata__"][schema.FORMAT_KEY] == "2" + assert header["__metadata__"][schema.SCHEMA_KEY] == original[schema.SCHEMA_KEY] + + +@pytest.mark.parametrize("location", ["state.0.candidate", "param_groups"]) +def test_legacy_payload_rejected_before_execution(tmp_path, location): + marker = tmp_path / "must_not_exist" + + class Payload: + def __reduce__(self): + return eval, (f"__import__('pathlib').Path({str(marker)!r}).write_text('executed')",) + + tensors = {location: torch.tensor(list(pickle.dumps(Payload(), protocol=2)), dtype=torch.uint8)} + path = tmp_path / "legacy.safetensors" + save_file(tensors, str(path), metadata={"non_tensor_keys": json.dumps([location])}) + for loader in ( + load_flat, + lambda file: GeneralCheckpointIO().load_optimizer(adam(), str(file)), + ): + with pytest.raises(LegacyOptimizerCheckpointError, match="migrat|Migrate"): + loader(path) + assert not marker.exists() + + +def test_torch_optimizer_load_is_explicitly_restricted(tmp_path): + marker = tmp_path / "must_not_exist" + + class Payload: + def __reduce__(self): + return eval, (f"__import__('pathlib').Path({str(marker)!r}).write_text('executed')",) + + path = tmp_path / "optimizer.bin" + torch.save(Payload(), path) + error = RuntimeError if Version(torch.__version__) < Version("2.10.0") else pickle.UnpicklingError + with pytest.raises(error): + load_optimizer_state(path) + assert not marker.exists() + + +def optimizer_file_case(kind): + from colossalai.checkpoint_io.utils import load_optim_or_model_shard + + state = adam().state_dict() + if kind == "full": + return state, load_optimizer_state + if kind == "param_groups": + return state["param_groups"], load_param_groups + return state["state"], lambda path: load_optim_or_model_shard(str(path), True, False) + + +@pytest.mark.parametrize("kind", ["full", "param_groups", "state_shard"]) +@pytest.mark.parametrize("version", ["2.5.1", "2.9.1", "2.10.0rc1"]) +def test_legacy_torch_version_rejected_before_deserialization(tmp_path, monkeypatch, kind, version): + payload, reader = optimizer_file_case(kind) + path = tmp_path / "legacy.bin" + torch.save(payload, path) + + def must_not_load(*args, **kwargs): + pytest.fail("Unsupported torch must be rejected before torch.load") + + monkeypatch.setattr(torch, "__version__", version) + monkeypatch.setattr(torch, "load", must_not_load) + with pytest.raises(RuntimeError, match="requires torch >= 2.10.0"): + reader(path) + + +@pytest.mark.parametrize("kind", ["full", "param_groups", "state_shard"]) +@pytest.mark.parametrize("version", ["2.10.0", "2.11.0+cu128"]) +def test_legacy_torch_version_policy_accepts_patched_versions(tmp_path, monkeypatch, kind, version): + # This checks version dispatch on the installed torch, not compatibility + # with an actual installation of each version string. + payload, reader = optimizer_file_case(kind) + path = tmp_path / "legacy.bin" + torch.save(payload, path) + monkeypatch.setattr(torch, "__version__", version) + assert_equal(payload, reader(path)) + + +@pytest.mark.parametrize("kind", ["full", "param_groups", "state_shard"]) +def test_v2_does_not_use_legacy_torch_version_gate(tmp_path, monkeypatch, kind): + payload, reader = optimizer_file_case(kind) + path = tmp_path / "optimizer.safetensors" + save_optimizer_file(path, {"param_groups": payload} if kind == "param_groups" else payload) + monkeypatch.setattr(torch, "__version__", "2.5.1") + assert_equal(payload, reader(path)) + + +def test_legacy_tensor_only_shard(tmp_path): + path = tmp_path / "legacy.safetensors" + save_file({"state.0.step": torch.tensor(1.0)}, str(path)) + assert_equal({0: {"step": torch.tensor(1.0)}}, load_flat(path, expected_kind="state_shard")) + with pytest.raises(LegacyOptimizerCheckpointError): + load_flat(path, expected_kind="full") + + +@pytest.mark.parametrize( + "mutation", + [ + "version", + "kind", + "legacy", + "unknown", + "duplicate_json", + "missing_tensor", + "extra_tensor", + "shape", + "dtype", + "bool_id", + ], +) +def test_invalid_schema_rejected(tmp_path, mutation): + tensors, metadata = _flatten_optim_state_dict(adam().state_dict()) + tree = json.loads(metadata[schema.SCHEMA_KEY]) + + def nodes(value): + if isinstance(value, dict): + yield value + for child in value.values(): + yield from nodes(child) + elif isinstance(value, list): + for child in value: + yield from nodes(child) + + if mutation == "version": + metadata[schema.FORMAT_KEY] = "3" + elif mutation == "kind": + metadata["kind"] = "param_groups" + elif mutation == "legacy": + metadata["non_tensor_keys"] = "[]" + elif mutation == "unknown": + tree["type"] = "python_class" + elif mutation == "duplicate_json": + metadata[schema.SCHEMA_KEY] = '{"type":"null","type":"dict"}' + elif mutation == "missing_tensor": + tensors.pop(next(iter(tensors))) + elif mutation == "extra_tensor": + tensors["unreferenced"] = torch.ones(1) + elif mutation == "shape": + next(n for n in nodes(tree) if n.get("type") == "tensor_ref")["shape"] = [100] + elif mutation == "dtype": + next(n for n in nodes(tree) if n.get("type") == "tensor_ref")["dtype"] = "torch.int64" + elif mutation == "bool_id": + n = next(n for n in nodes(tree) if n.get("type") == "int") + n["value"] = True + if mutation != "duplicate_json": + metadata[schema.SCHEMA_KEY] = json.dumps(tree) + path = tmp_path / "invalid.safetensors" + save_file(tensors, str(path), metadata=metadata) + with pytest.raises(ValueError): + load_flat(path, expected_kind="full") + + +def test_reader_schema_byte_boundary_before_tensor_fetch(tmp_path, monkeypatch): + from colossalai.utils.safetensors import _FileTensors + + state = adam().state_dict() + state["state"][0]["note"] = "中文" + tensors, metadata = _flatten_optim_state_dict(state) + metadata[schema.SCHEMA_KEY] = json.dumps(json.loads(metadata[schema.SCHEMA_KEY]), ensure_ascii=False) + size = len(metadata[schema.SCHEMA_KEY].encode("utf-8")) + assert size > len(metadata[schema.SCHEMA_KEY]) + path = tmp_path / "byte_boundary.safetensors" + save_file(tensors, str(path), metadata=metadata) + monkeypatch.setattr(schema, "MAX_SCHEMA_BYTES", size) + assert_equal(state, load_flat(path, expected_kind="full")) + + def must_not_fetch(self, key): + pytest.fail("Oversized metadata must be rejected before fetching tensors") + + monkeypatch.setattr(schema, "MAX_SCHEMA_BYTES", size - 1) + monkeypatch.setattr(_FileTensors, "__getitem__", must_not_fetch) + with pytest.raises(ValueError, match="oversized"): + load_flat(path, expected_kind="full") + + +@pytest.mark.parametrize("limit,boundary", [("MAX_DEPTH", 6), ("MAX_NODES", 12)]) +def test_reader_container_budget_boundary(tmp_path, monkeypatch, limit, boundary): + # Root -> groups list -> group dict -> three nested lists -> int: depth 6. + # Including dictionary keys and params=[0], this payload has 12 typed nodes. + state = {"param_groups": [{"params": [0], "extra": [[[1]]]}]} + path = tmp_path / "container_boundary.safetensors" + save_optimizer_file(path, state) + monkeypatch.setattr(schema, limit, boundary) + assert_equal(state, load_flat(path, expected_kind="param_groups")) + monkeypatch.setattr(schema, limit, boundary - 1) + with pytest.raises(ValueError, match="depth or node limit"): + load_flat(path, expected_kind="param_groups") + + +def test_budget_and_unsupported_writer_inputs(monkeypatch): + state = adam().state_dict() + state["state"][0]["object"] = object() + with pytest.raises(TypeError, match="Unsupported"): + _flatten_optim_state_dict(state) + state["state"][0].pop("object") + state["state"][0]["cycle"] = state + with pytest.raises(ValueError, match="Cyclic"): + _flatten_optim_state_dict(state) + state["state"][0].pop("cycle") + monkeypatch.setattr(schema, "MAX_NODES", 2) + with pytest.raises(ValueError, match="node limit"): + _flatten_optim_state_dict(state) + + +@pytest.mark.parametrize("prefetch", [False, True]) +def test_shards_groups_and_prefetch(tmp_path, prefetch): + state = adam().state_dict() + paths = [] + keys = [] + for param_id, values in state["state"].items(): + path = tmp_path / f"{param_id}.safetensors" + shard = {param_id: values} + flat, _ = _flatten_optim_state_dict(shard) + keys.append(set(flat)) + write_v2(path, shard) + paths.append(str(path)) + assert keys[0].isdisjoint(keys[1]) + loaded = {} + for shard in load_state_dict_shards(paths, True, True, low_cpu_mem_mode=not prefetch): + loaded.update(shard) + assert_equal(state["state"], loaded) + groups = tmp_path / "groups.safetensors" + save_optimizer_file(groups, {"param_groups": state["param_groups"]}) + assert_equal(state["param_groups"], load_param_groups(groups)) + + +def test_pinned_cache_accepts_new_keys_shapes_and_dtypes(monkeypatch): + from colossalai.checkpoint_io import utils + + # Buffer selection is a CPU unit; real pinned transfers run in GPU integration. + monkeypatch.setattr(utils, "_pin_tensor", lambda tensor: torch.empty_like(tensor)) + cache = {} + first = prepare_pinned_state_dict({"a": torch.ones(1)}, cache) + second = prepare_pinned_state_dict({"a": torch.ones(1), "b": torch.ones(2)}, cache) + assert first["a"] is second["a"] + third = prepare_pinned_state_dict({"a": torch.ones(3, dtype=torch.int64)}, cache) + assert third["a"].shape == (3,) and third["a"].dtype == torch.int64 + + +@pytest.mark.parametrize("escape", ["../outside.bin", "/outside.bin"]) +def test_index_path_escape_rejected(tmp_path, escape): + path = tmp_path / "optimizer.index.json" + path.write_text(json.dumps({"metadata": {}, "weight_map": {"0": escape}})) + with pytest.raises(ValueError, match="root|relative"): + CheckpointIndexFile.from_file(path, is_optimizer=True) + + +def test_model_snapshot_links_to_external_blobs(tmp_path): + snapshot = tmp_path / "snapshots" / "revision" + snapshot.mkdir(parents=True) + blobs = tmp_path / "blobs" + blobs.mkdir() + expected = torch.nn.Linear(2, 2).state_dict() + blob = blobs / "opaque_hash" + save_file(expected, str(blob)) + (snapshot / "model.safetensors").symlink_to(blob) + (snapshot / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": {name: "model.safetensors" for name in expected}, + } + ) + ) + restored = torch.nn.Linear(2, 2) + GeneralCheckpointIO().load_model(restored, str(snapshot), strict=True) + assert_equal(expected, restored.state_dict()) + + +@pytest.mark.parametrize("location", ["param_groups", "weight_map"]) +@pytest.mark.parametrize("reference", ["absolute", "parent", "symlink"]) +def test_optimizer_api_rejects_external_index_paths(tmp_path, location, reference): + directory = tmp_path / "checkpoint" + directory.mkdir() + outside = tmp_path / "outside.safetensors" + outside.write_bytes(b"never decoded") + if reference == "absolute": + name = str(outside) + elif reference == "parent": + name = "../outside.safetensors" + else: + (directory / "linked.safetensors").symlink_to(outside) + name = "linked.safetensors" + metadata = {"param_groups": name if location == "param_groups" else "groups.safetensors"} + weights = {"0": name} if location == "weight_map" else {} + (directory / "optimizer.index.json").write_text(json.dumps({"metadata": metadata, "weight_map": weights})) + with pytest.raises(ValueError, match="relative|root"): + GeneralCheckpointIO().load_optimizer(adam(), str(directory)) + + +def test_optimizer_index_preserves_internal_link_names(tmp_path, cpu_sync): + expected = adam().state_dict() + save_optimizer_file(tmp_path / "groups_blob", {"param_groups": expected["param_groups"]}) + save_optimizer_file(tmp_path / "states_blob", expected["state"]) + (tmp_path / "groups.safetensors").symlink_to(tmp_path / "groups_blob") + (tmp_path / "states.safetensors").symlink_to(tmp_path / "states_blob") + index = tmp_path / "optimizer.index.json" + index.write_text( + json.dumps( + { + "metadata": {"param_groups": "groups.safetensors"}, + "weight_map": {str(key): "states.safetensors" for key in expected["state"]}, + } + ) + ) + loaded_index = CheckpointIndexFile.from_file(index) + assert loaded_index.get_param_group_filename().endswith("groups.safetensors") + assert loaded_index.get_checkpoint_filenames()[0][0].endswith("states.safetensors") + restored = adam() + GeneralCheckpointIO().load_optimizer(restored, str(tmp_path)) + assert_equal(expected, restored.state_dict()) + + +@pytest.fixture +def cpu_sync(monkeypatch): + from colossalai.accelerator import get_accelerator + + accelerator = get_accelerator() + if accelerator.name == "cpu": + monkeypatch.setattr(accelerator, "synchronize", lambda: None) + + +@pytest.mark.skipif( + Version(torch.__version__) < Version("2.10.0"), reason="Legacy optimizer loading requires torch >= 2.10" +) +def test_general_sync_checkpoint_next_step(tmp_path, cpu_sync): + original = adam() + restored = adam() + io = GeneralCheckpointIO() + io.save_optimizer(original, str(tmp_path), shard=True, size_per_shard=1) + io.load_optimizer(restored, str(tmp_path)) + assert_equal(original.state_dict(), restored.state_dict()) + + +def test_migration_requires_explicit_trust_and_preserves_input(tmp_path): + state = adam().state_dict() + flat = {f"state.{pid}.{key}": tensor for pid, values in state["state"].items() for key, tensor in values.items()} + flat["param_groups"] = torch.tensor(list(pickle.dumps(state["param_groups"])), dtype=torch.uint8) + source = tmp_path / "legacy.safetensors" + output = tmp_path / "new.safetensors" + save_file(flat, str(source), metadata={"non_tensor_keys": '["param_groups"]'}) + original_bytes = source.read_bytes() + with pytest.raises(ValueError, match="Explicit"): + migrate_optimizer_checkpoint(source, output) + assert not output.exists() + report = migrate_optimizer_checkpoint(source, output, trust_legacy_pickle=True) + assert report["format_version"] == 2 and source.name in report["inputs"] + assert output.name in report["outputs"] + assert_equal(state, load_flat(output)) + assert source.read_bytes() == original_bytes + with pytest.raises(FileExistsError): + migrate_optimizer_checkpoint(source, output, trust_legacy_pickle=True) + + +def test_sharded_migration(tmp_path, cpu_sync): + state = adam().state_dict() + source = tmp_path / "old" + source.mkdir() + torch.save(state["param_groups"], source / "groups.bin") + torch.save(state["state"], source / "states.bin") + (source / "optimizer.index.json").write_text( + json.dumps( + { + "metadata": {"param_groups": "groups.bin"}, + "weight_map": {str(key): "states.bin" for key in state["state"]}, + } + ) + ) + output = tmp_path / "new" + migrate_optimizer_checkpoint(source, output, trust_legacy_pickle=True) + restored = adam() + GeneralCheckpointIO().load_optimizer(restored, str(output)) + assert_equal(state, restored.state_dict()) + + +@pytest.mark.parametrize("suffix", [".bin", ".safetensors"]) +def test_migrate_legacy_nested_tensors(tmp_path, suffix): + state = adam().state_dict() + state["param_groups"][0]["lr"] = torch.tensor(0.001) + state["state"][0]["history"] = [torch.arange(3), (None, torch.tensor(0.5))] + source = tmp_path / ("legacy" + suffix) + if suffix == ".bin": + torch.save(state, source) + else: + flat, objects = {}, [] + values = { + f"state.{pid}.{key}": value for pid, fields in state["state"].items() for key, value in fields.items() + } + values["param_groups"] = state["param_groups"] + for key, value in values.items(): + if not isinstance(value, torch.Tensor): + objects.append(key) + value = torch.tensor(list(pickle.dumps(value)), dtype=torch.uint8) + flat[key] = value + save_file(flat, str(source), metadata={"non_tensor_keys": json.dumps(objects)}) + output = tmp_path / "migrated.safetensors" + migrate_optimizer_checkpoint(source, output, trust_legacy_pickle=True) + assert_equal(state, load_flat(output, expected_kind="full")) + + +@pytest.mark.parametrize("suffix", [".bin", ""]) +def test_migration_rejects_destination_suffix_before_decoding(tmp_path, monkeypatch, suffix): + from colossalai.checkpoint_io import optimizer_migration + + source = tmp_path / "old.bin" + torch.save(adam().state_dict(), source) + destination = tmp_path / ("new" + suffix) + + def must_not_decode(path): + pytest.fail("Invalid destinations must be rejected before legacy decoding") + + monkeypatch.setattr(optimizer_migration, "_load_trusted_legacy", must_not_decode) + with pytest.raises(ValueError, match="destination must end with .safetensors"): + migrate_optimizer_checkpoint(source, destination, trust_legacy_pickle=True) + assert not destination.exists() + + +def test_moe_loads_migrated_shards_at_format_boundary(tmp_path, cpu_sync): + from types import SimpleNamespace + + from colossalai.checkpoint_io.moe_checkpoint import MoECheckpointIO + from colossalai.interface import OptimizerWrapper + + original = adam().state_dict() + source = tmp_path / "old" + source.mkdir() + torch.save(original["param_groups"], source / "groups.bin") + torch.save(original["state"], source / "states.bin") + (source / "optimizer.index.json").write_text( + json.dumps( + { + "metadata": {"param_groups": "groups.bin"}, + "weight_map": {"0": "states.bin", "1": "states.bin"}, + } + ) + ) + destination = tmp_path / "new" + migrate_optimizer_checkpoint(source, destination, trust_legacy_pickle=True) + opt = adam() + wrapped = OptimizerWrapper(opt) + params = [param for group in opt.param_groups for param in group["params"]] + wrapped.get_master_to_working_map = lambda: None + wrapped.param_info = { + "param2id": {id(param): index for index, param in enumerate(params)}, + "param2shape": {id(param): param.shape for param in params}, + } + io = object.__new__(MoECheckpointIO) + GeneralCheckpointIO.__init__(io) + io.coordinator = SimpleNamespace(is_master=lambda: False) + io.verbose = False + # Isolate file dispatch here; expert redistribution is tested on CUDA. + io.shard_from_complete_optimizer_state = lambda state, **kwargs: state + io.load_optimizer(wrapped, str(destination)) + assert_equal(original, opt.state_dict()) + + +@pytest.mark.parametrize( + "entry", + [ + "general", + "low_level_zero", + "fsdp", + "hybrid_full", + "hybrid_shard", + "shard_helper", + ], +) +def test_all_direct_loaders_reject_legacy_objects(tmp_path, entry): + from types import SimpleNamespace + + from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroCheckpointIO + from colossalai.booster.plugin.torch_fsdp_plugin import FSDPOptimizerWrapper, TorchFSDPCheckpointIO + from colossalai.checkpoint_io.hybrid_parallel_checkpoint_io import HybridParallelCheckpointIO + from colossalai.checkpoint_io.utils import load_optim_or_model_shard + from colossalai.interface import OptimizerWrapper + + path = tmp_path / "legacy.safetensors" + save_file( + {"state.0.step": torch.zeros(1, dtype=torch.uint8)}, + str(path), + metadata={"non_tensor_keys": '["state.0.step"]'}, + ) + opt = adam() + # Only exercise each real loader's format-rejection boundary here. Plugin + # training and state distribution are covered by GPU integration tests. + if entry == "general": + call = lambda: GeneralCheckpointIO().load_unsharded_optimizer(opt, str(path)) + elif entry == "low_level_zero": + io = object.__new__(LowLevelZeroCheckpointIO) + GeneralCheckpointIO.__init__(io) + call = lambda: io.load_unsharded_optimizer(opt, str(path)) + elif entry == "fsdp": + io = object.__new__(TorchFSDPCheckpointIO) + GeneralCheckpointIO.__init__(io) + wrapped = object.__new__(FSDPOptimizerWrapper) + call = lambda: io.load_unsharded_optimizer(wrapped, str(path)) + elif entry == "shard_helper": + call = lambda: load_optim_or_model_shard(str(path), True, True) + else: + io = object.__new__(HybridParallelCheckpointIO) + GeneralCheckpointIO.__init__(io) + io.coordinator = SimpleNamespace(is_master=lambda: False) + wrapped = OptimizerWrapper(opt) + wrapped.get_master_to_working_map = lambda: None + if entry == "hybrid_full": + call = lambda: io.load_unsharded_optimizer(wrapped, str(path)) + else: + params = [p for group in opt.param_groups for p in group["params"]] + wrapped.param_info = {"param2id": {id(p): i for i, p in enumerate(params)}} + save_optimizer_file( + tmp_path / "groups.safetensors", + {"param_groups": opt.state_dict()["param_groups"]}, + ) + index = tmp_path / "optimizer.index.json" + index.write_text( + json.dumps( + { + "metadata": {"param_groups": "groups.safetensors"}, + "weight_map": {"0": path.name}, + } + ) + ) + call = lambda: io.load_sharded_optimizer(wrapped, str(index), "") + with pytest.raises(LegacyOptimizerCheckpointError): + call() + + +def test_migration_cli_requires_flag(tmp_path): + from click.testing import CliRunner + + from colossalai.cli.cli import cli + + source = tmp_path / "old.bin" + expected = adam().state_dict() + torch.save(expected, source) + destination = tmp_path / "new.safetensors" + runner = CliRunner() + missing = runner.invoke(cli, ["checkpoint", "migrate-optimizer", str(source), str(destination)]) + assert missing.exit_code != 0 and not destination.exists() + result = runner.invoke( + cli, + [ + "checkpoint", + "migrate-optimizer", + "--trust-legacy-pickle", + str(source), + str(destination), + ], + ) + assert result.exit_code == 0, result.output + assert_equal(expected, load_flat(destination)) + + +@pytest.mark.skipif( + not HAS_TENSORNVME or not torch.cuda.is_available(), + reason="TensorNVMe requires a guarded CUDA device", +) +def test_real_async_writer(tmp_path): + state = adam().state_dict() + path = tmp_path / "optimizer.safetensors" + writer = save_nested(str(path), state) + writer.sync_before_step() + writer.synchronize() + assert_equal(state, load_flat(path)) + + +@pytest.mark.skipif( + not HAS_TENSORNVME or not torch.cuda.is_available(), + reason="Requires a guarded CUDA device", +) +@pytest.mark.parametrize("shard", [False, True]) +def test_general_repeated_async_save(tmp_path, shard): + def build(): + return torch.optim.Adam( + [ + {"params": [torch.nn.Parameter(torch.ones(2, 2))], "lr": 0.01}, + {"params": [torch.nn.Parameter(torch.ones(2))], "lr": 0.02}, + ] + ) + + optimizer = build() + io = GeneralCheckpointIO() + path = tmp_path / ("shards" if shard else "full.safetensors") + # Save before state is initialized, then after initialization, then with a + # longer parameter-group field. Pinned buffers must not encode object sizes. + for step in range(3): + if step: + for group in optimizer.param_groups: + for param in group["params"]: + param.grad = torch.ones_like(param) + optimizer.step() + optimizer.param_groups[0]["note"] = "updated" * (step * 100) + io.save_optimizer(optimizer, str(path), shard=shard, use_async=True, size_per_shard=0.00001) + io._sync_d2h() + io._sync_io() + restored = build() + io.load_optimizer(restored, str(path)) + assert_equal(optimizer.state_dict(), restored.state_dict())