From e57b0dec381b42b470ed8116990d6e896b5e7726 Mon Sep 17 00:00:00 2001 From: refinedheart <3102024448@qq.com> Date: Sun, 25 Oct 2026 03:21:40 +0800 Subject: [PATCH 1/2] [checkpoint] write data-only optimizer safetensors v2 Replace file-selected pickle decoding with a versioned typed JSON schema and native tensor references. Reject legacy embedded objects in training loaders and provide an explicit migration tool for trusted checkpoints. Use v2 parameter groups across checkpoint plugins, validate optimizer index paths without restricting model blob caches, and preserve pinned buffers across repeated saves. Add format, migration and plugin regressions and collect the CPU security tests in both CI entry points. --- .github/scripts/ci/run_colossalai_batch.sh | 1 + .github/scripts/ci/run_cpu_unit_tests.sh | 1 + colossalai/booster/plugin/gemini_plugin.py | 15 +- .../booster/plugin/low_level_zero_plugin.py | 9 +- .../booster/plugin/torch_fsdp_plugin.py | 21 +- .../checkpoint_io/general_checkpoint_io.py | 14 +- .../hybrid_parallel_checkpoint_io.py | 41 +- colossalai/checkpoint_io/index_file.py | 48 +- colossalai/checkpoint_io/moe_checkpoint.py | 19 +- .../checkpoint_io/optimizer_migration.py | 131 ++++ colossalai/checkpoint_io/utils.py | 52 +- colossalai/cli/checkpoint.py | 29 + colossalai/cli/cli.py | 2 + colossalai/utils/_optimizer_safetensors.py | 251 +++++++ colossalai/utils/safetensors.py | 167 ++--- docs/source/en/basics/booster_checkpoint.md | 17 + .../zh-Hans/basics/booster_checkpoint.md | 17 + .../test_optimizer_safetensors_plugins.py | 200 +++++ .../test_optimizer_safetensors_v2.py | 695 ++++++++++++++++++ 19 files changed, 1565 insertions(+), 165 deletions(-) create mode 100644 colossalai/checkpoint_io/optimizer_migration.py create mode 100644 colossalai/cli/checkpoint.py create mode 100644 colossalai/utils/_optimizer_safetensors.py create mode 100644 tests/test_checkpoint_io/test_optimizer_safetensors_plugins.py create mode 100644 tests/test_checkpoint_io/test_optimizer_safetensors_v2.py 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..45469ddb0032 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,19 @@ 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_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 torch.load(group_file_path, map_location="cpu", weights_only=True) + + +def load_optimizer_state(checkpoint): + if str(checkpoint).endswith(".safetensors"): + return load_flat(checkpoint, expected_kind="full") + return torch.load(checkpoint, map_location="cpu", weights_only=True) def merge_optimizer_param_group_defaults(saved_groups: List[dict], current_groups: List[dict]) -> None: @@ -706,7 +720,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 +783,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 +951,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,10 +1027,24 @@ 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) else: 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..b1ba24145705 100644 --- a/docs/source/en/basics/booster_checkpoint.md +++ b/docs/source/en/basics/booster_checkpoint.md @@ -53,6 +53,23 @@ 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. + +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..96a7f64c96d6 100644 --- a/docs/source/zh-Hans/basics/booster_checkpoint.md +++ b/docs/source/zh-Hans/basics/booster_checkpoint.md @@ -54,6 +54,23 @@ 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 对象会被拒绝。 + +来源可信的旧优化器文件或目录，可以显式迁移到一个不存在的新目标： + +```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..3148692a0f4d --- /dev/null +++ b/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py @@ -0,0 +1,695 @@ +import json +import math +import pickle + +import pytest +import torch +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) + with pytest.raises(pickle.UnpicklingError): + load_optimizer_state(path) + assert not marker.exists() + + +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) + + +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" + torch.save(adam().state_dict(), 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(load_optimizer_state(source), 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()) From a7434bd50c40476351dba88118960144d98f561d Mon Sep 17 00:00:00 2001 From: refinedheart <3102024448@qq.com> Date: Sun, 25 Oct 2026 03:21:46 +0800 Subject: [PATCH 2/2] [checkpoint] require patched torch for legacy optimizer files Require torch >= 2.10 before loading legacy optimizer files, parameter groups or state shards with weights_only=True. Earlier versions have known weights-only vulnerabilities. Keep this policy local to optimizer readers; v2 safetensors and the project dependency range are unchanged. Cover rejection before deserialization, patched-version dispatch and v2 loading without the legacy version gate. Document the compatibility boundary and trusted migration path. --- colossalai/checkpoint_io/utils.py | 17 ++++- docs/source/en/basics/booster_checkpoint.md | 2 + .../zh-Hans/basics/booster_checkpoint.md | 2 + .../test_optimizer_safetensors_v2.py | 69 +++++++++++++++++-- 4 files changed, 83 insertions(+), 7 deletions(-) diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py index 45469ddb0032..f259f9b31053 100644 --- a/colossalai/checkpoint_io/utils.py +++ b/colossalai/checkpoint_io/utils.py @@ -506,16 +506,27 @@ def save_param_groups(state_dict: dict, group_file_path: str) -> None: 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 torch.load(group_file_path, map_location="cpu", weights_only=True) + 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 torch.load(checkpoint, map_location="cpu", weights_only=True) + return _load_legacy_optimizer_file(checkpoint) def merge_optimizer_param_group_defaults(saved_groups: List[dict], current_groups: List[dict]) -> None: @@ -1046,7 +1057,7 @@ def load_optim_or_model_shard(path: str, is_optim: bool, use_safetensors: bool) if path.endswith(".safetensors"): 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/docs/source/en/basics/booster_checkpoint.md b/docs/source/en/basics/booster_checkpoint.md index b1ba24145705..89c321dd8d6e 100644 --- a/docs/source/en/basics/booster_checkpoint.md +++ b/docs/source/en/basics/booster_checkpoint.md @@ -59,6 +59,8 @@ Async optimizer checkpoints use a versioned, data-only safetensors format. Tenso 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 diff --git a/docs/source/zh-Hans/basics/booster_checkpoint.md b/docs/source/zh-Hans/basics/booster_checkpoint.md index 96a7f64c96d6..4f8893ed1d2c 100644 --- a/docs/source/zh-Hans/basics/booster_checkpoint.md +++ b/docs/source/zh-Hans/basics/booster_checkpoint.md @@ -60,6 +60,8 @@ booster.load(model, model_dir) 默认加载器拒绝包含 `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 diff --git a/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py b/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py index 3148692a0f4d..c5a11791e1b9 100644 --- a/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py +++ b/tests/test_checkpoint_io/test_optimizer_safetensors_v2.py @@ -4,6 +4,7 @@ import pytest import torch +from packaging.version import Version from safetensors.torch import save_file from colossalai.checkpoint_io import GeneralCheckpointIO @@ -188,11 +189,60 @@ def __reduce__(self): path = tmp_path / "optimizer.bin" torch.save(Payload(), path) - with pytest.raises(pickle.UnpicklingError): + 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)) @@ -428,6 +478,9 @@ def cpu_sync(monkeypatch): 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() @@ -636,16 +689,24 @@ def test_migration_cli_requires_flag(tmp_path): from colossalai.cli.cli import cli source = tmp_path / "old.bin" - torch.save(adam().state_dict(), source) + 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)] + cli, + [ + "checkpoint", + "migrate-optimizer", + "--trust-legacy-pickle", + str(source), + str(destination), + ], ) assert result.exit_code == 0, result.output - assert_equal(load_optimizer_state(source), load_flat(destination)) + assert_equal(expected, load_flat(destination)) @pytest.mark.skipif(