Skip to content

[checkpoint]: remove embedded pickle from async optimizer checkpoints #6479

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuxuan-lou merged 2 commits into from
Oct 25, 2026
Merged

[checkpoint]: remove embedded pickle from async optimizer checkpoints #6479

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e57b0de
[checkpoint] write data-only optimizer safetensors v2
refinedheart Oct 24, 2026
a7434bd
[checkpoint] require patched torch for legacy optimizer files
refinedheart Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/scripts/ci/run_colossalai_batch.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions .github/scripts/ci/run_cpu_unit_tests.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 10 additions & 5 deletions colossalai/booster/plugin/gemini_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand Down
9 changes: 5 additions & 4 deletions colossalai/booster/plugin/low_level_zero_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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()
Expand Down
21 changes: 14 additions & 7 deletions colossalai/booster/plugin/torch_fsdp_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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()
Expand All @@ -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 = {}
Expand Down
14 changes: 9 additions & 5 deletions colossalai/checkpoint_io/general_checkpoint_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand Down
41 changes: 24 additions & 17 deletions colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"]}
Expand Down Expand Up @@ -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"]}
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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)

Expand Down
Loading
Loading