import torch import logging import gc import functools import json import inspect from typing import Dict from tqdm import tqdm from collections import defaultdict import math logging.basicConfig(level=logging.ERROR) class ACIQ: def __init__(self, size): self.num_bits = size # TODO: expose as cmd line parameters self.stochastic = False self.int_exp = False self.enforce_true_zero = True #params['true_zero'] self.alpha_gaus = {2: 1.71, 3: 2.15, 4: 2.55, 5: 2.93, 6: 3.28, 7: 3.61, 8: 3.92} self.alpha_laplace = {2: 2.83, 3: 3.89, 4: 5.03, 5: 6.2, 6: 7.41, 7: 8.64, 8: 9.89} self.gaussian_const = (0.5 * 0.35) * (1 + (math.pi * math.log(4)) ** 0.5) def alpha2DeltaOffset(self, alpha, max_value, min_value, mean): max_range = max_value - min_value if alpha <= 0 or alpha >= max_range / 2: delta = max_range else: delta = 2 * alpha min_value = max(min_value, mean - delta / 2) return delta, min_value def gemmlowpClippingQuantize(self, input): min_value = input.min() max_value = input.max() mean = input.mean() alpha = self.get_alpha_gaus(input) # gaussian clipping delta, min_value = self.alpha2DeltaOffset(alpha, max_value, min_value, mean) return torch.stack([delta + min_value, min_value], 0) def get_max_min(self, x): if self.num_bits > 8: return torch.stack([x.max(), x.min()], 0) return self.gemmlowpClippingQuantize(x) def get_alpha_gaus(self, tensor): N = 1 for i in range(len(tensor.shape)): N *= tensor.shape[i] min_value = tensor.min() max_value = tensor.max() std = ((max_value - min_value) * self.gaussian_const) / ((2 * math.log(N)) ** 0.5) return self.alpha_gaus[self.num_bits] * std class SmoothQuantizer: def __init__( self, model, n_parallel_calib_samples=None, max_calib_samples=128, max_calib_seq_len=512, alpha=0.5, act_bit=8, act_sym=True, generate_for_npu=False ) -> None: self.act_sym = act_sym self.model = model self.tokenizer = model.tokenizer #self.w_bit = model.args.quant_bit self.act_bit = act_bit self.group_size = model.args.quant_block self.alpha = alpha self.generate_for_npu = generate_for_npu self.max_calib_samples = max_calib_samples self.max_calib_seq_len = max_calib_seq_len self.split = 'train' self.calib_data = 'wikitext' if model.args.calib_data is None else model.args.calib_data self.best_device = SmoothQuantizer.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)): SmoothQuantizer.to_device(self.modules[idx], "cpu") self.act_scales = [{} for _ in range(len(self.modules))] self.act_dict = [defaultdict(dict) for _ in range(len(self.modules))] self.n_parallel_calib_samples = n_parallel_calib_samples self.samples = self.init_quant( n_samples=self.max_calib_samples, max_seq_len=self.max_calib_seq_len, ) @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 # dataset = load_dataset(data, split=split) with open(data, 'r', encoding='utf-8') as f: dataset = f.read().splitlines() # dataset = dataset.shuffle(seed=42) else: raise NotImplementedError( "Either pass a string to a huggingface dataset" "that is preprocessed with one sample of text per element" ) samples = [] if custom_calib_data == False: dataset = dataset.shuffle(seed=42) for i in range(n_samples): input_ids = tokenizer( dataset[i]["text"], return_tensors="pt", max_length=max_seq_len, truncation=True ).input_ids samples.append(input_ids) else: for i in range(n_samples): 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 samples.append(input_ids) return samples @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 clear_memory(weight=None): if weight is not None: del weight gc.collect() torch.cuda.empty_cache() @staticmethod def clear_block_cache(block): if hasattr(block, 'self_attn') and block.self_attn is not None: attn = block.self_attn if hasattr(attn, 'past_key_value'): attn.past_key_value = None if hasattr(attn, 'conv_state'): attn.conv_state = None if hasattr(attn, 'rnn_state'): attn.rnn_state = None def clear_all_block_caches(self): for block in self.modules: self.clear_block_cache(block) def init_quant(self, n_samples=128, max_seq_len=512): samples = SmoothQuantizer.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): layer_kwargs = {} seq_len = sample.numel() new_tokens = 0 inps = self.model.embedding(sample).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) del sample SmoothQuantizer.clear_memory() return layer_kwargs, inps def _get_max_input(self, idx, layer, named_linears): def infer_feature_dim(module): if hasattr(module, "in_features"): return module.in_features if hasattr(module, "in_channels"): return module.in_channels if hasattr(module, "weight") and getattr(module, "weight", None) is not None: weight = module.weight if weight.dim() == 1: return weight.numel() return None def stat_tensor(name, tensor, module): feature_dim = infer_feature_dim(module) if tensor.dim() == 3 and feature_dim is not None: if tensor.shape[-1] == feature_dim: pass elif tensor.shape[1] == feature_dim: tensor = tensor.transpose(1, 2).contiguous() hidden_dim = tensor.shape[-1] tensor = tensor.reshape(-1, hidden_dim).abs().detach() comming_max = torch.max(tensor, dim=0)[0].float().cpu() if name in self.act_scales[idx]: self.act_scales[idx][name] = torch.max(self.act_scales[idx][name], comming_max) else: self.act_scales[idx][name] = comming_max def stat_input_hook(m, x, y, name): if isinstance(x, tuple): x = x[0] stat_tensor(name, x, m) handles = [] for name in named_linears: handles.append( named_linears[name].register_forward_hook( functools.partial(stat_input_hook, name=name) ) ) layer_kwargs = self._select_layer_kwargs(layer, self.module_kwargs) module_kwargs = self._sanitize_kwargs(layer_kwargs, layer) self.inps = self._module_forward(self.inps, layer, module_kwargs) for h in handles: h.remove() def _sanitize_kwargs(self, inputs_kwargs, module): """ Remove the arguments that are not supported in the module's forward pass to avoid breaking behaviour between different versions of transformers. Args: inputs_kwargs (`dict`): The input dictionary to pass to the model layer module (`torch.nn.Module`): Target module to quantize. """ 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): 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 @torch.no_grad() def _module_forward( self, x: torch.Tensor, module: torch.nn.Module, module_kwargs: Dict ) -> torch.Tensor: if self.n_parallel_calib_samples is None: # runs through all samples at once # print(module, x, module_kwargs); exit(0) module_output = module(x, **module_kwargs) if isinstance(module_output, tuple): module_output = module_output[0] else: # memory efficiently runs through all calibration samples # but only n_parallel_calib_samples at a time module_output = [] partitioned_inputs = torch.split(x, self.n_parallel_calib_samples) for x_partial in partitioned_inputs: partial_output = module(x_partial, **module_kwargs) if isinstance(partial_output, tuple): partial_output = partial_output[0] module_output.append(partial_output.cpu()) module_output = torch.cat(module_output, dim=0) return module_output @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 get_named_linears(module): linears = {} for name, m in module.named_modules(): # 兼容更多潜在的线性层类型 if isinstance(m, torch.nn.Linear) or m.__class__.__name__ == 'Linear': linears[name] = m return linears @staticmethod def get_all_leaf_modules(module): targets = {} for name, submod in module.named_modules(): if name == "": continue if len(list(submod.children())) == 0: if isinstance(submod, (torch.nn.Dropout, torch.nn.Identity)): continue targets[name] = submod return targets @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 @torch.no_grad() def smooth_ln_fcs(ln, fcs, act_scales, alpha=0.5): if not isinstance(fcs, list): fcs = [fcs] if not SmoothQuantizer.is_allowed_norms(ln): raise NotImplementedError( f"LayerNorm {ln} is not supported for smooth quantization." ) for fc in fcs: assert isinstance(fc, torch.nn.Linear) assert ln.weight.numel() == fc.in_features == act_scales.numel() device, dtype = fcs[0].weight.device, fcs[0].weight.dtype act_scales = act_scales.to(device=device, dtype=dtype) weight_scales = torch.cat( [fc.weight.abs().max(dim=0, keepdim=True)[0] for fc in fcs], dim=0 ) weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5) scales = ( (act_scales.pow(alpha) / weight_scales.pow(1 - alpha)) .clamp(min=1e-5) .to(device) .to(dtype) ) if SmoothQuantizer.is_offset_rmsnorm(ln): ln.weight += 1 ln.weight.div_(scales) ln.weight -= 1 else: ln.weight.div_(scales) if hasattr(ln, "bias") and ln.bias is not None: ln.bias.div_(scales) for fc in fcs: fc.weight.mul_(scales.view(1, -1)) @staticmethod def is_allowed_norms(op): if isinstance(op, torch.nn.LayerNorm): return True if any(t in str(type(op)) for t in ['LlamaRMSNorm', 'GemmaRMSNorm', 'CohereLayerNorm']): return True if "rmsnorm" in str(op.__class__).lower(): return True return False def _apply_scale(self, idx, module): model_type = getattr(self.model.config, "model_type", "") layer_type = getattr(module, "layer_type", None) if model_type in ("qwen3_5", "qwen3_5_moe"): if layer_type == "linear_attention" and hasattr(module, "linear_attn"): attn_ln = module.input_layernorm linear_attn = module.linear_attn fcs = [] for name in ("in_proj_qkv", "in_proj_a", "in_proj_b", "in_proj_z"): fc = getattr(linear_attn, name, None) if fc is not None: fcs.append(fc) if fcs and "linear_attn.in_proj_qkv" in self.act_scales[idx]: input_scales = self.act_scales[idx]["linear_attn.in_proj_qkv"] SmoothQuantizer.smooth_ln_fcs(attn_ln, fcs, input_scales, self.alpha) return if hasattr(module.self_attn, 'q_proj') and hasattr(module.self_attn, 'k_proj') and hasattr(module.self_attn, 'v_proj'): attn_ln = module.input_layernorm qkv = [ module.self_attn.q_proj, module.self_attn.k_proj, module.self_attn.v_proj, ] if "self_attn.q_proj" in self.act_scales[idx]: qkv_input_scales = self.act_scales[idx]["self_attn.q_proj"] SmoothQuantizer.smooth_ln_fcs(attn_ln, qkv, qkv_input_scales, self.alpha) return if hasattr(module.self_attn, 'q_proj') and hasattr(module.self_attn, 'k_proj') and hasattr(module.self_attn, 'v_proj'): attn_ln = module.input_layernorm qkv = [ module.self_attn.q_proj, module.self_attn.k_proj, module.self_attn.v_proj, ] qkv_input_scales = self.act_scales[idx]["self_attn.q_proj"] SmoothQuantizer.smooth_ln_fcs(attn_ln, qkv, qkv_input_scales, self.alpha) ffn_ln = module.post_attention_layernorm # feed forward norm fcs = [module.mlp.gate_proj, module.mlp.up_proj] ffn_input_scales = self.act_scales[idx]["mlp.gate_proj"] SmoothQuantizer.smooth_ln_fcs(ffn_ln, fcs, ffn_input_scales, self.alpha) @torch.no_grad() def _get_all_static_scales(self, idx, layer, named_linears): def stat_io_hook(m, x, y, name): if isinstance(x, tuple): x = x[0] max_min = self.act_quanter.get_max_min(x) if name not in self.act_dict[idx] or "input" not in self.act_dict[idx][name]: self.act_dict[idx][name]["input"] = max_min else: self.act_dict[idx][name]["input"] = max_min * (1-self.moment) + self.moment * self.act_dict[idx][name]["input"] if isinstance(y, tuple): y = y[0] max_min = self.act_quanter.get_max_min(y) if name not in self.act_dict[idx] or "output" not in self.act_dict[idx][name]: self.act_dict[idx][name]["output"] = max_min else: self.act_dict[idx][name]["output"] = 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, self.module_kwargs) module_kwargs = self._sanitize_kwargs(layer_kwargs, layer) self.inps = self._module_forward(self.inps, layer, module_kwargs) for h in handles: h.remove() @torch.no_grad() def _extract_static_scales(self): print("Extracting static 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 zero = 0.0 return [scale, zero] 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() # Assume has zeropoint 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: scale = 1.0 else: scale = (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) elif self.act_bit == 16: print("Error: CPU only supports 8 bit feature map quantized") return [scale, zero] if self.act_sym: func = compute_scale_sym else: func = compute_scale_zero_asym for idx in range(len(self.modules)): 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 quantize(self): for i in tqdm(range(len(self.samples)), desc="collecting data and computing scales..."): sample = self.samples[i] if sample.numel() == 0: continue self.clear_all_block_caches() self.module_kwargs, self.inps = self._get_first_input(sample) for idx in range(len(self.modules)): SmoothQuantizer.to_device(self.modules[idx], self.best_device) self.clear_block_cache(self.modules[idx]) if self.module_kwargs.get("position_ids", None) is not None: self.module_kwargs["position_ids"] = self.module_kwargs["position_ids"].to(self.best_device) if self.module_kwargs.get("attention_mask", None) is not None: self.module_kwargs["attention_mask"] = self.module_kwargs["attention_mask"].to(self.best_device) named_layers = SmoothQuantizer.get_all_leaf_modules(self.modules[idx]) self._get_max_input(idx, self.modules[idx], named_layers) if "cpu" != self.best_device: SmoothQuantizer.to_device(self.modules[idx], "cpu") for idx in tqdm(range(len(self.modules)), desc="applying scales..."): self._apply_scale(idx, self.modules[idx]) for i in tqdm(range(len(self.samples)), desc="collecting static activation scales..."): sample = self.samples[i] if sample.numel() == 0: continue self.clear_all_block_caches() self.module_kwargs, self.inps = self._get_first_input(sample) for idx in range(len(self.modules)): SmoothQuantizer.to_device(self.modules[idx], self.best_device) self.clear_block_cache(self.modules[idx]) if self.module_kwargs.get("position_ids", None) is not None: self.module_kwargs["position_ids"] = self.module_kwargs["position_ids"].to(self.best_device) if self.module_kwargs.get("attention_mask", None) is not None: self.module_kwargs["attention_mask"] = self.module_kwargs["attention_mask"].to(self.best_device) named_linears = SmoothQuantizer.get_all_leaf_modules(self.modules[idx]) self._get_all_static_scales(idx, self.modules[idx], named_linears) if "cpu" != self.best_device: SmoothQuantizer.to_device(self.modules[idx], "cpu") self._extract_static_scales() SmoothQuantizer.clear_memory() for idx in range(len(self.modules)): SmoothQuantizer.to_device(self.modules[idx], "cpu") def _find_match_in_dict(self, mnn_op_name, layer_act_dict): """ mnn_op_name: e.g., '/blocks.0/self_attn/q_norm/Mul_1_output_0' layer_act_dict: { 'self_attn.q_norm': {...}, 'self_attn.q_proj': {...} } """ 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): """ 量化参数传导机制。 通过图的拓扑结构，将已知的量化参数传递给相邻的未知 Tensor。 """ import copy # 定义透传算子：输入和输出共享 Scale/Zero # 这些算子不改变数值分布，只改变形状或排布 PASS_THROUGH_OPS = [ 'Reshape', 'Squeeze', 'Unsqueeze', 'Flatten', 'Transpose', 'Permute', 'ConvertTensor', 'Cast', 'Slice', 'StridedSlice', 'Split', 'Concat', 'Pack' ] # 定义特殊的单向或部分传导算子 # Gather: Output Scale == Input[0] (Data) Scale. (Input[1] 是 indices，不需要) 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 # ----------------------------------------------- # 策略 1: 透传算子 (双向传导) # Input <-> Output # ----------------------------------------------- if op_type in PASS_THROUGH_OPS: # 1. Forward: 任意 Input 有参数 -> 传给所有 Output # 通常取第一个有参数的 input 作为源 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 # 2. Backward: 任意 Output 有参数 -> 传给所有 Input # (仅当 Input 还没有参数时) 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 # ----------------------------------------------- # 策略 2: Gather 类 (仅数据输入 <-> 输出) # ----------------------------------------------- elif op_type in DATA_SELECT_OPS: data_idx = inputs[0] # 第0个是 params/data 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 # ----------------------------------------------- # 策略 3: BinaryOp (Add/Mul) - 谨慎处理 # 通常用于 Residual Add。 # ----------------------------------------------- elif op_type == 'BinaryOp': out_idx = outputs[0] # Backward: # 如果 Add 的输出已知（通常是因为连着下一个 Linear/Norm 的输入）， # 我们可以尝试推导输入的 Scale。 # 注意：对于 Add，如果 Input A 和 Input B 的范围差异巨大，直接回传可能有风险。 # 但在 Transformer 残差结构中，通常 Input 和 Output 的 Scale 是同数量级的。 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 # Forward: # 如果所有输入都有 Scale，取 Scale 最大的那个传给输出 # (保守策略，避免截断) 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: # 找到 Scale 最大的那个 input 的 info 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', '') if 'lm_head' in op_name: continue should_process = False if not self.generate_for_npu: should_process = (op_type == 'Convolution') else: should_process = (op_type not in npu_ignore_types) 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