From e02d487f6bd33cdda849758704189a1c475aba79 Mon Sep 17 00:00:00 2001 From: 08raviyadav <47456958+08raviyadav@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:36:33 +0545 Subject: [PATCH 1/2] Restrict unpickler in async optimizer checkpoint loading (fixes #6471) --- colossalai/utils/safetensors.py | 59 +++++++++++- .../test_safetensors_security.py | 91 +++++++++++++++++++ 2 files changed, 148 insertions(+), 2 deletions(-) create mode 100644 tests/test_checkpoint_io/test_safetensors_security.py diff --git a/colossalai/utils/safetensors.py b/colossalai/utils/safetensors.py index b15b92fcc861..8ee886614489 100644 --- a/colossalai/utils/safetensors.py +++ b/colossalai/utils/safetensors.py @@ -1,5 +1,6 @@ # a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214 import json +import pickle import warnings from dataclasses import asdict, dataclass from typing import Dict, List, Optional, Tuple @@ -20,7 +21,7 @@ _TYPES_INV = {v: k for k, v in _TYPES.items()} import io -from torch.distributed.distributed_c10d import _pickler, _unpickler +from torch.distributed.distributed_c10d import _pickler ASYNC_WRITE_ENTRIES = 32 @@ -31,6 +32,51 @@ def _require_async_file_writer(): return AsyncFileWriter +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler for the non-tensor values stored by :func:`_flatten_optim_state_dict`. + + Those values are plain optimizer data (``param_groups`` and non-tensor state + entries such as steps and hyperparameters). Only the small allow-list of + globals below may be referenced by the pickle stream; anything else -- for + example ``os.system`` referenced by a crafted checkpoint file -- is rejected + before it can execute. See issue #6471: previously the set of keys to unpickle + came from the checkpoint's own metadata and the unpickler was unrestricted, + so a malicious ``.safetensors`` optimizer checkpoint achieved arbitrary code + execution on load. + """ + + _ALLOWED_GLOBALS = { + ("builtins", "bytearray"), + ("builtins", "complex"), + ("builtins", "frozenset"), + ("builtins", "range"), + ("builtins", "set"), + ("builtins", "slice"), + ("collections", "Counter"), + ("collections", "OrderedDict"), + ("collections", "defaultdict"), + # Tensors nested inside a non-tensor container are rebuilt with the same + # constructors that torch.load(weights_only=True) permits. + ("torch", "Tensor"), + ("torch._utils", "_rebuild_meta_tensor_no_storage"), + ("torch._utils", "_rebuild_parameter"), + ("torch._utils", "_rebuild_parameter_with_state"), + ("torch._utils", "_rebuild_tensor"), + ("torch._utils", "_rebuild_tensor_v2"), + ("torch.storage", "TypedStorage"), + ("torch.storage", "UntypedStorage"), + ("torch.storage", "_load_from_bytes"), + } + + def find_class(self, module, name): + if (module, name) in self._ALLOWED_GLOBALS: + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"global '{module}.{name}' is forbidden when loading optimizer checkpoints; " + "only plain data and tensors may be restored from non-tensor entries" + ) + + def _object_to_tensor(obj, device): f = io.BytesIO() _pickler(f).dump(obj) @@ -45,7 +91,7 @@ def _object_to_tensor(obj, device): def _tensor_to_object(tensor, tensor_size): tensor = tensor.cpu() buf = tensor.numpy().tobytes()[:tensor_size] - return _unpickler(io.BytesIO(buf)).load() + return _RestrictedUnpickler(io.BytesIO(buf)).load() @dataclass @@ -107,6 +153,15 @@ def _unflatten_optim_state_dict(flat_dict: dict, metadata: Optional[dict] = None non_tensor_keys = json.loads(metadata["non_tensor_keys"]) else: non_tensor_keys = [] + # The metadata comes from the checkpoint file itself, so only accept the key + # shapes _flatten_optim_state_dict() can produce: "param_groups" and + # "state". + for key in non_tensor_keys: + if key == "param_groups": + continue + parts = key.split(seperator) + if len(parts) != 3 or parts[0] != "state" or not parts[1].isdigit(): + raise ValueError(f"invalid non-tensor key in checkpoint metadata: {key!r}") 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 diff --git a/tests/test_checkpoint_io/test_safetensors_security.py b/tests/test_checkpoint_io/test_safetensors_security.py new file mode 100644 index 000000000000..c56f4e4e88e7 --- /dev/null +++ b/tests/test_checkpoint_io/test_safetensors_security.py @@ -0,0 +1,91 @@ +"""Regression tests for issue #6471. + +The async optimizer-checkpoint loader used to unpickle whichever entries a +checkpoint file's own metadata named, with an unrestricted unpickler -- a +crafted .safetensors file could execute arbitrary code on load. These tests +check that (a) legitimate optimizer state dicts still round-trip and +(b) a malicious checkpoint is rejected without executing its payload. +""" + +import json +import os +import pickle + +import pytest +import torch +from safetensors.torch import save_file + +from colossalai.utils.safetensors import _flatten_optim_state_dict, _unflatten_optim_state_dict, load_flat + + +class _Payload: + def __reduce__(self): + return (os.system, ("echo pwned > PWNED.txt",)) + + +def _write_malicious_checkpoint(path, key): + payload = pickle.dumps(_Payload(), protocol=2) + tensors = { + "param_groups": torch.tensor([0], dtype=torch.uint8), + key: torch.tensor(list(payload), dtype=torch.uint8), + } + save_file(tensors, str(path), metadata={"non_tensor_keys": json.dumps([key])}) + + +def test_malicious_checkpoint_does_not_execute(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + checkpoint = tmp_path / "optimizer.safetensors" + # A key with the exact shape the writer produces, carrying pickle payload bytes. + _write_malicious_checkpoint(checkpoint, "state.0.candidate") + with pytest.raises(pickle.UnpicklingError): + load_flat(str(checkpoint)) + assert not (tmp_path / "PWNED.txt").exists() + + +def test_off_schema_non_tensor_key_rejected(tmp_path): + checkpoint = tmp_path / "optimizer.safetensors" + _write_malicious_checkpoint(checkpoint, "totally.bogus.key.extra") + with pytest.raises(ValueError): + load_flat(str(checkpoint)) + + +def test_optimizer_state_dict_round_trip(tmp_path): + model = torch.nn.Linear(4, 2) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01, betas=(0.9, 0.95)) + for _ in range(2): + loss = model(torch.randn(3, 4)).sum() + loss.backward() + optimizer.step() + optimizer.zero_grad() + + state_dict = optimizer.state_dict() + flat, metadata = _flatten_optim_state_dict(state_dict) + restored = _unflatten_optim_state_dict( + dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])} + ) + assert restored["param_groups"] == state_dict["param_groups"] + for idx, fields in state_dict["state"].items(): + for name, value in fields.items(): + restored_value = restored["state"][idx][name] + if isinstance(value, torch.Tensor): + assert torch.equal(restored_value, value) + else: + assert restored_value == value + + checkpoint = tmp_path / "optimizer.safetensors" + save_file(flat, str(checkpoint), metadata={"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])}) + assert load_flat(str(checkpoint))["param_groups"] == state_dict["param_groups"] + + +def test_tensor_nested_in_non_tensor_value_round_trips(): + state_dict = { + "state": {0: {"step": 3, "bundle": {"vec": torch.tensor([1.0, 2.0]), "name": "x"}}}, + "param_groups": [{"lr": 0.1, "params": [0]}], + } + flat, metadata = _flatten_optim_state_dict(state_dict) + restored = _unflatten_optim_state_dict( + dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])} + ) + bundle = restored["state"][0]["bundle"] + assert bundle["name"] == "x" + assert torch.equal(bundle["vec"], torch.tensor([1.0, 2.0])) From 1d58cb5b6e12ae6ee75647b6ea89b096948e094b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:56:20 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_checkpoint_io/test_safetensors_security.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_checkpoint_io/test_safetensors_security.py b/tests/test_checkpoint_io/test_safetensors_security.py index c56f4e4e88e7..ee7b9ca05084 100644 --- a/tests/test_checkpoint_io/test_safetensors_security.py +++ b/tests/test_checkpoint_io/test_safetensors_security.py @@ -60,9 +60,7 @@ def test_optimizer_state_dict_round_trip(tmp_path): state_dict = optimizer.state_dict() flat, metadata = _flatten_optim_state_dict(state_dict) - restored = _unflatten_optim_state_dict( - dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])} - ) + restored = _unflatten_optim_state_dict(dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])}) assert restored["param_groups"] == state_dict["param_groups"] for idx, fields in state_dict["state"].items(): for name, value in fields.items(): @@ -83,9 +81,7 @@ def test_tensor_nested_in_non_tensor_value_round_trips(): "param_groups": [{"lr": 0.1, "params": [0]}], } flat, metadata = _flatten_optim_state_dict(state_dict) - restored = _unflatten_optim_state_dict( - dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])} - ) + restored = _unflatten_optim_state_dict(dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])}) bundle = restored["state"][0]["bundle"] assert bundle["name"] == "x" assert torch.equal(bundle["vec"], torch.tensor([1.0, 2.0]))