#!/usr/bin/env python3 """ Q-TensorFormer Drop-In LLM Compressor. Converts dense feed-forward and projection layers in any standard Hugging Face Causal LM (Llama-3, Qwen-2.5, Mistral, TinyLlama, etc.) into Tensor-Train contracted layers with dynamic zero-SVD rank slicing and closed-loop PID control. Usage: python qtensorformer_compress.py \ --model-name-or-path meta-llama/Llama-3.2-1B \ --output-dir ./Llama-3.2-1B-QTensor \ --tt-rank 16 """ import argparse import math import os import sys from typing import List, Tuple import torch import torch.nn as nn from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from configuration_qtensorformer import QTensorFormerConfig, factorize_dim from modeling_qtensorformer import QTensorFormerForCausalLM, TensorTrainLinear def svd_factorize_matrix_to_tt( weight: torch.Tensor, in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], target_ranks: List[int], ) -> List[torch.Tensor]: """ Factorizes a dense 2D weight matrix W [in_features, out_features] into TT-cores via sequential truncated Singular Value Decompositions (TT-SVD algorithm). """ device = weight.device num_cores = len(in_shape) assert len(out_shape) == num_cores # Reshape W into tensor of shape (m_1, n_1, m_2, n_2, ..., m_d, n_d) w_tensor = weight.reshape(*in_shape, *out_shape) # Permute to interleave: (m_1, n_1, m_2, n_2, ...) perm = [] for i in range(num_cores): perm.append(i) perm.append(num_cores + i) w_tensor = w_tensor.permute(*perm) curr_matrix = w_tensor.reshape(in_shape[0] * out_shape[0], -1) cores: List[torch.Tensor] = [] r_prev = 1 for k in range(num_cores - 1): m_k = in_shape[k] n_k = out_shape[k] max_r = target_ranks[k + 1] U, S, Vh = torch.linalg.svd(curr_matrix, full_matrices=False) r_curr = min(len(S), max_r) U_trunc = U[:, :r_curr] S_trunc = S[:r_curr] Vh_trunc = Vh[:r_curr, :] # Form core G_k: shape [r_prev, m_k, n_k, r_curr] core_k = U_trunc.reshape(r_prev, m_k, n_k, r_curr) cores.append(core_k) # Remaining matrix for next step curr_matrix = (torch.diag(S_trunc) @ Vh_trunc).reshape( r_curr * in_shape[k + 1] * out_shape[k + 1], -1 ) r_prev = r_curr # Final core m_last = in_shape[-1] n_last = out_shape[-1] core_last = curr_matrix.reshape(r_prev, m_last, n_last, 1) cores.append(core_last) return cores def compress_dense_model_to_qtensorformer( model_name_or_path: str, output_dir: str, max_tt_rank: int = 16, device: str = "cpu", ): print(f"\n[Compressor] Loading source model from: {model_name_or_path}") source_config = AutoConfig.from_pretrained(model_name_or_path) source_model = AutoModelForCausalLM.from_pretrained( model_name_or_path, torch_dtype=torch.float32, device_map=device, ) # Derive QTensorFormer configuration matching the source model architecture hidden_size = getattr(source_config, "hidden_size", 2048) intermediate_size = getattr(source_config, "intermediate_size", 5632) num_layers = getattr(source_config, "num_hidden_layers", 16) num_heads = getattr(source_config, "num_attention_heads", 16) num_kv_heads = getattr(source_config, "num_key_value_heads", num_heads) vocab_size = getattr(source_config, "vocab_size", 32000) tt_ranks = [1, max_tt_rank, max_tt_rank, 1] print(f"[Compressor] Generating target Q-TensorFormer configuration...") target_config = QTensorFormerConfig( vocab_size=vocab_size, hidden_size=hidden_size, num_hidden_layers=num_layers, num_attention_heads=num_heads, num_key_value_heads=num_kv_heads, intermediate_size=intermediate_size, tt_ranks=tt_ranks, max_position_embeddings=getattr(source_config, "max_position_embeddings", 4096), enable_dynamic_routing=True, ) target_model = QTensorFormerForCausalLM(target_config).to(device) # Copy Embeddings & Final LM Head print(f"[Compressor] Transferring token embeddings and classification head...") with torch.no_grad(): if hasattr(source_model, "model") and hasattr(source_model.model, "embed_tokens"): target_model.model.embed_tokens.weight.copy_(source_model.model.embed_tokens.weight) if hasattr(source_model, "lm_head"): target_model.lm_head.weight.copy_(source_model.lm_head.weight) # SVD Factorize MLPs layer-by-layer print(f"[Compressor] Factorizing dense MLP layers into Tensor-Train cores...") for idx in range(min(num_layers, len(target_model.model.layers))): print(f" -> Converting Layer {idx + 1}/{num_layers}...") src_layer = source_model.model.layers[idx] if hasattr(source_model, "model") else None tgt_layer = target_model.model.layers[idx] if src_layer is not None and hasattr(src_layer, "mlp"): with torch.no_grad(): # Factorize gate_proj if hasattr(src_layer.mlp, "gate_proj") and hasattr(tgt_layer.mlp, "gate_proj"): w_gate = src_layer.mlp.gate_proj.weight.data cores_gate = svd_factorize_matrix_to_tt( w_gate, tgt_layer.mlp.gate_proj.in_shape, tgt_layer.mlp.gate_proj.out_shape, tt_ranks, ) for c_idx, c in enumerate(cores_gate): tgt_layer.mlp.gate_proj.cores[c_idx].copy_(c) # Factorize up_proj if hasattr(src_layer.mlp, "up_proj") and hasattr(tgt_layer.mlp, "up_proj"): w_up = src_layer.mlp.up_proj.weight.data cores_up = svd_factorize_matrix_to_tt( w_up, tgt_layer.mlp.up_proj.in_shape, tgt_layer.mlp.up_proj.out_shape, tt_ranks, ) for c_idx, c in enumerate(cores_up): tgt_layer.mlp.up_proj.cores[c_idx].copy_(c) # Factorize down_proj if hasattr(src_layer.mlp, "down_proj") and hasattr(tgt_layer.mlp, "down_proj"): w_down = src_layer.mlp.down_proj.weight.data cores_down = svd_factorize_matrix_to_tt( w_down, tgt_layer.mlp.down_proj.in_shape, tgt_layer.mlp.down_proj.out_shape, tt_ranks, ) for c_idx, c in enumerate(cores_down): tgt_layer.mlp.down_proj.cores[c_idx].copy_(c) # Save compressed model and tokenizer print(f"\n[Compressor] Saving compressed model to {output_dir}...") os.makedirs(output_dir, exist_ok=True) target_model.save_pretrained(output_dir) target_config.save_pretrained(output_dir) try: tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) tokenizer.save_pretrained(output_dir) except Exception as e: print(f"[Compressor] Warning: could not copy tokenizer ({e})") src_params = sum(p.numel() for p in source_model.parameters()) tgt_params = sum(p.numel() for p in target_model.parameters()) print(f"[Compressor] Compression Complete!") print(f" Source Model Parameters: {src_params:,}") print(f" Target Model Parameters: {tgt_params:,}") print(f" Overall Reduction: {(src_params - tgt_params) / src_params * 100:.2f}%") print(f" Model ready for load via: AutoModelForCausalLM.from_pretrained('{output_dir}', trust_remote_code=True)\n") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Q-TensorFormer Drop-in LLM Compressor") parser.add_argument("--model-name-or-path", type=str, required=True, help="HF model ID or local directory") parser.add_argument("--output-dir", type=str, required=True, help="Target directory for compressed model") parser.add_argument("--tt-rank", type=int, default=16, help="Maximum Tensor-Train rank (default: 16)") parser.add_argument("--device", type=str, default="cpu", help="Compute device for SVD (cpu or cuda)") args = parser.parse_args() compress_dense_model_to_qtensorformer( model_name_or_path=args.model_name_or_path, output_dir=args.output_dir, max_tt_rank=args.tt_rank, device=args.device, )