import torch import gc import functools import json import inspect from tqdm import tqdm from collections import defaultdict from torch import optim from .smooth_quantizer import ACIQ, SmoothQuantizer import torch.nn.functional as F class OmniQuantizer: def __init__( self, model, max_calib_samples=32, max_calib_seq_len=128, act_bit=8, act_sym=True, generate_for_npu=False, epochs=20, lr=5e-3, wd=0.0 ) -> None: self.model = model self.tokenizer = model.tokenizer self.act_bit = act_bit self.act_sym = act_sym self.generate_for_npu = generate_for_npu self.epochs = epochs self.lr = lr self.wd = wd self.max_calib_samples = max_calib_samples self.max_calib_seq_len = max_calib_seq_len self.calib_data = 'wikitext' if model.args.calib_data is None else model.args.calib_data self.split = 'train' self.best_device = self.get_best_device() self.modules = self.model.blocks self.act_quanter = ACIQ(act_bit) self.moment = 0.99 if "cpu" != self.best_device: for idx in range(len(self.modules)): self.to_device(self.modules[idx], "cpu") self.act_dict = [defaultdict(dict) for _ in range(len(self.modules))] @staticmethod def _is_offset_rmsnorm(op): type_name = str(type(op)) if any( t in type_name for t in [ "GemmaRMSNorm", "Qwen3_5RMSNorm", "Qwen3_5MoeRMSNorm", "Qwen3NextRMSNorm", ] ): return True return False @staticmethod def get_best_device(): if torch.backends.mps.is_available(): return "mps" elif torch.cuda.is_available(): return "cuda:0" else: return "cpu" @staticmethod def to_device(module, device): for child_name, child_module in module.named_children(): if child_name == 'self_attn': for sub_name, sub_child in child_module.named_children(): if sub_name != 'config': sub_child.to(device) else: child_module.to(device) @staticmethod def clear_memory(weight=None): if weight is not None: del weight gc.collect() torch.cuda.empty_cache() @staticmethod def get_calib_dataset(data, tokenizer=None, n_samples=128, max_seq_len=512, split="train"): custom_calib_data = False if isinstance(data, str): from datasets import load_dataset if data == "pileval": dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation") elif data == "wikitext": dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split=split) else: custom_calib_data = True with open(data, 'r', encoding='utf-8') as f: dataset = f.read().splitlines() else: raise NotImplementedError("Data loading error") samples = [] if custom_calib_data == False: dataset = dataset.shuffle(seed=42) count = 0 idx = 0 while count < n_samples and idx < len(dataset): try: text = dataset[idx]["text"] # skip empty lines if not text.strip(): idx += 1 continue input_ids = tokenizer( text, return_tensors="pt", max_length=max_seq_len, truncation=True ).input_ids # skip empty tokenized inputs if input_ids.numel() > 0: samples.append(input_ids) count += 1 except: pass idx += 1 else: for i in range(min(n_samples, len(dataset))): messages = [{"role": "system", "content": ""}, {"role": "user", "content": dataset[i]}] prompt = tokenizer.apply_chat_template(messages) input_ids = tokenizer( prompt, return_tensors="pt", max_length=max_seq_len, truncation=True ).input_ids if input_ids.numel() > 0: samples.append(input_ids) print(f"Collected {len(samples)} valid calibration samples.") return samples def init_quant(self, n_samples=128, max_seq_len=512): samples = self.get_calib_dataset( data=self.calib_data, tokenizer=self.tokenizer, n_samples=n_samples, max_seq_len=max_seq_len, split=self.split ) return samples def _get_first_input(self, sample): sample = sample.long() layer_kwargs = {} seq_len = sample.numel() new_tokens = 0 try: inps = self.model.embedding(sample) except RuntimeError: sample = sample.to(self.best_device) inps = self.model.embedding(sample) inps = inps.to(self.best_device) position_ids = self.model.get_position_ids(seq_len, new_tokens, sample) rotary_pos_emb = self.model.rotary(position_ids) attention_mask = self.model.get_attention_mask(seq_len, new_tokens) layer_kwargs["rotary_pos_emb"] = rotary_pos_emb.to(self.best_device) layer_kwargs["attention_mask"] = attention_mask.to(self.best_device) layer_kwargs["position_ids"] = position_ids.to(self.best_device) del sample self.clear_memory() return layer_kwargs, inps def _sanitize_kwargs(self, inputs_kwargs, module): module_signature = inspect.signature(module.forward).parameters sanitized_kwargs = {} for k, v in inputs_kwargs.items(): if k in module_signature: sanitized_kwargs[k] = v return sanitized_kwargs def _select_layer_kwargs(self, module, inputs_kwargs): """Select per-layer kwargs for mixed attention models.""" selected_kwargs = dict(inputs_kwargs) attention_mask = selected_kwargs.get("attention_mask", None) if attention_mask is None: return selected_kwargs if getattr(self.model.config, "attention_type", None) != "mix": return selected_kwargs if isinstance(attention_mask, torch.Tensor) and attention_mask.dim() >= 1 and attention_mask.shape[0] == 2: layer_type = getattr(module, "layer_type", None) is_sliding = layer_type in ("linear_attention", "sliding_attention") selected_kwargs["attention_mask"] = attention_mask[int(is_sliding)] return selected_kwargs def _clear_block_kv_cache(self, block): """Clear KV cache on the block's attention so each calibration sample is independent.""" if hasattr(block, "self_attn") and block.self_attn is not None: if hasattr(block.self_attn, "past_key_value"): block.self_attn.past_key_value = None if hasattr(block.self_attn, "conv_state"): block.self_attn.conv_state = None if hasattr(block.self_attn, "rnn_state"): block.self_attn.rnn_state = None def _safe_forward(self, x, module, module_kwargs): try: target_dtype = next(module.parameters()).dtype target_device = next(module.parameters()).device except StopIteration: target_dtype = torch.float32 target_device = x.device x = x.to(device=target_device, dtype=target_dtype) if "cuda" in str(target_device): with torch.cuda.amp.autocast(enabled=True, dtype=target_dtype): out = module(x, **module_kwargs) else: out = module(x, **module_kwargs) if isinstance(out, tuple): out = out[0] return out def _run_optimization(self, x_in, fcs, ln, act_max): device = self.best_device target_dtype = list(fcs[0].parameters())[0].dtype # Increase micro_batch_size for better GPU utilization micro_batch_size = 64 # Pre-move weights to GPU and keep there weights = torch.cat([fc.weight for fc in fcs], dim=0).to(device) act_max = act_max.to(device=device, dtype=target_dtype) weight_max_per_channel = torch.cat([fc.weight.abs().max(dim=0, keepdim=True)[0] for fc in fcs], dim=0) weight_max_per_channel = weight_max_per_channel.max(dim=0)[0].clamp(min=1e-5).to(device) scales_init = (act_max.pow(0.5) / weight_max_per_channel.pow(0.5)).clamp(min=1e-5) scales_init = scales_init.to(device=device, dtype=target_dtype) log_scale = torch.nn.Parameter(torch.log(scales_init)) with torch.no_grad(): w_init_smooth = weights * scales_init.view(1, -1) clip_init = w_init_smooth.abs().max(dim=1, keepdim=True)[0] clip_val = torch.nn.Parameter(clip_init) optimizer = optim.AdamW([ {'params': [log_scale], 'lr': self.lr}, {'params': [clip_val], 'lr': self.lr} ], weight_decay=self.wd) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.epochs, eta_min=self.lr * 0.1 ) # Pre-compute constants for quantization act_bit = self.act_bit act_sym = self.act_sym q_max_w = 2 ** (act_bit - 1) - 1 # Inline quantize functions to reduce function call overhead if act_sym: q_max_act = 2 ** (act_bit - 1) - 1 q_min_act = -q_max_act else: q_max_act = 2 ** act_bit - 1 q_min_act = 0 N = x_in.shape[0] num_steps = (N + micro_batch_size - 1) // micro_batch_size # Pre-load all data to GPU if it fits, otherwise use pinned memory try: # Try to fit all data on GPU x_in_gpu = x_in.to(device, dtype=target_dtype) use_gpu_data = True except RuntimeError: # Fall back to CPU with pinned memory for faster transfer x_in_gpu = x_in.pin_memory() if x_in.device.type == 'cpu' else x_in use_gpu_data = False # Pre-compute target outputs for all batches (computed once, not every epoch) with torch.no_grad(): y_targets = [] for i in range(0, N, micro_batch_size): if use_gpu_data: x_batch = x_in_gpu[i : i + micro_batch_size] else: x_batch = x_in_gpu[i : i + micro_batch_size].to(device, dtype=target_dtype, non_blocking=True) with torch.cuda.amp.autocast(enabled=True, dtype=target_dtype): y_target = F.linear(x_batch, weights) y_targets.append(y_target.float()) if not use_gpu_data: del x_batch for epoch in range(self.epochs): optimizer.zero_grad(set_to_none=True) # More efficient than zero_grad() total_loss = 0.0 for batch_idx, i in enumerate(range(0, N, micro_batch_size)): if use_gpu_data: x_micro = x_in_gpu[i : i + micro_batch_size] else: x_micro = x_in_gpu[i : i + micro_batch_size].to(device, dtype=target_dtype, non_blocking=True) y_micro_target = y_targets[batch_idx] scale = torch.exp(log_scale) s_view = scale.view(1, -1) x_sim = x_micro / s_view w_sim = weights * s_view if act_sym: act_scale = x_sim.abs().max() / q_max_act act_scale = torch.clamp(act_scale, min=1e-5) x_q = torch.round(x_sim / act_scale) x_q = torch.clamp(x_q, q_min_act, q_max_act) x_q = x_q * act_scale x_q = (x_q - x_sim).detach() + x_sim else: t_min, t_max = x_sim.min(), x_sim.max() act_scale = (t_max - t_min) / q_max_act act_scale = torch.clamp(act_scale, min=1e-5) zero = -t_min / act_scale x_q = torch.round(x_sim / act_scale + zero) x_q = torch.clamp(x_q, q_min_act, q_max_act) x_q = (x_q - zero) * act_scale x_q = (x_q - x_sim).detach() + x_sim # Inline quantize_weight_with_clip clip_v = F.relu(clip_val) + 1e-5 w_clamped = torch.clamp(w_sim, -clip_v, clip_v) w_scale = clip_v / q_max_w w_q = torch.round(w_clamped / w_scale) * w_scale w_q = (w_q - w_clamped).detach() + w_clamped x_q = x_q.to(dtype=target_dtype) w_q = w_q.to(dtype=target_dtype) with torch.cuda.amp.autocast(enabled=True, dtype=target_dtype): y_pred = F.linear(x_q, w_q) loss = F.mse_loss(y_pred.float(), y_micro_target) loss = loss / num_steps loss.backward() total_loss += loss.item() optimizer.step() scheduler.step() # Cleanup y_targets del y_targets with torch.no_grad(): final_scale = torch.exp(log_scale).detach().view(-1) if self._is_offset_rmsnorm(ln): ln.weight += 1 ln.weight.div_(final_scale) ln.weight -= 1 else: ln.weight.div_(final_scale) if hasattr(ln, "bias") and ln.bias is not None: ln.bias.div_(final_scale) final_clip = F.relu(clip_val).detach() current_idx = 0 for fc in fcs: num_out = fc.weight.shape[0] layer_clip = final_clip[current_idx : current_idx + num_out] fc.weight.mul_(final_scale.view(1, -1)) fc.weight.data = torch.clamp(fc.weight.data, -layer_clip, layer_clip) current_idx += num_out del x_in_gpu, weights, weight_max_per_channel, scales_init, log_scale, clip_val, optimizer, scheduler torch.cuda.empty_cache() def _get_robust_act_max(self, x): try: x_flat = x.reshape(-1, x.shape[-1]) if x_flat.shape[0] > 2048: if x_flat.shape[0] > 10000: indices = torch.randperm(x_flat.shape[0])[:10000] x_sample = x_flat[indices] else: x_sample = x_flat robust_max = torch.quantile(x_sample.abs().float(), 0.999, dim=0) return robust_max else: return x_flat.abs().max(dim=0)[0] except: return x.reshape(-1, x.shape[-1]).abs().max(dim=0)[0] def _extract_static_scales(self): print("OmniQuant: Extracting final JSON scales...") def compute_scale_sym(max_min): bit_scale = 2 ** (self.act_bit - 1) - 1 max_v = max_min.abs().max().item() scale = max_v / bit_scale return [scale, 0.0] def compute_scale_zero_asym(max_min): bit_scale = 2 ** (self.act_bit) - 1 max_v = max_min[0].item() min_v = max_min[1].item() if max_v < 0.0: max_v = 0.0 if min_v > 0.0: min_v = 0.0 scale = 1.0 if max_v == min_v else (max_v - min_v) / bit_scale zero = round(-min_v / scale - 2 ** (self.act_bit - 1)) if self.act_bit == 16 and self.generate_for_npu: zero = round(min_v / scale) return [scale, zero] func = compute_scale_sym if self.act_sym else compute_scale_zero_asym for idx in range(len(self.act_dict)): for name, input_output in self.act_dict[idx].items(): self.act_dict[idx][name]['input'] = func(input_output['input']) self.act_dict[idx][name]['output'] = func(input_output['output']) def _get_all_static_scales_safe(self, idx, layer, named_linears, x_in, module_kwargs): def stat_io_hook(m, x, y, name): if isinstance(x, tuple): x = x[0] if isinstance(y, tuple): y = y[0] inp_max_min = self.act_quanter.get_max_min(x.detach().float().to("cpu")) out_max_min = self.act_quanter.get_max_min(y.detach().float().to("cpu")) if name not in self.act_dict[idx] or "input" not in self.act_dict[idx][name]: self.act_dict[idx][name]["input"] = inp_max_min else: self.act_dict[idx][name]["input"] = inp_max_min * (1-self.moment) + self.moment * self.act_dict[idx][name]["input"] if name not in self.act_dict[idx] or "output" not in self.act_dict[idx][name]: self.act_dict[idx][name]["output"] = out_max_min else: self.act_dict[idx][name]["output"] = out_max_min * (1-self.moment) + self.moment * self.act_dict[idx][name]["output"] handles = [] for name in named_linears: handles.append(named_linears[name].register_forward_hook(functools.partial(stat_io_hook, name=name))) layer_kwargs = self._select_layer_kwargs(layer, module_kwargs) sanitized_kwargs = self._sanitize_kwargs(layer_kwargs, layer) with torch.no_grad(): self._safe_forward(x_in, layer, sanitized_kwargs) for h in handles: h.remove() def _prepare_calibration_data(self): """Prepare calibration samples and compute initial embeddings. Returns: layer_inputs: List of (input_tensor, kwargs) tuples for each sample """ # Check if we already have prepared data if hasattr(self, '_cached_layer_inputs') and self._cached_layer_inputs is not None: print("OmniQuant: Using cached calibration data...") return self._cached_layer_inputs print("OmniQuant: Initializing...") self.samples = self.init_quant( n_samples=self.max_calib_samples, max_seq_len=self.max_calib_seq_len, ) print(f"OmniQuant: Pre-computing embeddings for {len(self.samples)} samples...") layer_inputs = [] for sample in self.samples: # skip empty sample if sample.numel() == 0: continue kw, inp = self._get_first_input(sample) cpu_kw = {} for k, v in kw.items(): if isinstance(v, torch.Tensor): cpu_kw[k] = v.to("cpu") else: cpu_kw[k] = v layer_inputs.append((inp.to("cpu"), cpu_kw)) self.clear_memory() # Cache for potential reuse self._cached_layer_inputs = layer_inputs return layer_inputs def optimize_weights(self, collect_feature_map=False): """Phase 1: Optimize weights by adjusting LayerNorm and Linear layer weights. This phase applies smooth quantization optimization to reduce quantization error. Args: collect_feature_map: If True, also collect feature map info during this pass to avoid a second traversal. """ layer_inputs = self._prepare_calibration_data() if collect_feature_map: # Re-initialize act_dict for fresh collection self.act_dict = [defaultdict(dict) for _ in range(len(self.modules))] print(f"OmniQuant: Starting weight optimization (Epochs={self.epochs})...") for idx in tqdm(range(len(self.modules)), desc="OmniQuant: Optimize Weights"): block = self.modules[idx] self.to_device(block, self.best_device) attn_inputs_list = [] mlp_inputs_list = [] next_layer_outputs = [] def hook_attn_input(m, i, o): if isinstance(i, tuple) and len(i) > 0: inp = i[0] else: inp = i attn_inputs_list.append(inp.detach().view(-1, inp.shape[-1])) def hook_mlp_input(m, i, o): if isinstance(i, tuple) and len(i) > 0: inp = i[0] else: inp = i mlp_inputs_list.append(inp.detach().view(-1, inp.shape[-1])) attn_module = block.self_attn attn_linears = [] attn_hook_target = None if all(hasattr(attn_module, name) for name in ("q_proj", "k_proj", "v_proj")): attn_linears = [attn_module.q_proj, attn_module.k_proj, attn_module.v_proj] attn_hook_target = attn_module.q_proj elif hasattr(attn_module, "in_proj_qkv"): attn_linears = [attn_module.in_proj_qkv] for optional_name in ("in_proj_a", "in_proj_b", "in_proj_z"): optional_proj = getattr(attn_module, optional_name, None) if optional_proj is not None: attn_linears.append(optional_proj) attn_hook_target = attn_module.in_proj_qkv h1 = None if attn_hook_target is not None: h1 = attn_hook_target.register_forward_hook(hook_attn_input) h2 = block.mlp.gate_proj.register_forward_hook(hook_mlp_input) # Pre-compute sanitized kwargs once for this block sample_kw_gpu = {} for k, v in layer_inputs[0][1].items(): if isinstance(v, torch.Tensor): sample_kw_gpu[k] = v.to(self.best_device) else: sample_kw_gpu[k] = v sample_kw_gpu = self._select_layer_kwargs(block, sample_kw_gpu) sanitized_kw_template = self._sanitize_kwargs(sample_kw_gpu, block) # Single forward pass: collect hooks AND compute outputs with torch.no_grad(): for inp, kw in layer_inputs: # Clear KV cache so each sample is processed independently (no past_key_value from previous iteration) self._clear_block_kv_cache(block) inp_gpu = inp.to(self.best_device) kw = self._select_layer_kwargs(block, kw) # Reuse sanitized keys, only update tensor values kw_gpu = {} for k, v in kw.items(): if k in sanitized_kw_template: if isinstance(v, torch.Tensor): kw_gpu[k] = v.to(self.best_device) else: kw_gpu[k] = v out = self._safe_forward(inp_gpu, block, kw_gpu) # Store output for next layer next_layer_outputs.append((out.detach().to("cpu"), kw)) del inp_gpu, kw_gpu, out if h1 is not None: h1.remove() h2.remove() # Process collected attention inputs optimize_attn = len(attn_inputs_list) > 0 if optimize_attn: # Concatenate on GPU, then move to CPU once total_attn_in = torch.cat(attn_inputs_list, dim=0).to("cpu") del attn_inputs_list # Mixed-attention models such as Qwen3.5 are highly sensitive to # attention-side rescaling. Keep the generic optimization path for # other architectures only. ln_attn = block.input_layernorm robust_max_attn = self._get_robust_act_max(total_attn_in) self._run_optimization(total_attn_in, attn_linears, ln_attn, robust_max_attn) del robust_max_attn del total_attn_in else: del attn_inputs_list # Process collected MLP inputs optimize_mlp = len(mlp_inputs_list) > 0 if optimize_mlp: # Concatenate on GPU, then move to CPU once total_mlp_in = torch.cat(mlp_inputs_list, dim=0).to("cpu") del mlp_inputs_list fcs_mlp = [block.mlp.gate_proj, block.mlp.up_proj] ln_mlp = block.post_attention_layernorm robust_max_mlp = self._get_robust_act_max(total_mlp_in) self._run_optimization(total_mlp_in, fcs_mlp, ln_mlp, robust_max_mlp) del total_mlp_in, robust_max_mlp else: del mlp_inputs_list self.clear_memory() # Outputs already computed in the single forward pass above layer_inputs = next_layer_outputs del next_layer_outputs if "cpu" != self.best_device: self.to_device(block, "cpu") self.clear_memory() print("OmniQuant: Weight optimization completed.") # Save final layer outputs for potential reuse by collect_feature_map_info self._final_layer_outputs = layer_inputs for idx in range(len(self.modules)): self.to_device(self.modules[idx], "cpu") self.clear_memory() # If collect_feature_map is requested, do it now using optimized weights if collect_feature_map: self._collect_feature_map_optimized() def _collect_lm_head_info(self, calib_inputs): """Collect lm_head layer activation info for NPU.""" if not self.generate_for_npu: return lm_head_idx = len(self.modules) self.act_dict.append(defaultdict(dict)) if hasattr(self.model, 'lm') and hasattr(self.model.lm, 'lm'): lm_head = self.model.lm.lm elif hasattr(self.model, 'lm_head'): lm_head = self.model.lm_head else: lm_head = None print("Warning: lm_head not found in model, skipping lm_head calibration.") if lm_head is not None: lm_head.to(self.best_device) if hasattr(self.model, 'final_layernorm'): self.model.final_layernorm.to(self.best_device) lm_head_ops = {'lm_head': lm_head} for inp, kw in calib_inputs: inp_gpu = inp.to(self.best_device) with torch.no_grad(): if hasattr(self.model, 'final_layernorm'): hidden_states = self.model.final_layernorm(inp_gpu) else: hidden_states = inp_gpu self._get_all_static_scales_safe(lm_head_idx, lm_head, lm_head_ops, hidden_states, {}) del inp_gpu, hidden_states lm_head.to("cpu") if hasattr(self.model, 'final_layernorm'): self.model.final_layernorm.to("cpu") self.clear_memory() def _collect_feature_map_optimized(self): """Optimized feature map collection that reuses embedding computation. This uses cached layer outputs from optimize_weights() to avoid re-computing embeddings through all layers. """ print("OmniQuant: Collecting static activation scales (optimized)...") gc.collect() torch.cuda.empty_cache() # Re-initialize act_dict self.act_dict = [defaultdict(dict) for _ in range(len(self.modules))] # Use cached initial inputs calib_inputs = self._prepare_calibration_data() for idx in tqdm(range(len(self.modules)), desc="Collecting Feature Map Info"): block = self.modules[idx] self.to_device(block, self.best_device) target_ops = SmoothQuantizer.get_all_leaf_modules(block) next_calib_inputs = [] # Batch process samples for better GPU utilization batch_size = min(8, len(calib_inputs)) # Process multiple samples together for batch_start in range(0, len(calib_inputs), batch_size): batch_end = min(batch_start + batch_size, len(calib_inputs)) batch_items = calib_inputs[batch_start:batch_end] for inp, kw in batch_items: # Process each calibration sample independently to avoid KV cache from the previous sample causing dimension mismatch between attn_weights and attention_mask self._clear_block_kv_cache(block) inp_gpu = inp.to(self.best_device) kw = self._select_layer_kwargs(block, kw) kw_gpu = {k: (v.to(self.best_device) if isinstance(v, torch.Tensor) else v) for k, v in kw.items()} # First forward: collect activation scales self._get_all_static_scales_safe(idx, block, target_ops, inp_gpu, kw_gpu) sanitized_kw = self._sanitize_kwargs(kw_gpu, block) # Clear again before the second forward: _get_all_static_scales_safe has already run one forward and written past_key_value; if not cleared, key concatenation in this forward would cause dimension error (256 vs 128) self._clear_block_kv_cache(block) with torch.no_grad(): out = self._safe_forward(inp_gpu, block, sanitized_kw) next_calib_inputs.append((out.cpu(), kw)) del inp_gpu, kw_gpu, out calib_inputs = next_calib_inputs if "cpu" != self.best_device: self.to_device(block, "cpu") self.clear_memory() # Collect lm_head info if needed self._collect_lm_head_info(calib_inputs) del calib_inputs self.clear_memory() self._extract_static_scales() for idx in range(len(self.modules)): self.to_device(self.modules[idx], "cpu") self.clear_memory() print("OmniQuant: Feature map info collection completed.") def quantize(self, collect_feature_map=False): """Run the full OmniQuant quantization pipeline. Args: collect_feature_map: If True, collect feature map info after weight optimization. If False, only perform weight optimization. """ # Run weight optimization, optionally collecting feature map info in the same pass self.optimize_weights(collect_feature_map=collect_feature_map) def clear_cache(self): """Clear cached calibration data to free memory.""" if hasattr(self, '_cached_layer_inputs'): del self._cached_layer_inputs self._cached_layer_inputs = None if hasattr(self, '_final_layer_outputs'): del self._final_layer_outputs self._final_layer_outputs = None self.clear_memory() def _find_match_in_dict(self, mnn_op_name, layer_act_dict): best_match = None max_len = 0 for pt_name in layer_act_dict.keys(): pt_path = pt_name.replace('.', '/') if pt_path in mnn_op_name: if len(pt_path) > max_len: max_len = len(pt_path) best_match = pt_name return best_match def _propagate_quant_info(self, mnn_ops, quant_info_dict): import copy PASS_THROUGH_OPS = [ 'Reshape', 'Squeeze', 'Unsqueeze', 'Flatten', 'Transpose', 'Permute', 'ConvertTensor', 'Cast', 'Slice', 'StridedSlice', 'Split', 'Concat', 'Pack' ] DATA_SELECT_OPS = ['Gather', 'GatherV2', 'GatherND'] print("Start propagating quantization parameters...") changed = True pass_round = 0 while changed: changed = False pass_round += 1 update_count = 0 for op in mnn_ops: op_type = op.get('type', '') inputs = op.get('inputIndexes', []) outputs = op.get('outputIndexes', []) if not inputs or not outputs: continue if op_type in PASS_THROUGH_OPS: source_info = None for inp_idx in inputs: if inp_idx in quant_info_dict: source_info = quant_info_dict[inp_idx] break if source_info: for out_idx in outputs: if out_idx not in quant_info_dict: quant_info_dict[out_idx] = copy.deepcopy(source_info) quant_info_dict[out_idx]['index'] = out_idx # 修正 index changed = True update_count += 1 target_info = None for out_idx in outputs: if out_idx in quant_info_dict: target_info = quant_info_dict[out_idx] break if target_info: for inp_idx in inputs: if inp_idx not in quant_info_dict: quant_info_dict[inp_idx] = copy.deepcopy(target_info) quant_info_dict[inp_idx]['index'] = inp_idx changed = True update_count += 1 elif op_type in DATA_SELECT_OPS: data_idx = inputs[0] out_idx = outputs[0] # Forward: Data -> Output if data_idx in quant_info_dict and out_idx not in quant_info_dict: quant_info_dict[out_idx] = copy.deepcopy(quant_info_dict[data_idx]) quant_info_dict[out_idx]['index'] = out_idx changed = True update_count += 1 # Backward: Output -> Data if out_idx in quant_info_dict and data_idx not in quant_info_dict: quant_info_dict[data_idx] = copy.deepcopy(quant_info_dict[out_idx]) quant_info_dict[data_idx]['index'] = data_idx changed = True update_count += 1 elif op_type == 'BinaryOp': out_idx = outputs[0] if out_idx in quant_info_dict: target_info = quant_info_dict[out_idx] for inp_idx in inputs: if inp_idx not in quant_info_dict: quant_info_dict[inp_idx] = copy.deepcopy(target_info) quant_info_dict[inp_idx]['index'] = inp_idx changed = True update_count += 1 else: scales = [] valid_inputs = [] for inp_idx in inputs: if inp_idx in quant_info_dict: scales.append(quant_info_dict[inp_idx]['quantInfo']['scale']) valid_inputs.append(inp_idx) if len(valid_inputs) > 0: max_scale_idx = valid_inputs[scales.index(max(scales))] source = quant_info_dict[max_scale_idx] quant_info_dict[out_idx] = copy.deepcopy(source) quant_info_dict[out_idx]['index'] = out_idx changed = True update_count += 1 print(f" Pass {pass_round}: Updated {update_count} tensors.") return quant_info_dict def apply(self, base_path): mnn = json.load(open(base_path, 'rt')) mnn['extraTensorDescribe'] = [] max_val = 2 ** (self.act_bit - 1) - 1 min_val = -max_val data_type = 'DT_INT16' if self.act_bit <= 8: data_type = 'DT_INT8' elif self.act_bit > 8 and self.act_bit <= 16: data_type = 'DT_INT16' quant_info_dict = {} npu_ignore_types = {'Input', 'Const', 'Extra', 'Reshape', 'ConvertTensor'} for op in mnn['oplists']: op_name = op.get('name', '') op_type = op.get('type', '') should_process = False if not self.generate_for_npu: should_process = (op_type == 'Convolution') else: should_process = (op_type not in npu_ignore_types) # Handle lm_head separately using the dedicated index if 'lm_head' in op_name: if self.generate_for_npu and should_process: lm_head_idx = len(self.modules) # lm_head is stored at this index if lm_head_idx < len(self.act_dict) and len(self.act_dict[lm_head_idx]) > 0: # lm_head stats are stored with key 'lm_head' if 'lm_head' in self.act_dict[lm_head_idx]: stats = self.act_dict[lm_head_idx]['lm_head'] print("Quantize lm head for QNN") if 'input' in stats and len(op['inputIndexes']) > 0: tensor_idx = op['inputIndexes'][0] if tensor_idx not in quant_info_dict: scale, zero = stats['input'] quant_info_dict[tensor_idx] = { 'index': tensor_idx, 'quantInfo': { 'scale': scale, 'zero': zero, 'min': min_val, 'max': max_val, "type": data_type } } if 'output' in stats and len(op['outputIndexes']) > 0: tensor_idx = op['outputIndexes'][0] if tensor_idx not in quant_info_dict: scale, zero = stats['output'] quant_info_dict[tensor_idx] = { 'index': tensor_idx, 'quantInfo': { 'scale': scale, 'zero': zero, 'min': min_val, 'max': max_val, "type": data_type } } continue if should_process: try: import re match = re.search(r'(?:blocks|layers)\.(\d+)', op_name) if match: layer_idx = int(match.group(1)) else: continue except: continue if layer_idx >= len(self.act_dict): continue layer_act_dict = self.act_dict[layer_idx] matched_pt_name = self._find_match_in_dict(op_name, layer_act_dict) if matched_pt_name: stats = layer_act_dict[matched_pt_name] if 'input' in stats and len(op['inputIndexes']) > 0: tensor_idx = op['inputIndexes'][0] if tensor_idx not in quant_info_dict: scale, zero = stats['input'] quant_info_dict[tensor_idx] = { 'index': tensor_idx, 'quantInfo': { 'scale': scale, 'zero': zero, 'min': min_val, 'max': max_val, "type": data_type } } if 'output' in stats and len(op['outputIndexes']) > 0: tensor_idx = op['outputIndexes'][0] if tensor_idx not in quant_info_dict: scale, zero = stats['output'] quant_info_dict[tensor_idx] = { 'index': tensor_idx, 'quantInfo': { 'scale': scale, 'zero': zero, 'min': min_val, 'max': max_val, "type": data_type } } if self.generate_for_npu: print(f"Initial collected tensors: {len(quant_info_dict)}") self._propagate_quant_info(mnn['oplists'], quant_info_dict) print(f"final collected tensors: {len(quant_info_dict)}") mnn['extraTensorDescribe'] = list(quant_info_dict.values()) with open(base_path, 'w', encoding='utf-8') as f: json.dump(mnn, f, ensure_ascii=False, indent=4) return base_path