diff --git a/fedlab/utils/__init__.py b/fedlab/utils/__init__.py index 9737d42d..b79d173e 100644 --- a/fedlab/utils/__init__.py +++ b/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 \ No newline at end of file +from .serialization import SerializationTool diff --git a/fedlab/utils/compressor.py b/fedlab/utils/compressor.py new file mode 100644 index 00000000..9d650b18 --- /dev/null +++ b/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) diff --git a/tests/test_utils/test_compressor.py b/tests/test_utils/test_compressor.py new file mode 100644 index 00000000..2919308e --- /dev/null +++ b/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()