Add Int4BlockCompressor utility for quantized tensor serialization

fedlab/utils/__init__.py

@@ -14,6 +14,7 @@
 
 
 from .aggregator import Aggregators
+from .compressor import Int4BlockCompressor
 from .logger import Logger
 from .message_code import MessageCode
-from .serialization import SerializationTool
+from .serialization import SerializationTool

fedlab/utils/compressor.py

@@ -0,0 +1,197 @@
+# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group)
+
+# 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 torch
+import torch.nn as nn
+from typing import Tuple, Optional
+from .serialization import SerializationTool
+
+
+class Int4BlockCompressor(object):
+    """Block-wise INT4 tensor compressor with dynamic scaling and residual error feedback.
+
+    Pack two signed 4-bit values into a single uint8 byte, yielding an 8x reduction
+    in raw tensor transmission size for serialized model parameters and gradients.
+
+    Args:
+        block_size (int, optional): Size of each quantization block. Must be positive and even. Defaults to 64.
+        use_error_feedback (bool, optional): Whether to track local quantization residuals across rounds. Defaults to False.
+    """
+
+    def __init__(self, block_size: int = 64, use_error_feedback: bool = False):
+        if block_size <= 0 or block_size % 2 != 0:
+            raise ValueError(f"block_size must be a positive even integer, got {block_size}")
+        self.block_size = block_size
+        self.use_error_feedback = use_error_feedback
+        self.residual: Optional[torch.Tensor] = None
+
+    def reset_error(self) -> None:
+        """Clear cached residual buffer."""
+        self.residual = None
+
+    def compress(self, tensor: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
+        """Quantize and pack a tensor into INT4 nibbles with per-block scale factors.
+
+        Args:
+            tensor (torch.Tensor): Input tensor.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
+                - packed: uint8 tensor holding two 4-bit elements per byte.
+                - scale: float scale factors per block.
+                - orig_shape: original tensor shape.
+                - orig_dtype: original tensor dtype.
+        """
+        orig_shape = tensor.shape
+        orig_dtype = tensor.dtype
+        x = tensor.detach().float()
+
+        if self.use_error_feedback:
+            if self.residual is None or self.residual.shape != orig_shape or self.residual.device != tensor.device:
+                self.residual = torch.zeros_like(x)
+            x = x + self.residual
+
+        flat = x.flatten()
+        n = flat.numel()
+        pad_len = (self.block_size - (n % self.block_size)) % self.block_size
+        if pad_len > 0:
+            padded = torch.nn.functional.pad(flat, (0, pad_len), value=0.0)
+        else:
+            padded = flat
+
+        m = padded.numel() // self.block_size
+        blocks = padded.view(m, self.block_size)
+
+        max_abs = torch.max(torch.abs(blocks), dim=1, keepdim=True).values
+        scale = torch.clamp(max_abs / 7.0, min=1e-8)
+
+        quantized = torch.clamp(torch.round(blocks / scale), -8, 7).to(torch.int32)
+        nibbles = (quantized + 8).to(torch.uint8)
+
+        paired = nibbles.view(m, self.block_size // 2, 2)
+        packed = ((paired[:, :, 0] << 4) | (paired[:, :, 1] & 0x0F)).flatten()
+        scale_vec = scale.squeeze(-1)
+
+        if self.use_error_feedback:
+            decompressed = self._decompress_core(packed, scale_vec, orig_shape, orig_dtype)
+            self.residual = (x - decompressed).detach()
+
+        return packed, scale_vec, orig_shape, orig_dtype
+
+    def _decompress_core(
+        self,
+        packed: torch.Tensor,
+        scale: torch.Tensor,
+        orig_shape: torch.Size,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        m = scale.numel()
+        block_bytes = self.block_size // 2
+        packed_blocks = packed.view(m, block_bytes)
+
+        high = (packed_blocks >> 4) & 0x0F
+        low = packed_blocks & 0x0F
+
+        unpacked = torch.stack([high, low], dim=2).view(m, self.block_size)
+        values = unpacked.to(torch.float32) - 8.0
+
+        scaled = values * scale.view(m, 1)
+        total_elements = orig_shape.numel()
+        return scaled.flatten()[:total_elements].view(orig_shape).to(dtype)
+
+    def decompress(
+        self,
+        packed: torch.Tensor,
+        scale: torch.Tensor,
+        orig_shape: torch.Size,
+        dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        """Unpack INT4 nibbles and rescale back to floating-point tensor."""
+        return self._decompress_core(packed, scale, orig_shape, dtype)
+
+    @classmethod
+    def compress_tensor(
+        cls,
+        tensor: torch.Tensor,
+        block_size: int = 64,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
+        """Stateless compression helper for individual tensors."""
+        compressor = cls(block_size=block_size, use_error_feedback=False)
+        return compressor.compress(tensor)
+
+    @classmethod
+    def decompress_tensor(
+        cls,
+        packed: torch.Tensor,
+        scale: torch.Tensor,
+        orig_shape: torch.Size,
+        dtype: torch.dtype = torch.float32,
+        block_size: int = 64,
+    ) -> torch.Tensor:
+        """Stateless decompression helper for individual tensors."""
+        compressor = cls(block_size=block_size, use_error_feedback=False)
+        return compressor.decompress(packed, scale, orig_shape, dtype=dtype)
+
+    @classmethod
+    def compress_model(
+        cls,
+        model: nn.Module,
+        block_size: int = 64,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
+        """Serialize and compress full model parameter vector."""
+        flat_params = SerializationTool.serialize_model(model)
+        return cls.compress_tensor(flat_params, block_size=block_size)
+
+    @classmethod
+    def deserialize_model(
+        cls,
+        model: nn.Module,
+        packed: torch.Tensor,
+        scale: torch.Tensor,
+        orig_shape: torch.Size,
+        dtype: torch.dtype = torch.float32,
+        block_size: int = 64,
+        mode: str = "copy",
+    ) -> None:
+        """Decompress parameter vector and deserialize into model."""
+        decompressed = cls.decompress_tensor(
+            packed, scale, orig_shape, dtype=dtype, block_size=block_size
+        )
+        SerializationTool.deserialize_model(model, decompressed, mode=mode)
+
+    @classmethod
+    def compress_model_gradients(
+        cls,
+        model: nn.Module,
+        block_size: int = 64,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
+        """Serialize and compress model gradients."""
+        flat_grads = SerializationTool.serialize_model_gradients(model)
+        return cls.compress_tensor(flat_grads, block_size=block_size)
+
+    @classmethod
+    def deserialize_model_gradients(
+        cls,
+        model: nn.Module,
+        packed: torch.Tensor,
+        scale: torch.Tensor,
+        orig_shape: torch.Size,
+        dtype: torch.dtype = torch.float32,
+        block_size: int = 64,
+    ) -> None:
+        """Decompress gradient vector and deserialize into model grad attributes."""
+        decompressed = cls.decompress_tensor(
+            packed, scale, orig_shape, dtype=dtype, block_size=block_size
+        )
+        SerializationTool.deserialize_model_gradients(model, decompressed)

tests/test_utils/test_compressor.py

@@ -0,0 +1,174 @@
+# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group)
+
+# 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 unittest
+import torch
+import torch.nn as nn
+from fedlab.utils.compressor import Int4BlockCompressor
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features=64, hidden=32, out_features=10):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(in_features, hidden)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden, out_features)
+
+    def forward(self, x):
+        return self.fc2(self.relu(self.fc1(x)))
+
+
+class CompressorTestCase(unittest.TestCase):
+    def test_invalid_block_size(self):
+        with self.assertRaises(ValueError):
+            Int4BlockCompressor(block_size=63)
+        with self.assertRaises(ValueError):
+            Int4BlockCompressor(block_size=0)
+        with self.assertRaises(ValueError):
+            Int4BlockCompressor(block_size=-16)
+
+    def test_shape_and_dtype_preservation(self):
+        compressor = Int4BlockCompressor(block_size=64)
+
+        for dtype in [torch.float32, torch.float16]:
+            # Exact block multiple
+            t1 = torch.randn(256, dtype=dtype)
+            packed, scale, shape, orig_dtype = compressor.compress(t1)
+            recon1 = compressor.decompress(packed, scale, shape, dtype=orig_dtype)
+            self.assertEqual(recon1.shape, t1.shape)
+            self.assertEqual(recon1.dtype, dtype)
+
+            # Requires padding
+            t2 = torch.randn(173, dtype=dtype)
+            packed, scale, shape, orig_dtype = compressor.compress(t2)
+            recon2 = compressor.decompress(packed, scale, shape, dtype=orig_dtype)
+            self.assertEqual(recon2.shape, t2.shape)
+            self.assertEqual(recon2.dtype, dtype)
+
+            # Multi-dimensional
+            t3 = torch.randn(2, 5, 17, dtype=dtype)
+            packed, scale, shape, orig_dtype = compressor.compress(t3)
+            recon3 = compressor.decompress(packed, scale, shape, dtype=orig_dtype)
+            self.assertEqual(recon3.shape, t3.shape)
+            self.assertEqual(recon3.dtype, dtype)
+
+    def test_compression_ratio(self):
+        n = 4096
+        tensor = torch.randn(n, dtype=torch.float32)
+        raw_bytes = tensor.numel() * tensor.element_size()
+
+        packed, scale, _, _ = Int4BlockCompressor.compress_tensor(tensor, block_size=64)
+        packed_bytes = packed.numel() * packed.element_size()
+        scale_bytes = scale.numel() * scale.element_size()
+
+        # 4 bits per element gives an exact 8x payload reduction
+        self.assertEqual(packed_bytes, raw_bytes // 8)
+        # With scales, total footprint is still >7x smaller than float32
+        self.assertGreater(raw_bytes / (packed_bytes + scale_bytes), 7.0)
+
+    def test_edge_tensors(self):
+        compressor = Int4BlockCompressor(block_size=64)
+
+        # All zeros
+        zeros = torch.zeros(128)
+        packed, scale, shape, dtype = compressor.compress(zeros)
+        recon_zeros = compressor.decompress(packed, scale, shape, dtype)
+        self.assertTrue(torch.allclose(recon_zeros, zeros, atol=1e-5))
+
+        # Constant non-zero
+        const = torch.full((128,), 3.5)
+        packed, scale, shape, dtype = compressor.compress(const)
+        recon_const = compressor.decompress(packed, scale, shape, dtype)
+        self.assertTrue(torch.allclose(recon_const, const, atol=0.6))
+
+    def test_reconstruction_fidelity(self):
+        tensor = torch.randn(1024)
+        packed, scale, shape, dtype = Int4BlockCompressor.compress_tensor(tensor, block_size=64)
+        recon = Int4BlockCompressor.decompress_tensor(packed, scale, shape, dtype, block_size=64)
+
+        mse = torch.mean((tensor - recon) ** 2).item()
+        self.assertLess(mse, 0.03)
+
+    def test_error_feedback_drift(self):
+        compressor_ef = Int4BlockCompressor(block_size=64, use_error_feedback=True)
+        compressor_raw = Int4BlockCompressor(block_size=64, use_error_feedback=False)
+
+        torch.manual_seed(0)
+        rounds = 25
+        dim = 256
+
+        target = torch.zeros(dim)
+        accum_ef = torch.zeros(dim)
+        accum_raw = torch.zeros(dim)
+
+        for _ in range(rounds):
+            update = torch.randn(dim) * 0.05
+            target += update
+
+            p_ef, s_ef, sh_ef, dt_ef = compressor_ef.compress(update)
+            accum_ef += compressor_ef.decompress(p_ef, s_ef, sh_ef, dt_ef)
+
+            p_raw, s_raw, sh_raw, dt_raw = compressor_raw.compress(update)
+            accum_raw += compressor_raw.decompress(p_raw, s_raw, sh_raw, dt_raw)
+
+        drift_ef = torch.norm(target - accum_ef).item()
+        drift_raw = torch.norm(target - accum_raw).item()
+
+        self.assertLess(drift_ef, drift_raw)
+
+        compressor_ef.reset_error()
+        self.assertIsNone(compressor_ef.residual)
+
+    def test_model_gradients_pipeline(self):
+        model = MLP(in_features=32, hidden=16, out_features=4)
+        x = torch.randn(4, 32)
+        model(x).sum().backward()
+
+        packed, scale, shape, dtype = Int4BlockCompressor.compress_model_gradients(
+            model, block_size=32
+        )
+
+        target = MLP(in_features=32, hidden=16, out_features=4)
+        for p in target.parameters():
+            p.grad = torch.zeros_like(p.data)
+
+        Int4BlockCompressor.deserialize_model_gradients(
+            target, packed, scale, shape, dtype=dtype, block_size=32
+        )
+
+        for p_src, p_dst in zip(model.parameters(), target.parameters()):
+            self.assertIsNotNone(p_dst.grad)
+            cos_sim = torch.cosine_similarity(
+                p_src.grad.flatten().unsqueeze(0),
+                p_dst.grad.flatten().unsqueeze(0),
+            ).item()
+            self.assertGreater(cos_sim, 0.95)
+
+    def test_model_parameters_pipeline(self):
+        model = MLP(in_features=32, hidden=16, out_features=4)
+        packed, scale, shape, dtype = Int4BlockCompressor.compress_model(model, block_size=32)
+
+        target = MLP(in_features=32, hidden=16, out_features=4)
+        Int4BlockCompressor.deserialize_model(
+            target, packed, scale, shape, dtype=dtype, block_size=32, mode="copy"
+        )
+
+        for p_src, p_dst in zip(model.parameters(), target.parameters()):
+            diff = torch.norm(p_src.data - p_dst.data).item()
+            norm = torch.norm(p_src.data).item()
+            self.assertLess(diff / norm, 0.15)
+
+
+if __name__ == "__main__":
+    unittest.main()