diff --git a/docs/checkpoint.md b/docs/checkpoint.md index a5b36803..ee574648 100644 --- a/docs/checkpoint.md +++ b/docs/checkpoint.md @@ -1,5 +1,14 @@ # Checkpoint: Save and Restore System State +## Optional integer packing for SimpleStorage + +Set `backend.SimpleStorage.checkpoint_uint8_fields` to a list of field names +(for example `[routed_experts]`) to store contiguous CPU int16 tensors as uint8 +when every value lies in `[0, 255]`. This setting defaults to an empty list. +Other tensors and all remaining fields, including pending prompts, stay intact. +Loading restores int16 values automatically, independently of the current setting. +Compact checkpoints require a reader with this codec; older readers reject them. + > Last updated: 07/03/2026 ## Overview @@ -165,4 +174,4 @@ client.load_controller_checkpoint(...) # (2) controller restored second If step (1) partially succeeds and step (2) fails, the system is left in a mixed state: some storage units hold checkpoint data while the controller still reflects its pre-restore state. There is no rollback path. -**Workaround**: If `load_checkpoint` raises, call `tq.init()` again to reset the system to a clean state before retrying. \ No newline at end of file +**Workaround**: If `load_checkpoint` raises, call `tq.init()` again to reset the system to a clean state before retrying. diff --git a/tests/e2e/test_checkpoint_integer_encoding.py b/tests/e2e/test_checkpoint_integer_encoding.py new file mode 100644 index 00000000..a02a2235 --- /dev/null +++ b/tests/e2e/test_checkpoint_integer_encoding.py @@ -0,0 +1,103 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import tempfile + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import NonTensorStack, TensorDict + +import transfer_queue as tq + + +@pytest.mark.parametrize("compact_fields", [[], ["routed_experts"]]) +def test_native_checkpoint_preserves_finished_and_pending_rows(tmp_path, compact_fields): + ray.init( + address="local", + num_cpus=3, + num_gpus=0, + include_dashboard=False, + object_store_memory=128 * 1024**2, + _memory=1024**3, + _temp_dir=tempfile.mkdtemp(prefix="tq-codec-"), + ) + try: + tq.init( + OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 20, + "num_data_storage_units": 1, + "checkpoint_uint8_fields": compact_fields, + }, + }, + } + ) + ) + keys = ["finished", "pending"] + prompt = [{"role": "user", "content": "The complete pending prompt and tool state."}] + routes = torch.arange(16384, dtype=torch.int16).remainder(16).reshape(2, 128, 16, 4) + fields = TensorDict( + { + "routed_experts": routes, + "input_ids": torch.arange(256).reshape(2, 128), + "loss_mask": torch.ones(2, 128, dtype=torch.bool), + "rollout_log_probs": torch.zeros(2, 128), + "rm_scores": torch.tensor([[1.0], [0.0]]), + "response_weight_versions": torch.full((2, 128), 2), + "raw_prompt": NonTensorStack(copy.deepcopy(prompt), copy.deepcopy(prompt)), + "tools_kwargs": NonTensorStack({"task": {"id": "a"}}, {"task": {"id": "b"}}), + }, + batch_size=2, + ) + tq.kv_batch_put( + keys=keys, + partition_id="train", + fields=fields, + tags=[ + {"is_prompt": True, "status": "finished", "global_steps": 2}, + {"is_prompt": True, "status": "running", "global_steps": 2}, + ], + ) + tags_before = copy.deepcopy(tq.kv_list("train")) + before = tq.kv_batch_get(keys=keys, partition_id="train") + checkpoint = tmp_path / "checkpoint" + tq.save_checkpoint(checkpoint) + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + ray.get(controller.clear_partition.remote("train")) + tq.load_checkpoint(checkpoint) + assert tq.kv_list("train") == tags_before + after = tq.kv_batch_get(keys=keys, partition_id="train") + for field in fields.keys(): + if field in ("raw_prompt", "tools_kwargs"): + assert list(after[field]) == list(before[field]) + else: + assert after[field].dtype == before[field].dtype + if after[field].is_nested: + assert torch.equal(after[field].offsets(), before[field].offsets()) + for actual, expected in zip(after[field].unbind(), before[field].unbind(), strict=True): + assert actual.shape == expected.shape + assert torch.equal(actual, expected) + else: + assert torch.equal(after[field], before[field]) + finally: + tq.close() + ray.shutdown() diff --git a/tests/test_checkpoint_codec.py b/tests/test_checkpoint_codec.py new file mode 100644 index 00000000..637ed745 --- /dev/null +++ b/tests/test_checkpoint_codec.py @@ -0,0 +1,85 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import io +import pickle + +import pytest +import torch + +from transfer_queue.storage.checkpoint_codec import dump_checkpoint + + +def test_compact_checkpoint_restores_dtype_values_and_other_fields(): + routes = torch.arange(32768, dtype=torch.int16).remainder(256).reshape(128, 16, 16) + pending = [{"role": "user", "content": "Keep the complete pending prompt."}] + state = { + "storage_unit_id": "storage", + "storage_unit_size": 10, + "active_keys": {2, 9}, + "field_data": { + "routed_experts": {2: routes}, + "raw_prompt": {9: pending}, + "metrics": {2: {"duration": 1.25}}, + "input_ids": {2: torch.arange(128)}, + "response_weight_versions": {2: torch.full((128,), 2)}, + }, + } + stream = io.BytesIO() + dump_checkpoint(state, stream, ["routed_experts"]) + assert stream.tell() < len(pickle.dumps(state, protocol=pickle.HIGHEST_PROTOCOL)) - routes.numel() // 2 + restored = pickle.loads(stream.getvalue()) + assert restored["active_keys"] == {2, 9} + assert restored["field_data"]["raw_prompt"][9] == pending + assert restored["field_data"]["metrics"] == state["field_data"]["metrics"] + for field in ("routed_experts", "input_ids", "response_weight_versions"): + expected = state["field_data"][field][2] + actual = restored["field_data"][field][2] + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert torch.equal(actual, expected) + assert state["field_data"]["routed_experts"][2] is routes + assert routes.dtype == torch.int16 + + +@pytest.mark.parametrize( + "value", + [ + torch.tensor([-1, 0, 15], dtype=torch.int16), + torch.tensor([0, 255, 256], dtype=torch.int16), + torch.tensor([0.0, 1.5]), + torch.tensor([0, 15], dtype=torch.int64), + torch.empty(0, dtype=torch.int16), + torch.arange(16, dtype=torch.int16).reshape(4, 4).t(), + ], +) +def test_unsupported_tensor_keeps_legacy_representation(value): + state = {"field_data": {"routed_experts": {0: value}}} + stream = io.BytesIO() + dump_checkpoint(state, stream, ["routed_experts"]) + assert stream.getvalue() == pickle.dumps(state, protocol=pickle.HIGHEST_PROTOCOL) + restored = pickle.loads(stream.getvalue())["field_data"]["routed_experts"][0] + assert restored.dtype == value.dtype + assert restored.stride() == value.stride() + assert torch.equal(restored, value) + + +def test_default_and_unselected_fields_keep_legacy_representation(): + routes = torch.arange(16, dtype=torch.int16) + state = {"field_data": {"routed_experts": {0: routes}}} + for fields in ((), ["another_field"]): + stream = io.BytesIO() + dump_checkpoint(state, stream, fields) + assert stream.getvalue() == pickle.dumps(state, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index bd83a599..1510e5ac 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -32,6 +32,9 @@ backend: num_data_storage_units: 2 # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. required_node_resource: null + # Opt-in checkpoint encoding for CPU int16 tensors with values in [0, 255]. + # Loads restore int16; other tensors and fields are unchanged. + checkpoint_uint8_fields: [] # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 381341ec..376aadf1 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -50,6 +50,7 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: name=f"TransferQueueStorageUnit#{storage_unit_rank}", ).remote( storage_unit_size=storage_unit_size, + checkpoint_uint8_fields=conf.backend.SimpleStorage.get("checkpoint_uint8_fields", []), ) simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node logger.info( diff --git a/transfer_queue/storage/checkpoint_codec.py b/transfer_queue/storage/checkpoint_codec.py new file mode 100644 index 00000000..46a18594 --- /dev/null +++ b/transfer_queue/storage/checkpoint_codec.py @@ -0,0 +1,62 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Opt-in lossless integer encoding at the SimpleStorage checkpoint boundary.""" + +import pickle + +import torch + + +def _restore_int16_from_uint8_v1(value): + # A named reconstruction function makes older readers reject this encoding + # before installing checkpoint data, rather than exposing packed uint8 values. + if type(value) is not torch.Tensor or value.dtype != torch.uint8: + raise ValueError("Invalid compact int16 checkpoint tensor") + return value.to(torch.int16) + + +class _CheckpointPickler(pickle.Pickler): + def __init__(self, stream, state, fields): + super().__init__(stream, protocol=pickle.HIGHEST_PROTOCOL) + self._selected = { + id(value) + for field in fields + for value in state["field_data"].get(field, {}).values() + if type(value) is torch.Tensor + } + + def reducer_override(self, value): + if ( + id(value) in self._selected + and type(value) is torch.Tensor + and value.dtype == torch.int16 + and value.device.type == "cpu" + and value.layout == torch.strided + and value.is_contiguous() + and value.numel() + and value.min().item() >= 0 + and value.max().item() <= 255 + ): + return _restore_int16_from_uint8_v1, (value.to(torch.uint8),) + return NotImplemented + + +def dump_checkpoint(state, stream, uint8_fields=()): + """Pack eligible int16 tensors from explicit fields; retain all other values.""" + if uint8_fields: + _CheckpointPickler(stream, state, uint8_fields).dump(state) + else: + pickle.dump(state, stream, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index aa23855f..a0b31e55 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -25,6 +25,7 @@ import ray import zmq +from transfer_queue.storage.checkpoint_codec import dump_checkpoint from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger @@ -154,15 +155,22 @@ class SimpleStorageUnit: zmq_server_info: ZMQ connection information for clients. """ - def __init__(self, storage_unit_size: int | None = None): + def __init__(self, storage_unit_size: int | None = None, checkpoint_uint8_fields: list[str] | None = None): """Initialize a SimpleStorageUnit with the specified size. Args: storage_unit_size: Maximum number of elements that can be stored in this storage unit. If None, the storage unit has unlimited capacity. + checkpoint_uint8_fields: Fields eligible for lossless int16 checkpoint packing. + Empty by default; loading always restores the original dtype. """ self.storage_unit_id = f"TQ_STORAGE_UNIT_{uuid4().hex[:8]}" self.storage_unit_size = storage_unit_size + if isinstance(checkpoint_uint8_fields, str) or any( + not isinstance(field, str) or not field for field in (checkpoint_uint8_fields or ()) + ): + raise ValueError("checkpoint_uint8_fields must be a list of nonempty field names") + self.checkpoint_uint8_fields = tuple(checkpoint_uint8_fields or ()) self.storage_data = StorageUnitData(self.storage_unit_size) @@ -569,7 +577,7 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: "active_keys": self.storage_data._active_keys, } with open(path, "wb") as f: - pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + dump_checkpoint(state, f, self.checkpoint_uint8_fields) logger.info(f"[{self.storage_unit_id}]: saved checkpoint to {path}") return ZMQMessage.create( request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type]