[refinedheart]

📌 Checklist before creating the PR

• I have linked an existing issue for traceability: [BUG]: Async optimizer checkpoint loader unpickles any tensor named by the file's own metadata (arbitrary code execution on load_flat / load_optimizer) #6471, reported by @3em0.
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install

🚨 Issue number

Fixes #6471.

Related PR: #6478, which retains the existing format with a restricted unpickler. This PR proposes a versioned format without embedded pickle, with explicit migration for trusted legacy files.

📝 What does this PR do?

Background

Async optimizer checkpoints store non-tensor values, including param_groups, as pickle bytes in uint8 tensors. load_flat() uses the checkpoint's own non_tensor_keys metadata to select which tensors to unpickle. A valid safetensors container can therefore reach unrestricted object decoding through Booster.load_optimizer(), before optimizer-state validation.

Changes

• Add optimizer safetensors v2: native tensors and a typed JSON schema for scalars, containers and tensor references. Validate the version, payload kind, schema limits and references before returning optimizer state.
• Write parameter-group companions in v2 and use the common reader across General, Hybrid, MoE, Gemini, LowLevelZero and FSDP checkpoint paths. Preserve model snapshot/blob paths while constraining optimizer index paths.
• Reject legacy embedded-pickle files in normal loaders. Add colossalai checkpoint migrate-optimizer --trust-legacy-pickle SOURCE DESTINATION for trusted old files or directories, preserving inputs and refusing existing destinations.
• Require torch >=2.10.0 for normal restricted loading of legacy optimizer PyTorch files. Keep this check local to optimizer readers; the project's torch dependency range is unchanged.
• Cover repeated saves, format validation, migration and loader boundaries, and include the CPU security tests in both CPU CI entry points.

User and compatibility impact

• Checkpoint format: new async optimizer files and parameter-group companions use v2. Old embedded-pickle files need explicit migration, including valid files produced by older ColossalAI versions. Older readers cannot read v2; retain original checkpoints for rollback. Synchronous PyTorch state files retain their format.
• Loading/API: public Booster entry points remain unchanged; migration is a new opt-in CLI. v2 supports fixed basic/container types and dense non-quantized tensors, not arbitrary custom classes.
• PyTorch versions: legacy optimizer PyTorch files require >=2.10.0 because of known weights-only vulnerabilities. The legacy-file gate is not applied to v2 safetensors. Migration still performs legacy decoding and must only receive trusted input.
• Numerics: expert TP/EP/moe_dp sharding algorithms are unchanged. Distributed restore remains subject to the unexecuted tests below; no performance claim is made.

Test environment

Python 3.10.12, torch 2.11.0+cu128, one NVIDIA H20 for single-GPU tests. Tested tree: a7434bd5, based on main 0ac6a922.

Transformers	PEFT	Diffusers
5.17.0	0.21.2	0.35.2
4.51.3	0.13.2	0.29.0

Test commands and results

CPU targets, run separately in each environment:

CUDA_VISIBLE_DEVICES="" python -m pytest -p no:cacheprovider -ra \
  tests/test_checkpoint_io/test_optimizer_safetensors_v2.py \
  tests/test_checkpoint_io/test_optimizer_safetensors_plugins.py
bash .github/scripts/ci/run_colossalai_batch.sh cpu
bash .github/scripts/ci/run_cpu_unit_tests.sh

Scope	Each Transformers version
Format/security CPU suite	87 passed, 10 CUDA-only skips
Local Layer 1 CPU batch	110 passed, 3 CUDA-only skips
Local fork-safe CPU entry	114 passed, 3 CUDA-only skips
Single-GPU TensorNVMe, repeated saves and selected General/safetensors IO nodes	16 passed

The scopes overlap. These are local results, not GitHub Actions results. Version-gate controls produced nine expected failures on the previous implementation in each environment and pass after the gate.

Known issues / unexecuted tests

• Two-rank plugin and real migrated MoE optimizer restore/step tests are included but not run: two free GPUs on one host were unavailable. Full 4/8-GPU and multi-stage PP coverage is also pending.
• Actual execution used torch 2.11 only. Version-gate tests mock version strings to check dispatch and rejection order; they do not validate other torch installations. Other Python versions were not tested.

Risk and rollback

The main risk is checkpoint compatibility. Migration preserves source files, and reverting code does not convert v2 files back to the old format. Keep original checkpoints and the previous validated environment when planning rollback. Migration is not an authenticity check for untrusted input.

Documentation and release impact

Update the English and Chinese checkpoint documentation with the v2 format, legacy loading policy and migration command. No dependency, image or release-version change.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• The linked issue describes the problem and includes a reproducer.
• I have performed a self-review of my code
• I have added format, security, migration and loader tests; executed results and pending distributed coverage are listed above.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.