Module neuralops.compression_engine.huffman_coder

Huffman Coder & Entropy Compression Engine for Weight and KV Cache Quantization. Registered as 'huffman'. Production-hardened with vectorized encode/decode, exact byte accounting, safe LoRA freezing, and robust tokenizer forward checks.

Functions

def huffman_factory(model: nn.Module, comp_config: Any = None, **kwargs: Any) ‑> torch.nn.modules.module.Module

Expand source code

@COMPRESSOR_REGISTRY.register("huffman")
@COMPRESSOR_REGISTRY.register("huffman_coder")
def huffman_factory(model: nn.Module, comp_config: Any = None, **kwargs: Any) -> nn.Module:
    """Huffman compressor factory function."""
    bits = kwargs.pop(
        "quantize_bits",
        kwargs.pop(
            "bits",
            _cfg_get(comp_config, "huffman_bits", _cfg_get(comp_config, "quantize_bits", 8)),
        ),
    )
    max_error = kwargs.pop(
        "max_error",
        _cfg_get(comp_config, "huffman_max_error", _cfg_get(comp_config, "max_error", 0.20)),
    )
    target_attention_only = kwargs.pop(
        "target_attention_only",
        _cfg_get(
            comp_config,
            "huffman_target_attention_only",
            _cfg_get(comp_config, "target_attention_only", False),
        ),
    )
    fallback = kwargs.pop(
        "fallback_to_original",
        _cfg_get(
            comp_config,
            "huffman_fallback_to_original",
            _cfg_get(comp_config, "fallback_to_original", True),
        ),
    )
    compressor = HuffmanCompressor(
        model,
        quantize_bits=bits,
        max_error=max_error,
        target_attention_only=target_attention_only,
        fallback_to_original=fallback,
        cfg=comp_config,
        **kwargs,
    )
    compressed_model, _ = compressor.compress()
    return compressed_model

Huffman compressor factory function.

Classes

class HuffmanCoder

Expand source code

class HuffmanCoder:
    """
    Vectorized lossless entropy coder for quantized weight and KV cache tensors.
    Supports continuous tensor quantization (e.g. 4-bit / 8-bit), symbol frequency modeling,
    binary prefix tree construction, numpy bitstream packing, and vectorized decoding.
    """

    @staticmethod
    def build_tree(frequencies: dict[int, int]) -> HuffmanNode | None:
        """Build Huffman binary prefix tree from symbol frequencies."""
        if not frequencies:
            return None

        heap: list[HuffmanNode] = [HuffmanNode(sym, freq) for sym, freq in frequencies.items()]
        heapq.heapify(heap)

        if len(heap) == 1:
            root = HuffmanNode(freq=heap[0].freq, left=heap[0])
            return root

        while len(heap) > 1:
            node1 = heapq.heappop(heap)
            node2 = heapq.heappop(heap)
            merged = HuffmanNode(freq=node1.freq + node2.freq, left=node1, right=node2)
            heapq.heappush(heap, merged)

        return heap[0]

    @staticmethod
    def generate_code_table(
        root: HuffmanNode | None,
        current_code: str = "",
        code_table: dict[int, str] | None = None,
    ) -> dict[int, str]:
        """Traverse Huffman tree to generate prefix bitstring per symbol."""
        if code_table is None:
            code_table = {}
        if root is None:
            return code_table

        if root.symbol is not None:
            code_table[root.symbol] = current_code if current_code else "0"
            return code_table

        if root.left is not None:
            HuffmanCoder.generate_code_table(root.left, current_code + "0", code_table)
        if root.right is not None:
            HuffmanCoder.generate_code_table(root.right, current_code + "1", code_table)

        return code_table

    @staticmethod
    def encode(
        tensor: torch.Tensor,
        quantize_bits: int | None = 8,
    ) -> tuple[bytes, dict[int, str], dict[str, Any]]:
        """
        Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.
        """
        flat = tensor.detach().float().view(-1)
        numel = flat.numel()
        if numel == 0:
            return b"", {}, {"total_elements": 0, "entropy": 0.0, "compressed_bytes": 0}

        min_val = flat.min().item()
        max_val = flat.max().item()

        # Handle high quantize_bits properly: skip quantization / identity mapping
        if quantize_bits is None or quantize_bits >= 16:
            if tensor.dtype == torch.bfloat16:
                raw_bytes = tensor.detach().cpu().view(torch.int16).numpy().tobytes()
            else:
                raw_bytes = tensor.detach().cpu().numpy().tobytes()
            metadata = {
                "scale": 1.0,
                "min_val": 0.0,
                "orig_shape": list(tensor.shape),
                "orig_dtype": str(tensor.dtype),
                "total_elements": numel,
                "entropy": 0.0,
                "pad_len": 0,
                "compressed_bytes": len(raw_bytes),
                "is_raw": True,
            }
            return raw_bytes, {}, metadata

        # Quantization
        num_levels = 2 ** int(quantize_bits)
        scale = (max_val - min_val) / max(num_levels - 1, 1)
        scale = max(scale, 1e-12)
        quantized = torch.clamp(torch.round((flat - min_val) / scale), 0, num_levels - 1).long()

        q_np = quantized.cpu().numpy()
        symbols, counts = np.unique(q_np, return_counts=True)
        frequencies = {int(s): int(c) for s, c in zip(symbols, counts)}

        root = HuffmanCoder.build_tree(frequencies)
        code_table = HuffmanCoder.generate_code_table(root)

        probs = counts / numel
        entropy = -float(np.sum(probs * np.log2(probs + 1e-12)))

        # Vectorized bitstream packing using LUT mapping and numpy.packbits
        max_sym = int(symbols.max())
        lut = [code_table.get(i, "0") for i in range(max_sym + 1)]
        bit_chunks = [lut[v] for v in q_np]
        bitstream_str = "".join(bit_chunks)
        pad_len = (8 - len(bitstream_str) % 8) % 8
        padded = bitstream_str + "0" * pad_len

        # Fast bit packing using np.packbits
        bit_arr = np.frombuffer(padded.encode("ascii"), dtype=np.uint8) - ord("0")
        packed_np = np.packbits(bit_arr)
        byte_array = packed_np.tobytes()

        metadata = {
            "scale": scale,
            "min_val": min_val,
            "orig_shape": list(tensor.shape),
            "orig_dtype": str(tensor.dtype),
            "total_elements": numel,
            "entropy": round(entropy, 3),
            "pad_len": pad_len,
            "compressed_bytes": len(byte_array),
            "is_raw": False,
        }
        return byte_array, code_table, metadata

    @staticmethod
    def decode(
        packed_bytes: bytes,
        code_table: dict[int, str],
        metadata: dict[str, Any],
        device: torch.device | None = None,
    ) -> torch.Tensor:
        """Vectorized decode bitstream back into reconstructed tensor."""
        total_elements = metadata.get("total_elements", 0)
        if total_elements == 0:
            return torch.empty(metadata.get("orig_shape", [0]), device=device)

        if metadata.get("is_raw", False):
            orig_dtype_str = metadata.get("orig_dtype", "torch.float32")
            target_dtype = getattr(torch, orig_dtype_str.replace("torch.", ""), torch.float32)
            if target_dtype == torch.bfloat16:
                arr = np.frombuffer(packed_bytes, dtype=np.int16)
                tensor_out = (
                    torch.from_numpy(arr.copy())
                    .view(torch.bfloat16)
                    .reshape(metadata.get("orig_shape"))
                )
            else:
                dtype_str = orig_dtype_str.replace("torch.", "")
                np_dtype = getattr(np, dtype_str if hasattr(np, dtype_str) else "float32")
                arr = np.frombuffer(packed_bytes, dtype=np_dtype)
                tensor_out = (
                    torch.from_numpy(arr.copy())
                    .reshape(metadata.get("orig_shape"))
                    .to(dtype=target_dtype)
                )
            return tensor_out.to(device=device)

        scale = metadata.get("scale", 1.0)
        min_val = metadata.get("min_val", 0.0)
        shape = metadata.get("orig_shape", [total_elements])

        # Vectorized unpacking of bits
        pad_len = metadata.get("pad_len", 0)
        raw_arr = np.frombuffer(packed_bytes, dtype=np.uint8)
        bits = np.unpackbits(raw_arr)
        if pad_len > 0:
            bits = bits[:-pad_len]

        # Fast flat-array binary trie traversal
        node_left = [-1]
        node_right = [-1]
        node_sym = [-1]

        for sym, code in code_table.items():
            curr = 0
            for b in code:
                if b == "0":
                    if node_left[curr] == -1:
                        node_left[curr] = len(node_left)
                        node_left.append(-1)
                        node_right.append(-1)
                        node_sym.append(-1)
                    curr = node_left[curr]
                else:
                    if node_right[curr] == -1:
                        node_right[curr] = len(node_right)
                        node_left.append(-1)
                        node_right.append(-1)
                        node_sym.append(-1)
                    curr = node_right[curr]
            node_sym[curr] = sym

        nl = np.array(node_left, dtype=np.int32)
        nr = np.array(node_right, dtype=np.int32)
        ns = np.array(node_sym, dtype=np.int32)

        decoded = np.empty(total_elements, dtype=np.float32)
        curr = 0
        idx = 0

        # Fast traversal loop
        for b in bits:
            curr = nr[curr] if b else nl[curr]
            sym = ns[curr]
            if sym != -1:
                decoded[idx] = sym
                idx += 1
                curr = 0
                if idx == total_elements:
                    break

        decoded = decoded * scale + min_val
        return torch.from_numpy(decoded).reshape(shape).to(device=device)

Vectorized lossless entropy coder for quantized weight and KV cache tensors. Supports continuous tensor quantization (e.g. 4-bit / 8-bit), symbol frequency modeling, binary prefix tree construction, numpy bitstream packing, and vectorized decoding.

Static methods

def build_tree(frequencies: dict[int, int]) ‑> HuffmanNode | None

Expand source code

@staticmethod
def build_tree(frequencies: dict[int, int]) -> HuffmanNode | None:
    """Build Huffman binary prefix tree from symbol frequencies."""
    if not frequencies:
        return None

    heap: list[HuffmanNode] = [HuffmanNode(sym, freq) for sym, freq in frequencies.items()]
    heapq.heapify(heap)

    if len(heap) == 1:
        root = HuffmanNode(freq=heap[0].freq, left=heap[0])
        return root

    while len(heap) > 1:
        node1 = heapq.heappop(heap)
        node2 = heapq.heappop(heap)
        merged = HuffmanNode(freq=node1.freq + node2.freq, left=node1, right=node2)
        heapq.heappush(heap, merged)

    return heap[0]

Build Huffman binary prefix tree from symbol frequencies.

def decode(packed_bytes: bytes,
code_table: dict[int, str],
metadata: dict[str, Any],
device: torch.device | None = None) ‑> torch.Tensor

Expand source code

@staticmethod
def decode(
    packed_bytes: bytes,
    code_table: dict[int, str],
    metadata: dict[str, Any],
    device: torch.device | None = None,
) -> torch.Tensor:
    """Vectorized decode bitstream back into reconstructed tensor."""
    total_elements = metadata.get("total_elements", 0)
    if total_elements == 0:
        return torch.empty(metadata.get("orig_shape", [0]), device=device)

    if metadata.get("is_raw", False):
        orig_dtype_str = metadata.get("orig_dtype", "torch.float32")
        target_dtype = getattr(torch, orig_dtype_str.replace("torch.", ""), torch.float32)
        if target_dtype == torch.bfloat16:
            arr = np.frombuffer(packed_bytes, dtype=np.int16)
            tensor_out = (
                torch.from_numpy(arr.copy())
                .view(torch.bfloat16)
                .reshape(metadata.get("orig_shape"))
            )
        else:
            dtype_str = orig_dtype_str.replace("torch.", "")
            np_dtype = getattr(np, dtype_str if hasattr(np, dtype_str) else "float32")
            arr = np.frombuffer(packed_bytes, dtype=np_dtype)
            tensor_out = (
                torch.from_numpy(arr.copy())
                .reshape(metadata.get("orig_shape"))
                .to(dtype=target_dtype)
            )
        return tensor_out.to(device=device)

    scale = metadata.get("scale", 1.0)
    min_val = metadata.get("min_val", 0.0)
    shape = metadata.get("orig_shape", [total_elements])

    # Vectorized unpacking of bits
    pad_len = metadata.get("pad_len", 0)
    raw_arr = np.frombuffer(packed_bytes, dtype=np.uint8)
    bits = np.unpackbits(raw_arr)
    if pad_len > 0:
        bits = bits[:-pad_len]

    # Fast flat-array binary trie traversal
    node_left = [-1]
    node_right = [-1]
    node_sym = [-1]

    for sym, code in code_table.items():
        curr = 0
        for b in code:
            if b == "0":
                if node_left[curr] == -1:
                    node_left[curr] = len(node_left)
                    node_left.append(-1)
                    node_right.append(-1)
                    node_sym.append(-1)
                curr = node_left[curr]
            else:
                if node_right[curr] == -1:
                    node_right[curr] = len(node_right)
                    node_left.append(-1)
                    node_right.append(-1)
                    node_sym.append(-1)
                curr = node_right[curr]
        node_sym[curr] = sym

    nl = np.array(node_left, dtype=np.int32)
    nr = np.array(node_right, dtype=np.int32)
    ns = np.array(node_sym, dtype=np.int32)

    decoded = np.empty(total_elements, dtype=np.float32)
    curr = 0
    idx = 0

    # Fast traversal loop
    for b in bits:
        curr = nr[curr] if b else nl[curr]
        sym = ns[curr]
        if sym != -1:
            decoded[idx] = sym
            idx += 1
            curr = 0
            if idx == total_elements:
                break

    decoded = decoded * scale + min_val
    return torch.from_numpy(decoded).reshape(shape).to(device=device)

Vectorized decode bitstream back into reconstructed tensor.

def encode(tensor: torch.Tensor, quantize_bits: int | None = 8) ‑> tuple[bytes, dict[int, str], dict[str, typing.Any]]

Expand source code

@staticmethod
def encode(
    tensor: torch.Tensor,
    quantize_bits: int | None = 8,
) -> tuple[bytes, dict[int, str], dict[str, Any]]:
    """
    Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.
    """
    flat = tensor.detach().float().view(-1)
    numel = flat.numel()
    if numel == 0:
        return b"", {}, {"total_elements": 0, "entropy": 0.0, "compressed_bytes": 0}

    min_val = flat.min().item()
    max_val = flat.max().item()

    # Handle high quantize_bits properly: skip quantization / identity mapping
    if quantize_bits is None or quantize_bits >= 16:
        if tensor.dtype == torch.bfloat16:
            raw_bytes = tensor.detach().cpu().view(torch.int16).numpy().tobytes()
        else:
            raw_bytes = tensor.detach().cpu().numpy().tobytes()
        metadata = {
            "scale": 1.0,
            "min_val": 0.0,
            "orig_shape": list(tensor.shape),
            "orig_dtype": str(tensor.dtype),
            "total_elements": numel,
            "entropy": 0.0,
            "pad_len": 0,
            "compressed_bytes": len(raw_bytes),
            "is_raw": True,
        }
        return raw_bytes, {}, metadata

    # Quantization
    num_levels = 2 ** int(quantize_bits)
    scale = (max_val - min_val) / max(num_levels - 1, 1)
    scale = max(scale, 1e-12)
    quantized = torch.clamp(torch.round((flat - min_val) / scale), 0, num_levels - 1).long()

    q_np = quantized.cpu().numpy()
    symbols, counts = np.unique(q_np, return_counts=True)
    frequencies = {int(s): int(c) for s, c in zip(symbols, counts)}

    root = HuffmanCoder.build_tree(frequencies)
    code_table = HuffmanCoder.generate_code_table(root)

    probs = counts / numel
    entropy = -float(np.sum(probs * np.log2(probs + 1e-12)))

    # Vectorized bitstream packing using LUT mapping and numpy.packbits
    max_sym = int(symbols.max())
    lut = [code_table.get(i, "0") for i in range(max_sym + 1)]
    bit_chunks = [lut[v] for v in q_np]
    bitstream_str = "".join(bit_chunks)
    pad_len = (8 - len(bitstream_str) % 8) % 8
    padded = bitstream_str + "0" * pad_len

    # Fast bit packing using np.packbits
    bit_arr = np.frombuffer(padded.encode("ascii"), dtype=np.uint8) - ord("0")
    packed_np = np.packbits(bit_arr)
    byte_array = packed_np.tobytes()

    metadata = {
        "scale": scale,
        "min_val": min_val,
        "orig_shape": list(tensor.shape),
        "orig_dtype": str(tensor.dtype),
        "total_elements": numel,
        "entropy": round(entropy, 3),
        "pad_len": pad_len,
        "compressed_bytes": len(byte_array),
        "is_raw": False,
    }
    return byte_array, code_table, metadata

Vectorized encode input tensor into packed bitstream bytes and Huffman codebook.

def generate_code_table(root: HuffmanNode | None,
current_code: str = '',
code_table: dict[int, str] | None = None) ‑> dict[int, str]

Expand source code

@staticmethod
def generate_code_table(
    root: HuffmanNode | None,
    current_code: str = "",
    code_table: dict[int, str] | None = None,
) -> dict[int, str]:
    """Traverse Huffman tree to generate prefix bitstring per symbol."""
    if code_table is None:
        code_table = {}
    if root is None:
        return code_table

    if root.symbol is not None:
        code_table[root.symbol] = current_code if current_code else "0"
        return code_table

    if root.left is not None:
        HuffmanCoder.generate_code_table(root.left, current_code + "0", code_table)
    if root.right is not None:
        HuffmanCoder.generate_code_table(root.right, current_code + "1", code_table)

    return code_table

Traverse Huffman tree to generate prefix bitstring per symbol.

class HuffmanCompressor (model: nn.Module,
quantize_bits: int = 8,
max_error: float = 0.2,
target_attention_only: bool = False,
fallback_to_original: bool = True,
cfg: Any | None = None,
**kwargs: Any)

Expand source code

@COMPRESSOR_REGISTRY.register("huffman")
@COMPRESSOR_REGISTRY.register("huffman_coder")
class HuffmanCompressor(BaseCompressor):
    """
    HuffmanCompressor: Model-wide entropy compression profiler and simulator.
    """

    def __init__(
        self,
        model: nn.Module,
        quantize_bits: int = 8,
        max_error: float = 0.20,
        target_attention_only: bool = False,
        fallback_to_original: bool = True,
        cfg: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model, name="huffman", **kwargs)
        cfg = cfg or kwargs.get("comp_config")
        self.cfg = cfg

        self.quantize_bits = int(
            _cfg_get(
                cfg,
                "huffman_quantize_bits",
                _cfg_get(cfg, "quantize_bits", _cfg_get(cfg, "bits", quantize_bits)),
            )
        )
        self.max_error = float(
            _cfg_get(cfg, "huffman_max_error", _cfg_get(cfg, "max_error", max_error))
        )
        self.target_attention_only = bool(
            _cfg_get(
                cfg,
                "huffman_target_attention_only",
                _cfg_get(cfg, "target_attention_only", target_attention_only),
            )
        )
        self.fallback_to_original = bool(
            _cfg_get(
                cfg,
                "huffman_fallback_to_original",
                _cfg_get(cfg, "fallback_to_original", fallback_to_original),
            )
        )

        self.compressed_layers: list[str] = []
        self.fallback_layers: list[str] = []
        self.skipped_layers: list[str] = []
        self.layer_errors: dict[str, float] = {}

    def _is_target_module(self, name: str, module: nn.Module) -> bool:
        name_lower = name.lower()
        parts = name.split(".")
        if any(part in parts for part in ["lora_A", "lora_B", "base_layer"]):
            return False
        if any(
            skip in name_lower
            for skip in [
                "classifier",
                "score",
                "lm_head",
                "pooler",
                "embedding",
                "router",
                "gate",
            ]
        ):
            return False

        if self.target_attention_only:
            if not any(attn_kw in name_lower for attn_kw in _ATTN_NAMES):
                return False

        target_mod = getattr(module, "base_layer", module)
        cls_name = target_mod.__class__.__name__.lower()
        is_linear = isinstance(target_mod, nn.Linear)
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        if not (is_linear or is_conv1d):
            return False
        if not should_compress_linear(self.model, name):
            return False
        return True

    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Profile and apply Huffman entropy coding across all linear layers."""
        first_p = next(self.model.parameters(), None)
        _model_device = first_p.device if first_p is not None else torch.device("cpu")

        total_orig_bytes = 0
        total_comp_bytes = 0
        total_entropy = 0.0

        self.compressed_layers = []
        self.fallback_layers = []
        self.skipped_layers = []
        self.layer_errors = {}
        seen_weights: set[int] = set()

        # Freeze skipped non-target layers if target_attention_only is enabled
        if self.target_attention_only:
            for name, module in self.model.named_modules():
                if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                    continue
                target_mod = getattr(module, "base_layer", module)
                cls_name = target_mod.__class__.__name__.lower()
                is_linear = isinstance(target_mod, nn.Linear)
                is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
                if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                    _freeze_base_only(module)
                    self.skipped_layers.append(name)

        for name, module in self.model.named_modules():
            if not self._is_target_module(name, module):
                continue

            target_mod = getattr(module, "base_layer", module)
            if not hasattr(target_mod, "weight") or target_mod.weight is None:
                continue

            w_id = id(target_mod.weight)
            if w_id in seen_weights:
                self.skipped_layers.append(name)
                continue
            seen_weights.add(w_id)

            cls_name = target_mod.__class__.__name__.lower()
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            raw_w = target_mod.weight.data
            orig_shape = raw_w.shape
            orig_device = raw_w.device
            orig_dtype = raw_w.dtype

            # Byte accounting based on actual element_size()
            elem_bytes = raw_w.element_size() if hasattr(raw_w, "element_size") else 2
            layer_orig_bytes = raw_w.numel() * elem_bytes

            W_orig = raw_w.t().float() if is_conv1d else raw_w.float()
            norm_w = torch.norm(W_orig, p="fro").item()

            try:
                packed_b, code_map, meta = HuffmanCoder.encode(
                    W_orig, quantize_bits=self.quantize_bits
                )
                W_rec = HuffmanCoder.decode(packed_b, code_map, meta, device=orig_device).to(
                    dtype=torch.float32
                )

                if norm_w < 1e-12:
                    err = 0.0
                else:
                    err = float(torch.norm(W_orig - W_rec, p="fro").item() / norm_w)

                self.layer_errors[name] = err

                d_out, d_in = W_orig.shape
                fam = getattr(getattr(self.model, "config", None), "model_type", None)
                eff_max_err = max(self.max_error, adaptive_max_error(None, d_in, d_out, family=fam))

                if err > eff_max_err:
                    logger.warning(
                        "Huffman error %.4f on %s exceeded max_error %.4f (base max_err=%.4f)",
                        err,
                        name,
                        eff_max_err,
                        self.max_error,
                    )
                    if self.fallback_to_original:
                        self.fallback_layers.append(name)
                        _freeze_base_only(module)
                        total_orig_bytes += layer_orig_bytes
                        total_comp_bytes += layer_orig_bytes
                        continue

                # Apply reconstructed weight
                if is_conv1d:
                    res_w = W_rec.t().to(dtype=orig_dtype, device=orig_device)
                else:
                    res_w = W_rec.to(dtype=orig_dtype, device=orig_device)

                assert res_w.shape == orig_shape, (
                    f"Shape mismatch on {name}: {res_w.shape} vs {orig_shape}"
                )
                target_mod.weight.data.copy_(res_w)

                # Preserve LoRA adapter trainability
                for lora_attr in ("lora_A", "lora_B", "lora_a", "lora_b"):
                    if hasattr(module, lora_attr):
                        l_mod = getattr(module, lora_attr)
                        if (
                            isinstance(l_mod, nn.Module)
                            and hasattr(l_mod, "weight")
                            and l_mod.weight is not None
                        ):
                            l_mod.weight.requires_grad = True
                        elif isinstance(l_mod, nn.Parameter):
                            l_mod.requires_grad = True

                self.compressed_layers.append(name)
                total_orig_bytes += layer_orig_bytes
                total_comp_bytes += meta["compressed_bytes"]
                total_entropy += meta["entropy"]

            except Exception as e:
                logger.debug("Huffman exception on %s: %s", name, e)
                self.fallback_layers.append(name)
                _freeze_base_only(module)
                total_orig_bytes += layer_orig_bytes
                total_comp_bytes += layer_orig_bytes

        # Ensure all parameters on target device
        self.model.to(device=_model_device)

        forward_check_success = None
        try:
            with torch.no_grad():
                dev = next(self.model.parameters(), torch.tensor(0.0)).device
                tokenizer = getattr(self.model, "tokenizer", None) or getattr(
                    getattr(self.model, "config", None), "tokenizer", None
                )
                if tokenizer is not None:
                    tok_out = tokenizer("Hello world", return_tensors="pt")
                    tok_out = {k: v.to(dev) for k, v in tok_out.items()}
                    self.model(**tok_out)
                    forward_check_success = True
                else:
                    # Look for token embeddings
                    has_embed = any(
                        any(kw in n.lower() for kw in ("embed_tokens", "wte", "word_embeddings"))
                        for n, _ in self.model.named_modules()
                    )
                    if has_embed:
                        dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                        dummy_mask = torch.ones((1, 8), device=dev)
                        try:
                            self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
                        except TypeError:
                            self.model(dummy_ids)
                        forward_check_success = True
                    else:
                        first_linear = next(
                            (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                            None,
                        )
                        if first_linear is not None:
                            in_dim = getattr(first_linear, "in_features", 8)
                            self.model(torch.randn(1, 8, in_dim, device=dev))
                            forward_check_success = True
                        else:
                            forward_check_success = None
        except Exception as fwd_exc:
            logger.debug("Huffman forward check failed: %s", fwd_exc)
            forward_check_success = None

        stats = self.get_compression_stats()
        avg_err = (
            (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
            if self.layer_errors
            else 0.0
        )
        stats["quantize_bits"] = self.quantize_bits
        stats["avg_entropy_bits"] = round(total_entropy / max(len(self.compressed_layers), 1), 2)
        stats["total_compressed_bytes"] = total_comp_bytes
        stats["huffman_compression_ratio"] = round(total_orig_bytes / max(total_comp_bytes, 1), 2)
        stats["theoretical_compression_ratio"] = stats["huffman_compression_ratio"]
        stats["num_compressed"] = len(self.compressed_layers)
        stats["num_fallback"] = len(self.fallback_layers)
        stats["num_skipped"] = len(self.skipped_layers)
        stats["compressed_layers"] = self.compressed_layers
        stats["fallback_layers"] = self.fallback_layers
        stats["skipped_layers"] = self.skipped_layers
        stats["avg_reconstruction_error"] = avg_err
        stats["forward_check_success"] = forward_check_success
        logger.info(
            "Compressor huffman finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
            len(self.compressed_layers),
            len(self.fallback_layers),
            avg_err,
            self.max_error,
        )
        return self.model, stats

HuffmanCompressor: Model-wide entropy compression profiler and simulator.

Ancestors

• BaseCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Profile and apply Huffman entropy coding across all linear layers."""
    first_p = next(self.model.parameters(), None)
    _model_device = first_p.device if first_p is not None else torch.device("cpu")

    total_orig_bytes = 0
    total_comp_bytes = 0
    total_entropy = 0.0

    self.compressed_layers = []
    self.fallback_layers = []
    self.skipped_layers = []
    self.layer_errors = {}
    seen_weights: set[int] = set()

    # Freeze skipped non-target layers if target_attention_only is enabled
    if self.target_attention_only:
        for name, module in self.model.named_modules():
            if any(part in name.split(".") for part in ["lora_A", "lora_B", "base_layer"]):
                continue
            target_mod = getattr(module, "base_layer", module)
            cls_name = target_mod.__class__.__name__.lower()
            is_linear = isinstance(target_mod, nn.Linear)
            is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
            if (is_linear or is_conv1d) and not self._is_target_module(name, module):
                _freeze_base_only(module)
                self.skipped_layers.append(name)

    for name, module in self.model.named_modules():
        if not self._is_target_module(name, module):
            continue

        target_mod = getattr(module, "base_layer", module)
        if not hasattr(target_mod, "weight") or target_mod.weight is None:
            continue

        w_id = id(target_mod.weight)
        if w_id in seen_weights:
            self.skipped_layers.append(name)
            continue
        seen_weights.add(w_id)

        cls_name = target_mod.__class__.__name__.lower()
        is_conv1d = "conv1d" in cls_name or hasattr(target_mod, "nf")
        raw_w = target_mod.weight.data
        orig_shape = raw_w.shape
        orig_device = raw_w.device
        orig_dtype = raw_w.dtype

        # Byte accounting based on actual element_size()
        elem_bytes = raw_w.element_size() if hasattr(raw_w, "element_size") else 2
        layer_orig_bytes = raw_w.numel() * elem_bytes

        W_orig = raw_w.t().float() if is_conv1d else raw_w.float()
        norm_w = torch.norm(W_orig, p="fro").item()

        try:
            packed_b, code_map, meta = HuffmanCoder.encode(
                W_orig, quantize_bits=self.quantize_bits
            )
            W_rec = HuffmanCoder.decode(packed_b, code_map, meta, device=orig_device).to(
                dtype=torch.float32
            )

            if norm_w < 1e-12:
                err = 0.0
            else:
                err = float(torch.norm(W_orig - W_rec, p="fro").item() / norm_w)

            self.layer_errors[name] = err

            d_out, d_in = W_orig.shape
            fam = getattr(getattr(self.model, "config", None), "model_type", None)
            eff_max_err = max(self.max_error, adaptive_max_error(None, d_in, d_out, family=fam))

            if err > eff_max_err:
                logger.warning(
                    "Huffman error %.4f on %s exceeded max_error %.4f (base max_err=%.4f)",
                    err,
                    name,
                    eff_max_err,
                    self.max_error,
                )
                if self.fallback_to_original:
                    self.fallback_layers.append(name)
                    _freeze_base_only(module)
                    total_orig_bytes += layer_orig_bytes
                    total_comp_bytes += layer_orig_bytes
                    continue

            # Apply reconstructed weight
            if is_conv1d:
                res_w = W_rec.t().to(dtype=orig_dtype, device=orig_device)
            else:
                res_w = W_rec.to(dtype=orig_dtype, device=orig_device)

            assert res_w.shape == orig_shape, (
                f"Shape mismatch on {name}: {res_w.shape} vs {orig_shape}"
            )
            target_mod.weight.data.copy_(res_w)

            # Preserve LoRA adapter trainability
            for lora_attr in ("lora_A", "lora_B", "lora_a", "lora_b"):
                if hasattr(module, lora_attr):
                    l_mod = getattr(module, lora_attr)
                    if (
                        isinstance(l_mod, nn.Module)
                        and hasattr(l_mod, "weight")
                        and l_mod.weight is not None
                    ):
                        l_mod.weight.requires_grad = True
                    elif isinstance(l_mod, nn.Parameter):
                        l_mod.requires_grad = True

            self.compressed_layers.append(name)
            total_orig_bytes += layer_orig_bytes
            total_comp_bytes += meta["compressed_bytes"]
            total_entropy += meta["entropy"]

        except Exception as e:
            logger.debug("Huffman exception on %s: %s", name, e)
            self.fallback_layers.append(name)
            _freeze_base_only(module)
            total_orig_bytes += layer_orig_bytes
            total_comp_bytes += layer_orig_bytes

    # Ensure all parameters on target device
    self.model.to(device=_model_device)

    forward_check_success = None
    try:
        with torch.no_grad():
            dev = next(self.model.parameters(), torch.tensor(0.0)).device
            tokenizer = getattr(self.model, "tokenizer", None) or getattr(
                getattr(self.model, "config", None), "tokenizer", None
            )
            if tokenizer is not None:
                tok_out = tokenizer("Hello world", return_tensors="pt")
                tok_out = {k: v.to(dev) for k, v in tok_out.items()}
                self.model(**tok_out)
                forward_check_success = True
            else:
                # Look for token embeddings
                has_embed = any(
                    any(kw in n.lower() for kw in ("embed_tokens", "wte", "word_embeddings"))
                    for n, _ in self.model.named_modules()
                )
                if has_embed:
                    dummy_ids = torch.randint(0, 100, (1, 8), device=dev)
                    dummy_mask = torch.ones((1, 8), device=dev)
                    try:
                        self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
                    except TypeError:
                        self.model(dummy_ids)
                    forward_check_success = True
                else:
                    first_linear = next(
                        (m for m in self.model.modules() if isinstance(m, nn.Linear)),
                        None,
                    )
                    if first_linear is not None:
                        in_dim = getattr(first_linear, "in_features", 8)
                        self.model(torch.randn(1, 8, in_dim, device=dev))
                        forward_check_success = True
                    else:
                        forward_check_success = None
    except Exception as fwd_exc:
        logger.debug("Huffman forward check failed: %s", fwd_exc)
        forward_check_success = None

    stats = self.get_compression_stats()
    avg_err = (
        (sum(self.layer_errors.values()) / max(1, len(self.layer_errors)))
        if self.layer_errors
        else 0.0
    )
    stats["quantize_bits"] = self.quantize_bits
    stats["avg_entropy_bits"] = round(total_entropy / max(len(self.compressed_layers), 1), 2)
    stats["total_compressed_bytes"] = total_comp_bytes
    stats["huffman_compression_ratio"] = round(total_orig_bytes / max(total_comp_bytes, 1), 2)
    stats["theoretical_compression_ratio"] = stats["huffman_compression_ratio"]
    stats["num_compressed"] = len(self.compressed_layers)
    stats["num_fallback"] = len(self.fallback_layers)
    stats["num_skipped"] = len(self.skipped_layers)
    stats["compressed_layers"] = self.compressed_layers
    stats["fallback_layers"] = self.fallback_layers
    stats["skipped_layers"] = self.skipped_layers
    stats["avg_reconstruction_error"] = avg_err
    stats["forward_check_success"] = forward_check_success
    logger.info(
        "Compressor huffman finished: %d compressed, %d fallback, avg_err=%.4f, max_err_threshold=%.4f",
        len(self.compressed_layers),
        len(self.fallback_layers),
        avg_err,
        self.max_error,
    )
    return self.model, stats

Profile and apply Huffman entropy coding across all linear layers.

Inherited members

• BaseCompressor:
  • get_compression_stats
  • summary

class HuffmanNode (symbol: int | None = None,
freq: int = 0,
left: HuffmanNode | None = None,
right: HuffmanNode | None = None)

Expand source code

class HuffmanNode:
    """Node in the Huffman prefix tree."""

    def __init__(
        self,
        symbol: int | None = None,
        freq: int = 0,
        left: HuffmanNode | None = None,
        right: HuffmanNode | None = None,
    ) -> None:
        self.symbol = symbol
        self.freq = freq
        self.left = left
        self.right = right

    def __lt__(self, other: HuffmanNode) -> bool:
        return self.freq < other.freq

Node in the Huffman prefix tree.