src/Sub1BitLLM.py

import torch
import torch.nn as nn
from dataclasses import dataclass
from typing import Dict, Optional, List, Iterator, Tuple, Any
from pathlib import Path
import os
import numpy as np

from .gguf_writer import GGUFWriter, GGML_TYPES
from .quantization import dequantize_factor


@dataclass
class Sub1BitConfig:
    codebook_dim: int = 128
    energy_threshold: float = 0.95
    rank: int = 16
    U_bits: float = 0.5
    S_bits: float = 2.0
    Vt_bits: float = 0.5
    model_name: str = "llama-2-7b-sub1bit"
    architecture: str = "llama"


class LowRankFactor(torch.nn.Module):
    def __init__(self, U: torch.Tensor, S: torch.Tensor, Vt: torch.Tensor):
        super().__init__()
        self.register_buffer('U', U.half())
        self.register_buffer('S', S.half())
        self.register_buffer('Vt', Vt.half())
        self.rank = U.shape[1]

    def forward(self) -> torch.Tensor:
        return torch.matmul(self.U * self.S.unsqueeze(0), self.Vt)

    def forward_lowrank(self, x: torch.Tensor) -> torch.Tensor:
        return torch.matmul(torch.matmul(x, self.Vt.T) * self.S.unsqueeze(0), self.U.T)


class TernaryQuantizedFactor(torch.nn.Module):
    def __init__(self, data: torch.Tensor, scale: torch.Tensor, rank: int):
        super().__init__()
        self.register_buffer('data', data)
        self.register_buffer('scale', scale)
        self.rank = rank

    def forward(self) -> torch.Tensor:
        return self.data.float() * self.scale


class Sub1BitLLM(torch.nn.Module):
    def __init__(
        self,
        model_path: str,
        config: Optional[Sub1BitConfig] = None,
        device: str = "cuda" if torch.cuda.is_available() else "cpu"
    ):
        super().__init__()
        self.model_path = model_path
        self.config = config or Sub1BitConfig()
        self.device = device
        self.layers: Dict[int, LowRankFactor] = {}
        self.metadata: Dict = {}

    @classmethod
    def from_fp16(
        cls,
        model_path: str,
        config: Optional[Sub1BitConfig] = None,
        checkpoint_dir: Optional[str] = None,
        device: str = "cuda" if torch.cuda.is_available() else "cpu"
    ) -> "Sub1BitLLM":
        instance = cls(model_path, config, device)
        if checkpoint_dir is None:
            checkpoint_dir = Path(model_path).parent / "checkpoints"
        else:
            checkpoint_dir = Path(checkpoint_dir)
        if not checkpoint_dir.exists():
            raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")
        for ckpt_file in sorted(checkpoint_dir.glob("layer_*.pt")):
            layer_idx = int(ckpt_file.stem.split("_")[1])
            factor = torch.load(ckpt_file, weights_only=False, map_location=device)
            U = torch.from_numpy(factor['U']).to(device)
            S = torch.from_numpy(factor['S']).to(device)
            Vt = torch.from_numpy(factor['Vt']).to(device)
            instance.layers[layer_idx] = LowRankFactor(U, S, Vt)
        instance.metadata = {
            'num_layers': len(instance.layers),
            'rank': instance.config.rank,
            'energy_threshold': instance.config.energy_threshold
        }
        return instance

    def load_checkpoint(self, checkpoint_path: str) -> "Sub1BitLLM":
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        if 'layers' in checkpoint:
            for layer_idx, factor_data in checkpoint['layers'].items():
                self.layers[int(layer_idx)] = LowRankFactor(
                    torch.from_numpy(factor_data['U']).to(self.device),
                    torch.from_numpy(factor_data['S']).to(self.device),
                    torch.from_numpy(factor_data['Vt']).to(self.device),
                )
        return self

    def state_dict(self) -> Dict[str, torch.Tensor]:
        state = {}
        for layer_idx, layer in self.layers.items():
            state[f'layers.{layer_idx}.U'] = layer.U
            state[f'layers.{layer_idx}.S'] = layer.S
            state[f'layers.{layer_idx}.Vt'] = layer.Vt
        return state

    def forward(self, x: torch.Tensor, layer_indices: Optional[List[int]] = None) -> Dict[int, torch.Tensor]:
        outputs = {}
        indices = layer_indices if layer_indices is not None else list(self.layers.keys())
        for idx in indices:
            if idx in self.layers:
                outputs[idx] = self.layers[idx].forward_lowrank(x)
        return outputs

    def get_weight(self, layer_idx: int) -> torch.Tensor:
        if layer_idx not in self.layers:
            raise KeyError(f"Layer {layer_idx} not found")
        return self.layers[layer_idx]()

    def iter_layers(self) -> Iterator[Tuple[int, LowRankFactor]]:
        for idx in sorted(self.layers.keys()):
            yield idx, self.layers[idx]

    def compression_stats(self) -> Dict[str, float]:
        total_original = 0
        total_factor = 0
        for _, layer in self.iter_layers():
            orig_size = layer.U.shape[0] * layer.Vt.shape[1]
            factor_size = layer.U.numel() + layer.S.numel() + layer.Vt.numel()
            total_original += orig_size
            total_factor += factor_size
        return {
            'compression_ratio': total_original / total_factor if total_factor > 0 else 0,
            'avg_rank': sum(l.rank for _, l in self.iter_layers()) / max(len(self.layers), 1)
        }

    def to_gguf(self, output_path: str, metadata: Optional[Dict] = None):
        writer = GGUFWriter(output_path)
        writer.add_key_value("general.architecture", self.config.architecture)
        writer.add_key_value("general.name", self.config.model_name)
        writer.add_key_value("quantization.type", "sub1bit_lowrank")
        writer.add_key_value("quantization.U_bits", self.config.U_bits)
        writer.add_key_value("quantization.S_bits", self.config.S_bits)
        writer.add_key_value("quantization.Vt_bits", self.config.Vt_bits)
        if metadata:
            for key, value in metadata.items():
                writer.add_key_value(key, value)
        for layer_idx, layer in self.iter_layers():
            writer.add_tensor(
                f"model.layers.{layer_idx}.U",
                layer.U.cpu().numpy().astype(np.float16),
                GGML_TYPES['float16']
            )
            writer.add_tensor(
                f"model.layers.{layer_idx}.S",
                layer.S.cpu().numpy().astype(np.float16),
                GGML_TYPES['float16']
            )
            writer.add_tensor(
                f"model.layers.{layer_idx}.Vt",
                layer.Vt.cpu().numpy().astype(np.float16),
                GGML_TYPES['float16']
            )
            writer.add_tensor(
                f"model.layers.{layer_idx}.rank",
                np.array([layer.rank], dtype=np.int32),
                GGML_TYPES['int32']
            )
        writer.write()
        return os.path.getsize(output_path)


def from_fp16(
    model_path: str,
    config: Optional[Sub1BitConfig] = None,
    checkpoint_dir: Optional[str] = None,
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
) -> Sub1BitLLM:
    return Sub1BitLLM.from_fp16(model_path, config, checkpoint_dir, device)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Sub1BitLLM API Demo")
    parser.add_argument("--model", type=str, required=True, help="Path to model weights")
    parser.add_argument("--checkpoint-dir", type=str, default=None, help="Path to checkpoint directory")
    parser.add_argument("--device", type=str, default="cuda", help="Device")
    args = parser.parse_args()
    config = Sub1BitConfig(
        codebook_dim=128,
        energy_threshold=0.95,
        model_name="llama-2-7b-sub1bit"
    )
    model = from_fp16(args.model, config=config, checkpoint_dir=args.checkpoint_dir, device=args.device)
    print(f"Loaded Sub1BitLLM with {len(model.layers)} layers")
    print(f"Compression stats: {model.compression_stats()}")