""" DyQuant - 动态混合精度量化工具 支持层/头级别的不同精度混合（4/8/16 bit），在保持精度的同时提升吞吐 20%-30%。 使用方法： from inference.dyquant import DyQuantConverter, QuantConfig # 1. 创建量化配置 config = QuantConfig( model_path="fusion-8b-base", bits=4, # 默认 4-bit mixed_precision=True, # 混合精度 calib_samples=512, # 校准样本数 ) # 2. 转换模型 converter = DyQuantConverter(config) quantized_model = converter.convert() # 3. 保存量化模型 converter.save("fusion-8b-dyquant") # 4. 推理 output = quantized_model.generate(...) 作者：zhan1206 项目：Fusion - 六边形开源大模型 许可证：Apache 2.0 """ import torch import torch.nn as nn from typing import Dict, List, Optional, Tuple from dataclasses import dataclass import json from pathlib import Path import math @dataclass class QuantConfig: """ 量化配置 属性： model_path: 模型路径 bits: 默认量化位数（4/8） mixed_precision: 是否启用混合精度 calib_samples: 校准样本数 calib_data: 校准数据路径 output_path: 输出路径 per_head: 是否按头量化（True=更精细） """ model_path: str bits: int = 4 mixed_precision: bool = True calib_samples: int = 512 calib_data: Optional[str] = None output_path: Optional[str] = None per_head: bool = False class DyQuantConverter: """ 动态混合精度量化转换器 支持： - 动态量化（Dynamic Quantization）：int8，对延迟敏感场景 - 静态量化（Static Quantization）：int8/uint8，需校准数据 - 量化感知训练（QAT）：finetune-aware，适用于高精度需求 - 混合精度：不同层使用不同位数 - 按头量化：注意力头级别精度分配 """ def __init__(self, config: QuantConfig): """ 初始化转换器 参数： config: 量化配置 """ self.config = config self.model = None self.quant_layers = {} self.quantization_type = "dynamic" # 默认使用动态量化 print(f"[DyQuant] 初始化量化工具") print(f" 模型：{config.model_path}") print(f" 默认位数：{config.bits}-bit") print(f" 混合精度：{config.mixed_precision}") print(f" 按头量化：{config.per_head}") def load_model(self): """加载模型（优先 FusionModel/FusionMini，回退 AutoModelForCausalLM）""" print(f"\n[DyQuant] 加载模型：{self.config.model_path}") self.model = None # Try FusionMini first try: from models.fusion_mini import FusionMini self.model = FusionMini._load_from_safetensors(self.config.model_path) self.model.eval() print(f"[DyQuant] FusionMini 加载成功") return self.model except Exception as e1: pass # fallback # Try FusionModel try: from models.fusion_model import FusionModel self.model = FusionModel.from_pretrained(self.config.model_path) self.model.eval() print(f"[DyQuant] FusionModel 加载成功") return self.model except Exception as e2: pass # fallback # Fallback to AutoModelForCausalLM try: from transformers import AutoModelForCausalLM self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, torch_dtype=torch.bfloat16, device_map="cpu", trust_remote_code=True, ) self.model.eval() print(f"[DyQuant] AutoModelForCausalLM 加载成功（回退模式）") print(f"[DyQuant] 警告：非 Fusion 模型，SBLA/ThinkingDial 量化路径未验证") return self.model except Exception as e3: print(f"[DyQuant] 模型加载失败：{e3}") import traceback; traceback.print_exc() self.model = None return None def analyze_sensitivity(self) -> Dict[str, float]: """ 分析层敏感度 通过权重分布和激活值分析，确定哪些层对量化更敏感 返回： 层名称 -> 敏感度分数（0-1，越高越敏感） """ print(f"\n[DyQuant] 分析层敏感度...") sensitivity = {} if self.model is None: # 模拟模式 print(f"[DyQuant] 模拟模式：假设 32 层") for i in range(32): layer_name = f"model.layers.{i}" if i < 4 or i >= 28: sensitivity[layer_name] = 0.8 elif i < 8 or i >= 24: sensitivity[layer_name] = 0.5 else: sensitivity[layer_name] = 0.2 else: # 真实分析 print(f"[DyQuant] 计算真实敏感度...") # 遍历模型的所有 Linear 层 for name, module in self.model.named_modules(): if isinstance(module, nn.Linear): # 计算敏感度指标 sensitivity_score = self._compute_layer_sensitivity( module, name ) sensitivity[name] = sensitivity_score print(f"[DyQuant] 分析完成，共 {len(sensitivity)} 个量化层") # 统计 high_sens = sum(1 for v in sensitivity.values() if v > 0.6) mid_sens = sum(1 for v in sensitivity.values() if 0.3 < v <= 0.6) low_sens = sum(1 for v in sensitivity.values() if v <= 0.3) print(f" 高敏感层：{high_sens}") print(f" 中敏感层：{mid_sens}") print(f" 低敏感层：{low_sens}") return sensitivity def _compute_layer_sensitivity( self, layer: nn.Module, name: str, ) -> float: """ 计算单个层的敏感度 综合以下因素： 1. 权重分布的分散程度（std/mean） 2. 输出激活值的范围 3. 层的位置（首尾层更敏感） 参数： layer: 层模块 name: 层名称 返回： 敏感度分数（0-1） """ # 获取权重 weight = layer.weight.data # 1. 权重分散度（值域越大越敏感） w_std = weight.float().std().item() w_abs_max = weight.float().abs().max().item() w_dispersion = min(w_std / (w_abs_max + 1e-8), 1.0) # 2. 权重稀疏度（越稀疏越敏感） w_zero_ratio = (weight.float() == 0).sum().item() / weight.numel() # 3. 层位置因子（基于层名称推断） position_factor = 0.5 if "embeddings" in name or "output" in name: position_factor = 0.8 # 首尾层更敏感 elif "layers.0" in name or "layers.1" in name: position_factor = 0.7 # 前几层 # 综合敏感度 sensitivity = ( 0.4 * w_dispersion + 0.2 * (1 - w_zero_ratio) + 0.4 * position_factor ) return min(sensitivity, 1.0) def assign_precision( self, sensitivity: Dict[str, float], ) -> Dict[str, int]: """ 根据敏感度分配量化位数 策略： - 高敏感（>0.6）：8-bit 或 16-bit - 中敏感（0.3-0.6）：8-bit - 低敏感（<0.3）：4-bit 参数： sensitivity: 敏感度字典 返回： 层名称 -> 量化位数 """ print(f"\n[DyQuant] 分配量化精度...") precision_map = {} for layer_name, sens in sensitivity.items(): if sens > 0.6: precision_map[layer_name] = 8 # 高敏感用 8-bit elif sens > 0.3: precision_map[layer_name] = 8 # 中敏感用 8-bit else: precision_map[layer_name] = self.config.bits # 低敏感用配置位数 # 统计 num_8bit = sum(1 for b in precision_map.values() if b == 8) num_4bit = sum(1 for b in precision_map.values() if b == 4) num_other = len(precision_map) - num_8bit - num_4bit print(f" 8-bit 层：{num_8bit}") print(f" 4-bit 层：{num_4bit}") print(f" 其他精度：{num_other}") return precision_map def quantize_layer( self, layer: nn.Module, bits: int, per_head: bool = False, ) -> nn.Module: """ 量化单个层 使用 PyTorch 动态量化 API： - int8 对称量化 - 支持 Linear 层的动态量化 参数： layer: 待量化层 bits: 量化位数（4 或 8） per_head: 是否按头量化 返回： 量化后的层 """ if bits == 8: # PyTorch 动态量化（int8） quantized = torch.quantization.quantize_dynamic( layer, {nn.Linear}, dtype=torch.qint8, ) return quantized elif bits == 4: # 4-bit 对称量化（per-channel scale + zero_point） # 注意：此为训练感知近似量化，非推理加速专用 INT4 kernel return self._quantize_to_nbit(layer, 4) else: # 不量化 return layer def _quantize_to_nbit( self, layer: nn.Module, bits: int, ) -> nn.Module: """ 量化到指定位数 使用自定义的量化实现，支持任意位数 参数： layer: 待量化层 bits: 目标位数 返回： 量化后的层 """ class QuantizedLinear(nn.Module): def __init__(self, original_layer, bits): super().__init__() self.in_features = original_layer.in_features self.out_features = original_layer.out_features # 获取原始权重 weight = original_layer.weight.data.float() bias = original_layer.bias.data.float() if original_layer.bias is not None else None # 量化权重 q_weight, scale, zero_point = self._quantize_weight(weight, bits) self.register_buffer('q_weight', q_weight) self.register_buffer('scale', scale) self.register_buffer('zero_point', zero_point) if bias is not None: self.bias = nn.Parameter(bias) else: self.bias = None def _quantize_weight(self, weight, bits): """对称量化（per-channel）""" # 计算缩放因子（per-channel） max_vals = weight.abs().max(dim=1, keepdim=True).values # (out_features, 1) # 避免除零 max_vals = torch.where(max_vals == 0, torch.ones_like(max_vals), max_vals) qmax = 2 ** (bits - 1) - 1 scales = max_vals / qmax # (out_features, 1) # 量化 q_weight = torch.round(weight / scales).clamp(-qmax, qmax) return q_weight.to(torch.int8), scales.to(weight.dtype), torch.tensor(0, dtype=weight.dtype, device=weight.device) def forward(self, x): # 反量化 + 矩阵乘法 weight = self.q_weight.float() * self.scale return nn.functional.linear(x, weight, self.bias) return QuantizedLinear(layer, bits) def _quantize_layer_per_head( self, layer: nn.Module, num_heads: int, bits: int, ) -> nn.Module: """ 按头量化（用于注意力层） 将 weight matrix 按 head 分割，每个 head 独立量化 参数： layer: 待量化层 num_heads: 注意力头数 bits: 量化位数 返回： 量化后的层 """ # 按头量化实现 head_dim = layer.out_features // num_heads class PerHeadQuantizedLinear(nn.Module): def __init__(self, original_layer, num_heads, bits): super().__init__() self.num_heads = num_heads self.head_dim = head_dim weight = original_layer.weight.data.float() bias = original_layer.bias.data.float() if original_layer.bias is not None else None # 按头量化 q_weights = [] scales = [] for i in range(num_heads): head_weight = weight[:, i*head_dim:(i+1)*head_dim] q_w, scale, _ = self._quantize_head(head_weight, bits) q_weights.append(q_w) scales.append(scale) # 拼接 self.register_buffer('q_weights', torch.cat(q_weights, dim=1)) self.register_buffer('scales', torch.stack(scales)) if bias is not None: self.bias = nn.Parameter(bias) else: self.bias = None def _quantize_head(self, weight, bits): max_val = weight.abs().max() qmax = 2 ** (bits - 1) - 1 scale = max_val / qmax if max_val > 0 else torch.tensor(1.0) q_weight = torch.round(weight / scale).clamp(-2**(bits-1), 2**(bits-1)-1) return q_weight.to(torch.int8), scale.to(weight.dtype) def forward(self, x): # 解码每个头 weight_list = [] for i in range(self.num_heads): w = self.q_weights[:, i*self.head_dim:(i+1)*self.head_dim].float() w = w * self.scales[i] weight_list.append(w) weight = torch.cat(weight_list, dim=1) return nn.functional.linear(x, weight, self.bias) return PerHeadQuantizedLinear(layer, num_heads, bits) def convert(self) -> nn.Module: """ 执行量化转换 返回： 量化后的模型 """ print(f"\n[DyQuant] 开始量化转换...") # 1. 加载模型 if self.model is None: result = self.load_model() if result is not None: self.model = result if self.model is None: print(f"[DyQuant] 无法加载模型，返回 None") return None # 2. 分析敏感度 sensitivity = self.analyze_sensitivity() # 3. 分配精度 if self.config.mixed_precision: precision_map = self.assign_precision(sensitivity) else: precision_map = { layer: self.config.bits for layer in sensitivity.keys() } # 4. 逐层量化 print(f"\n[DyQuant] 逐层量化...") quantized_layers = {} for name, module in self.model.named_modules(): if name in precision_map: bits = precision_map[name] if isinstance(module, nn.Linear): if self.config.per_head and "q_proj" in name: # 按头量化（适用于 attention 投影） quantized = self._quantize_layer_per_head( module, num_heads=32, bits=bits ) else: quantized = self.quantize_layer(module, bits) quantized_layers[name] = quantized print(f" 量化 {name} -> {bits}-bit") # 5. 替换原模型层 for name, quantized_layer in quantized_layers.items(): self._replace_layer(self.model, name, quantized_layer) self.quant_layers = precision_map print(f"\n[DyQuant] 量化完成") print(f" 量化层数：{len(self.quant_layers)}") return self.model def _replace_layer( self, model: nn.Module, layer_name: str, new_layer: nn.Module, ): """替换模型中的指定层""" parts = layer_name.split('.') # 找到父模块 parent = model for part in parts[:-1]: if part.isdigit(): parent = parent[int(part)] else: parent = getattr(parent, part) # 替换 setattr(parent, parts[-1], new_layer) def save(self, output_path: str): """ 保存量化模型 参数： output_path: 输出路径 """ print(f"\n[DyQuant] 保存量化模型到：{output_path}") if self.model is None: print(f"[DyQuant] 模型为空，无法保存") return output_dir = Path(output_path) output_dir.mkdir(parents=True, exist_ok=True) # 保存模型 - extract state_dict from custom quantized layers # Custom QuantizedLinear layers are not HF-compatible, use safetensors directly try: import safetensors.torch as st state = self.model.state_dict() # Convert non-tensor values (scales/zeros) to tensors for serialization clean_state = {} for k, v in state.items(): if isinstance(v, torch.Tensor): clean_state[k] = v.contiguous() else: clean_state[k] = torch.tensor(v) if v is not None else torch.tensor(0.0) st.save_file(clean_state, str(output_dir / "model.safetensors")) except ImportError: # Fallback to torch.save if safetensors unavailable torch.save(self.model.state_dict(), output_dir / "pytorch_model.bin") # 保存量化配置 quant_config = { "quantization_type": self.quantization_type, "layers": self.quant_layers, "config": { "bits": self.config.bits, "mixed_precision": self.config.mixed_precision, "per_head": self.config.per_head, } } with open(output_dir / "quant_config.json", 'w', encoding='utf-8') as f: json.dump(quant_config, f, indent=2) print(f"[DyQuant] 保存完成") def get_model_size(self) -> int: """ 计算量化后的模型大小 返回： 模型大小（MB） """ if self.model is None: return 0 total_bytes = 0 for name, module in self.model.named_modules(): if isinstance(module, (nn.Linear, torch.nn.quantized.dynamic.Linear)): # 估计量化后的大小 if hasattr(module, 'qweight'): # int8: 1 byte per param total_bytes += module.qweight.numel() else: # float16: 2 bytes per param total_bytes += module.weight.numel() * 2 return total_bytes / (1024 * 1024) # ============================================================ # 便捷函数 # ============================================================ def quantize_fusion_model( model_path: str, output_path: str, bits: int = 4, mixed_precision: bool = True, ) -> str: """ 一键量化 Fusion 模型 参数： model_path: 模型路径 output_path: 输出路径 bits: 量化位数 mixed_precision: 是否混合精度 返回： 输出路径 """ config = QuantConfig( model_path=model_path, bits=bits, mixed_precision=mixed_precision, output_path=output_path, ) converter = DyQuantConverter(config) converter.convert() converter.save(output_path) size_mb = converter.get_model_size() print(f"\n量化完成！模型大小：{size_mb:.1f} MB") return output_path # ============================================================ # QAT (Quantization-Aware Training) Integration # ============================================================ class QATTrainer: """ Quantization-Aware Training trainer for Fusion models. Inserts fake-quantization nodes into the model during training, so the model learns to be robust to quantization noise. After training, the model can be quantized with minimal accuracy loss. Usage: from inference.dyquant import QATTrainer, QuantConfig config = QuantConfig(model_path="fusion-8b-base", bits=4) trainer = QATTrainer(config, train_data="data/train.json") trainer.train(epochs=3, lr=1e-5) trainer.save("fusion-8b-qat") """ def __init__( self, config: QuantConfig, train_data: Optional[str] = None, learning_rate: float = 1e-5, warmup_steps: int = 100, ): self.config = config self.train_data = train_data self.lr = learning_rate self.warmup_steps = warmup_steps self.converter = DyQuantConverter(config) self.model = None self.qat_model = None def prepare(self) -> Optional[nn.Module]: """Load model and insert fake-quantization nodes.""" self.model = self.converter.load_model() if self.model is None: print("[DyQuant] QAT: model load failed, cannot prepare") return None self.qat_model = self._insert_fake_quant(self.model) return self.qat_model def _insert_fake_quant(self, model: nn.Module) -> nn.Module: """Insert fake-quantization observers into all Linear layers.""" for name, module in model.named_modules(): if isinstance(module, nn.Linear): # Use PyTorch native fake quantization for all Linear layers module.qconfig = torch.ao.quantization.get_default_qat_qconfig('x86') torch.ao.quantization.prepare_qat(model, inplace=True) return model def train( self, epochs: int = 3, lr: Optional[float] = None, batch_size: int = 4, max_seq_len: int = 2048, ): """ Run QAT fine-tuning. Args: epochs: Number of training epochs lr: Learning rate (defaults to self.lr) batch_size: Training batch size max_seq_len: Maximum sequence length """ if self.qat_model is None: self.prepare() actual_lr = lr or self.lr device = next(self.qat_model.parameters()).device optimizer = torch.optim.AdamW(self.qat_model.parameters(), lr=actual_lr) scheduler = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.1, total_iters=self.warmup_steps ) print(f"[QAT] Starting QAT training: epochs={epochs}, lr={actual_lr}") # Load training data if provided if self.train_data and Path(self.train_data).exists(): train_dataset = self._load_dataset(self.train_data, max_seq_len) else: print("[QAT] Warning: No training data provided, using random calibration") train_dataset = self._generate_calib_data(batch_size * 10, max_seq_len) dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, shuffle=True ) self.qat_model.train() step = 0 for epoch in range(epochs): total_loss = 0.0 for batch in dataloader: input_ids = batch.to(device) attention_mask = torch.ones_like(input_ids) labels = input_ids.clone() outputs = self.qat_model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss if hasattr(outputs, 'loss') else outputs['loss'] loss.backward() optimizer.step() scheduler.step() optimizer.zero_grad() total_loss += loss.item() step += 1 avg_loss = total_loss / len(dataloader) print(f"[QAT] Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f}") print(f"[QAT] Training complete ({step} steps)") def _load_dataset(self, data_path: str, max_seq_len: int): """Load JSON training data and tokenize with model's tokenizer.""" import json from transformers import AutoTokenizer with open(data_path, 'r', encoding='utf-8') as f: data = json.load(f) texts = [item.get('text', item.get('prompt', '')) + ' ' + item.get('response', '') for item in data] # Load tokenizer from model path try: tokenizer = AutoTokenizer.from_pretrained( self.config.model_path, trust_remote_code=True, ) except Exception as e: print(f"[QAT] Warning: Cannot load tokenizer: {e}, using character-level fallback") # Fallback to character-level encoding encoded = [list(t.encode('utf-8'))[:max_seq_len] for t in texts] padded = [ seq + [0] * (max_seq_len - len(seq)) if len(seq) < max_seq_len else seq[:max_seq_len] for seq in encoded ] return torch.utils.data.TensorDataset(torch.tensor(padded, dtype=torch.long)) # Tokenize with proper tokenizer encoded = [] for text in texts: # Encode text to token IDs token_ids = tokenizer.encode(text, max_length=max_seq_len, truncation=True) # Pad or truncate to max_seq_len if len(token_ids) < max_seq_len: token_ids = token_ids + [tokenizer.pad_token_id or 0] * (max_seq_len - len(token_ids)) else: token_ids = token_ids[:max_seq_len] encoded.append(token_ids) return torch.utils.data.TensorDataset(torch.tensor(encoded, dtype=torch.long)) def _generate_calib_data(self, num_samples: int, seq_len: int): """Generate random calibration data.""" data = torch.randint(0, 1000, (num_samples, seq_len)) return torch.utils.data.TensorDataset(data) def save(self, output_path: str): """Convert QAT model to final quantized model and save.""" # Remove fake-quant nodes and convert to actual quantized model final_model = torch.ao.quantization.convert(self.qat_model, inplace=False) output_dir = Path(output_path) output_dir.mkdir(parents=True, exist_ok=True) torch.save(final_model.state_dict(), output_dir / "qat_model.pt") # Also save as regular quantized model self.config.output_path = output_path quantized = self.converter.convert() self.converter.save(output_path) print(f"[QAT] Saved to {output_path}") # ============================================================ # Main Entry Point # ============================================================ if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="DyQuant 模型量化工具") parser.add_argument("--model_path", type=str, required=True, help="模型路径") parser.add_argument("--output_path", type=str, required=True, help="输出路径") parser.add_argument("--bits", type=int, default=4, help="量化位数（4/8）") parser.add_argument("--mixed_precision", action="store_true", help="启用混合精度") parser.add_argument("--per_head", action="store_true", help="按头量化") args = parser.parse_args() # 创建量化配置 config = QuantConfig( model_path=args.model_path, bits=args.bits, mixed_precision=args.mixed_precision, per_head=args.per_head, output_path=args.output_path, ) # 量化 converter = DyQuantConverter(config) converter.convert() converter.save(args.output_path) print(f"\n量化完成！") print(f"输出路径：{args.output_path}") print(f"模型大小：{converter.get_model_size():.1f} MB")