Add opt-in lossless int16 checkpoint encoding

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.
+**Workaround**: If `load_checkpoint` raises, call `tq.init()` again to reset the system to a clean state before retrying.

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()

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)

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
 

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(

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)

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]