import os import sys import copy import json import torch import numpy as np from .torch_utils import quant as torch_quant from .torch_utils import onnx_export from tqdm import tqdm from .spinner import spinner_run from .gptq import GPTQ from .lora import LoRA EXPORT_LOG = '.export.log' class MNNConverter: def __init__(self, exporter, weight_ops = None): self.weight_ops = weight_ops self.exporter = exporter self.args = exporter.args self.mnn_weight_offset = 0 if os.path.exists(self.args.mnnconvert): self.mnnconvert = self.args.mnnconvert else: self.mnnconvert = None self.lm_weight = None self.tie_embeddings_info = None def transformer_c4_args(self): if getattr(self.args, 'transformer_c4', False): return ['--transformerFuseC4=1'] return [] def convert(self, convert_args): import contextlib log_fp = open(EXPORT_LOG, "a") sfd = None try: sfd = os.dup(1) log_fd = log_fp.fileno() # mnnconvert ... > .export.log os.dup2(log_fd, 1) except Exception: if sfd is not None: os.close(sfd) sfd = None try: with contextlib.redirect_stdout(log_fp): sys.argv = convert_args sys.argc = len(convert_args) if self.mnnconvert is None: from MNN.tools import mnnconvert mnnconvert.main() else: convert_args[0] = self.mnnconvert cmd = ' '.join(convert_args) message = os.popen(cmd).read() print(message) sys.argv = [] finally: if sfd is not None: os.dup2(sfd, 1) os.close(sfd) log_fp.close() @spinner_run(f'convert onnx model to ') def onnx2mnn(self, onnx_path, mnn_path, args = [], transformer_fuse = True, group_conv_native = False, weight_sym = False, save_external_data = True): convert_args = [ '', '-f', 'ONNX', '--modelFile', str(onnx_path), '--MNNModel', str(mnn_path), '--allowCustomOp' ] if transformer_fuse: convert_args += ['--transformerFuse'] if group_conv_native: convert_args += ['--groupConvNative'] if weight_sym: convert_args += ['--weightQuantAsymmetric=0'] if save_external_data: convert_args += ['--saveExternalData'] if self.args.hqq: convert_args += ['--hqq'] convert_args += self.transformer_c4_args() convert_args += args self.convert(convert_args) return mnn_path def mnn2json(self, mnn_path, json_path): convert_args = [ '', '-f', 'MNN', '--modelFile', str(mnn_path), '--JsonFile', str(json_path) ] self.convert(convert_args) return json_path def json2mnn(self, json_path, mnn_path): convert_args = [ '', '-f', 'JSON', '--modelFile', str(json_path), '--MNNModel', str(mnn_path) ] convert_args += self.transformer_c4_args() self.convert(convert_args) return mnn_path def removeDupOps(self, mnn_path): convert_args = [ '', '-f', 'MNN', '--modelFile', str(mnn_path), '--MNNModel', str(mnn_path), '--optimizeLevel=1' ] convert_args += self.transformer_c4_args() self.convert(convert_args) return mnn_path def export(self, onnx_path, quant_bit = None, quant_block = None, transformer_fuse = True, group_conv_native = False, weight_sym = None): self.onnx_model_path = onnx_path self.mnn_name = os.path.basename(onnx_path).replace('.onnx', '.mnn') self.mnn_model_path = os.path.join(self.args.dst_path, self.mnn_name) self.mnn_weight_path = f'{self.mnn_model_path}.weight' if self.weight_ops is None: if quant_bit is None: quant_bit = self.args.quant_bit if quant_block is None: quant_block = self.args.quant_block if weight_sym is None: weight_sym = self.args.sym if quant_bit == 16: quant_args = ['--fp16'] else: quant_args = [ '--weightQuantBits', str(quant_bit), '--weightQuantBlock', str(quant_block) ] if quant_bit == 32: quant_args = [] self.onnx2mnn(self.onnx_model_path, self.mnn_model_path, quant_args, transformer_fuse=transformer_fuse, group_conv_native=group_conv_native, weight_sym=weight_sym) else: mnn_json = f'{self.mnn_model_path}.json' self.onnx2mnn(self.onnx_model_path, self.mnn_model_path, transformer_fuse=transformer_fuse, group_conv_native=group_conv_native, weight_sym=weight_sym) self.mnn2json(self.mnn_model_path, mnn_json) self.rebuild(mnn_json) self.json2mnn(mnn_json, self.mnn_model_path) self.removeDupOps(self.mnn_model_path) self.mnn2json(self.mnn_model_path, mnn_json) if self.args.gptq_path is not None: self.apply_gptq(mnn_json) if self.args.lora_path is not None and self.args.lora_split: self.export_lora(mnn_json) if self.args.omni: self.export_omni_quant(mnn_json) if self.args.smooth: self.export_smooth_quant(mnn_json) return self.tie_embeddings_info def get_experts_graphs(self, experts): hidden_states = torch.randn((1, self.exporter.config.hidden_size)) layers_num = len(experts) expert_num = len(experts[0]) dummy_expert = experts[0][0] onnx_model = f'{self.exporter.onnx_path}/expert.onnx' onnx_export( dummy_expert, (hidden_states), onnx_model, input_names=['hidden_states'], output_names=['hidden_states']) mnn_model = f'{onnx_model}.mnn' mnn_json = f'{mnn_model}.json' self.onnx2mnn(onnx_model, mnn_model) self.mnn2json(mnn_model, mnn_json) expert_graph = json.load(open(mnn_json, 'rt')) tensors = expert_graph['tensorName'] nodes = expert_graph['oplists'] # get input and output inputs = [] outputs = [] for node in nodes: if node['type'] == 'Input': inputs.append(node['outputIndexes'][0]) for output_name in expert_graph['outputName']: outputs.append(tensors.index(output_name)) # Use actual layer indices (for models where not all layers have MoE) expert_layer_ids = getattr(self.exporter, 'expert_layer_ids', list(range(layers_num))) subgraphs = [] for i in range(layers_num): layer_idx = expert_layer_ids[i] for j in range(expert_num): ijnodes = copy.deepcopy(nodes) for op in ijnodes: if op['type'] == 'Extra': for attr in op['main']['attr']: if attr['key'] == 'name': names = attr['s'].split('/') names[2] = f'{layer_idx}_{j}' attr['s'] = '/'.join(names) subgraph = { 'name': f'/expert/{layer_idx}_{j}', 'inputs': inputs, 'outputs': outputs, 'tensors': copy.deepcopy(tensors), 'nodes': ijnodes } subgraphs.append(subgraph) return subgraphs @spinner_run(f'apply gptq to ') def apply_gptq(self, mnn_json): GPTQ(self.args.gptq_path).apply(mnn_json, self.mnn_weight_path) return self.mnn_weight_path @spinner_run(f'export split lora to ') def export_lora(self, mnn_json): lora_model = os.path.join(self.args.dst_path, 'lora.mnn') lora_json = f'{lora_model}.json' LoRA(self.args.lora_path).apply(mnn_json, lora_json) self.json2mnn(lora_json, lora_model) if os.path.exists(lora_json): os.remove(lora_json) return lora_model @spinner_run(f'export smooth quant scale to ') def export_smooth_quant(self, mnn_json): self.exporter.smooth_quantizer.apply(mnn_json) self.json2mnn(mnn_json, self.mnn_model_path) return self.mnn_model_path @spinner_run(f'export omni quant scale to ') def export_omni_quant(self, mnn_json): self.exporter.omni_quantizer.apply(mnn_json) self.json2mnn(mnn_json, self.mnn_model_path) return self.mnn_model_path @spinner_run(f'quant model weight to ', True) def rebuild(self, json_path): mnn_graph = json.load(open(json_path, 'rt')) has_experts = hasattr(self.exporter, 'experts') and len(self.exporter.experts) > 0 if has_experts: subgraphs = self.get_experts_graphs(self.exporter.experts) mnn_graph['subgraphs'] = subgraphs new_ops = [] # Load layernorm weight from external with open(self.mnn_weight_path, 'rb') as f: for op in tqdm(mnn_graph['oplists'], 'Load LayerNorm data'): if op['type'] == 'LayerNorm' and 'external' in op['main']: external = op['main']['external'] f.seek(external[0]) op['main']['gamma'] = np.frombuffer(f.read(external[1]), np.float32).tolist() op['main']['beta'] = np.frombuffer(f.read(external[2]), np.float32).tolist() del op['main']['external'] if op['type'] == 'Const' and 'external' in op['main']: external = op['main']['external'] f.seek(external[0]) op['main']['float32s'] = np.frombuffer(f.read(external[1]), np.float32).tolist() del op['main']['external'] # Rebuild ops with open(self.mnn_weight_path, 'wb') as self.mnn_weight: for op in tqdm(mnn_graph['oplists'], 'Quant weights'): if op['type'] == 'Extra' or op['type'] == 'LayerNorm': new_ops += self.rebuild_op(op, mnn_graph) else: new_ops.append(op) mnn_graph['oplists'] = new_ops if has_experts and 'subgraphs' in mnn_graph: for subgraph in tqdm(mnn_graph['subgraphs'], 'Quant subgraphs weights'): new_subops = [] for op in subgraph['nodes']: if op['type'] == 'Extra' or op['type'] == 'LayerNorm': new_subops += self.rebuild_op(op, subgraph) else: new_subops.append(op) subgraph['nodes'] = new_subops with open(json_path, 'w', encoding='utf-8') as file: json.dump(mnn_graph, file, ensure_ascii=False, indent=4) return self.mnn_weight_path def quant(self, weight, quant_bit, quant_block, symmetric): if self.exporter.args.skip_weight: # Skip expensive quantization when skip_weight is enabled oc, ic = weight.shape if quant_block == 0: block_size = ic else: block_size = quant_block block_num = ic // block_size # alpha: oc * block_num (symmetric) or oc * block_num * 2 (asymmetric) alpha_num = oc * block_num * (1 if symmetric else 2) alpha = torch.zeros(alpha_num, dtype=torch.float32) # q_weight: raw size is oc * ic, packed size depends on quant_bit # bits < 8 packing logic matches repack_low_bits q_weight_num = (oc * ic * quant_bit + 7) // 8 q_weight = torch.zeros(q_weight_num, dtype=torch.uint8) return q_weight, alpha q_weight, alpha = torch_quant(weight.cpu(), quant_bit, quant_block, symmetric, self.args.awq, self.args.hqq) return q_weight, alpha def write_weight(self, data): if isinstance(data, torch.Tensor): data = data.numpy() if isinstance(data, list): data = np.array(data).astype(np.float32) return self.mnn_weight.write(data.tobytes()) def write_header(self, ic, oc, quant_bit): dim_num = self.mnn_weight.write(b'\x02') shape_dtype = np.int16 if oc > 65535 or ic > 65535: shape_dtype = np.int32 dim_length = self.write_weight(np.array([oc, ic]).astype(shape_dtype)) offset = 1 << (quant_bit - 1) weight_map = [i for i in range(-offset, offset)] if len(weight_map) == 256: weight_map.insert(0, 0) else: weight_map.insert(0, len(weight_map)) map_length = self.write_weight(np.array(weight_map, dtype=np.int8)) header_length = dim_num + dim_length + map_length return header_length, shape_dtype == np.int32 def build_weight(self, linear, quant_bit, quant_block, symmetric): ic, oc = linear.in_features, linear.out_features if quant_bit == 16: if self.exporter.args.skip_weight: # Use a small dummy buffer and skip full weight loading/conversion weight_len = (ic * oc * 2) self.mnn_weight.seek(weight_len, 1) else: half_weight = linear.weight.data.flatten().half() weight_len = self.write_weight(half_weight) alpha_len, q_min, shape_int32, header_len = 0, 0, False, 0 else: q_min = 1 assert(quant_bit in (1, 2, 3, 4, 8)) q_weight, alpha = self.quant(linear.weight.data, quant_bit, quant_block, symmetric) header_len, shape_int32 = self.write_header(ic, oc, quant_bit) scale_fp16 = (self.args.scale_bit == 16) alpha_dtype_size = 2 if scale_fp16 else 4 if self.exporter.args.skip_weight: weight_len = len(q_weight) + header_len self.mnn_weight.seek(len(q_weight), 1) alpha_len = len(alpha) * alpha_dtype_size self.mnn_weight.seek(alpha_len, 1) else: weight_len = self.write_weight(q_weight) + header_len if scale_fp16: alpha_np = alpha.numpy() if hasattr(alpha, 'numpy') else np.asarray(alpha) alpha_fp16 = alpha_np.astype(np.float16) alpha_len = self.write_weight(alpha_fp16) else: alpha_len = self.write_weight(alpha) if linear.bias is not None: bias_length = (oc * 4) if self.exporter.args.skip_weight: self.mnn_weight.seek(bias_length, 1) else: bias = linear.bias.data.flatten().float() bias_length = self.write_weight(bias) else: bias_length = 0 external = [self.mnn_weight_offset, weight_len, alpha_len, bias_length, 0] self.mnn_weight_offset += (weight_len + alpha_len + bias_length) return external, q_min, shape_int32, header_len def build_tensor(self, graph, tensor_name): tensor_key = 'tensorName' if tensor_key not in graph and 'tensors' in graph: tensor_key = 'tensors' tensor_idx = [len(graph[tensor_key])] graph[tensor_key].append(tensor_name) return tensor_idx def rebuild_op(self, op, graph): if "type" in op['main']: op_type = op['main']['type'] else: op_type = op['type'] if op_type == 'FakeLinear': return self.rebuild_linear(op, graph) if op_type == 'FusedAttention': return self.rebuild_attnention(op, graph) if op_type == 'FusedRoPE': return self.rebuild_rope(op, graph) if op_type == 'FusedLinearAttention': return self.rebuild_linear_attnention(op, graph) if op_type == "LayerNorm": return self.rebuild_layernorm(op, graph) if op_type == 'MoE': return self.rebuild_moe(op, graph) return None def const_float_data(self, graph, tensor_index): op_key = 'oplists' if 'oplists' in graph else 'nodes' for op in graph[op_key]: if tensor_index not in op.get('outputIndexes', []): continue if op.get('type') != 'Const': break main = op.get('main', {}) if 'float32s' in main: return main['float32s'] break return None def rebuild_rope(self, op, graph): attrs = op['main']['attr'] name = op['name'] rope_cut_head_dim = 0 q_norm = False k_norm = False q_norm_eps = 0.0 k_norm_eps = 0.0 for attr in attrs: if attr['key'] == 'name': name = attr['s'] elif attr['key'] == 'rope_cut_head_dim': rope_cut_head_dim = attr['i'] elif attr['key'] == 'q_norm': q_norm = bool(attr['i']) elif attr['key'] == 'k_norm': k_norm = bool(attr['i']) elif attr['key'] == 'q_norm_eps': q_norm_eps = attr['f'] elif attr['key'] == 'k_norm_eps': k_norm_eps = attr['f'] rope_param = { "rope_cut_head_dim": rope_cut_head_dim, } input_indexes = op['inputIndexes'] if q_norm or k_norm: if len(input_indexes) < 6: raise RuntimeError(f'FusedRoPE {name} misses q/k norm inputs') if q_norm: q_gamma = self.const_float_data(graph, input_indexes[4]) if q_gamma is None: raise RuntimeError(f'FusedRoPE {name} misses q_norm gamma const') rope_param["q_norm"] = { "axis": [-1], "epsilon": q_norm_eps, "gamma": q_gamma, "useRMSNorm": True } if k_norm: k_gamma = self.const_float_data(graph, input_indexes[5]) if k_gamma is None: raise RuntimeError(f'FusedRoPE {name} misses k_norm gamma const') rope_param["k_norm"] = { "axis": [-1], "epsilon": k_norm_eps, "gamma": k_gamma, "useRMSNorm": True } rope_op = { "inputIndexes": input_indexes[:4], "main_type": "RoPEParam", "main": rope_param, "name": name, "outputIndexes": op['outputIndexes'], "type": "RoPE", "defaultDimentionFormat": op['defaultDimentionFormat'] } return [rope_op] def rebuild_moe(self, op, graph): moe = copy.deepcopy(op) moe['main'] = { 'attr': moe['main']['attr'][:3] } moe['type'] = 'MoE' return [moe] def rebuild_layernorm(self, op, graph): if "gamma" not in op['main'] or "beta" not in op['main']: return [op] attr = op['main'] gamma = attr['gamma'] beta = attr['beta'] gamma_len = self.write_weight(gamma) beta_len = self.write_weight(beta) del attr['gamma'] del attr['beta'] external = [self.mnn_weight_offset, gamma_len, beta_len] self.mnn_weight_offset += (gamma_len + beta_len) attr['external'] = external layernorm_op = { "name": op['name'], "inputIndexes": op['inputIndexes'], "outputIndexes": op['outputIndexes'], "type": "LayerNorm", "main_type": "LayerNorm", "main": attr, "defaultDimentionFormat": op['defaultDimentionFormat'] } return [layernorm_op] def rebuild_attnention(self, op, graph): attrs = op['main']['attr'] layer_index = -1 kv_shared_layer_index = -1 for attr in attrs: if attr['key'] == 'name': name = attr['s'] elif attr['key'] == 'kv_cache': kv_cache = attr['i'] elif attr['key'] == 'layer_index': layer_index = attr.get('i', -1) elif attr['key'] == 'kv_shared_layer_index': kv_shared_layer_index = attr.get('i', -1) origin_input = op['inputIndexes'] origin_output = op['outputIndexes'] main_dict = { "kv_cache": bool(kv_cache), "layer_index": layer_index, "kv_shared_layer_index": kv_shared_layer_index, } fused_attention = { "inputIndexes": origin_input, "main_type": "AttentionParam", "main": main_dict, "name": name, "outputIndexes": origin_output, "type": "Attention", "defaultDimentionFormat": "NHWC" } return [fused_attention] def rebuild_linear_attnention(self, op, graph): attrs = op['main']['attr'] num_k_heads = 0 num_v_heads = 0 head_k_dim = 0 head_v_dim = 0 attn_type = "gated_delta_rule" use_qk_l2norm = False name = "" # Parse attributes from Custom Op for attr in attrs: if attr['key'] == 'name': name = attr['s'] elif attr['key'] == 'num_k_heads': num_k_heads = attr['i'] elif attr['key'] == 'num_v_heads': num_v_heads = attr['i'] elif attr['key'] == 'head_k_dim': head_k_dim = attr['i'] elif attr['key'] == 'head_v_dim': head_v_dim = attr['i'] elif attr['key'] == 'attn_type': attn_type = attr['s'] elif attr['key'] == 'use_qk_l2norm': use_qk_l2norm = bool(attr['i']) input_indexes = op['inputIndexes'] output_indexes = op['outputIndexes'] linear_attention_param = { "attn_type": attn_type, "num_k_heads": num_k_heads, "num_v_heads": num_v_heads, "head_k_dim": head_k_dim, "head_v_dim": head_v_dim, "use_qk_l2norm": use_qk_l2norm } fused_linear_attention = { "inputIndexes": input_indexes, "main_type": "LinearAttentionParam", "main": linear_attention_param, "name": name, "outputIndexes": output_indexes, "type": "LinearAttention", "defaultDimentionFormat": "NHWC" } return [fused_linear_attention] def get_extra_attr(self, op, key, default=None): for attr in op.get('main', {}).get('attr', []): if attr.get('key') != key: continue if 's' in attr: return attr['s'] if 'i' in attr: return attr['i'] if 'f' in attr: return attr['f'] return default def rebuild_linear(self, op, graph): attrs = op['main']['attr'] for attr in attrs: if attr['key'] == 'name': name = attr['s'] elif attr['key'] == "in_features": ic = attr["i"] elif attr['key'] == "out_features": oc = attr["i"] elif attr['key'] == "has_bias": has_bias = attr["i"] linear = self.weight_ops[name] assert(linear.in_features == ic and linear.out_features == oc and (linear.bias is not None) == has_bias) is_lm = 'lm_head' in name quant_bit = self.args.lm_quant_bit if is_lm else self.args.quant_bit quant_block = self.args.lm_quant_block if is_lm else self.args.quant_block quant_sym = self.args.sym if self.args.quant_config is not None: with open(self.args.quant_config, 'r') as f: quant_config = json.load(f) if name in quant_config: op_config = quant_config[name] quant_bit = op_config.get('bits', quant_bit) quant_block = op_config.get('block_size', quant_block) quant_sym = op_config.get('symmetric', quant_sym) block_size = ic if quant_block == 0 else quant_block if is_lm and self.lm_weight is not None: external, q_min, shape_int32, header_len = self.lm_weight else: external, q_min, shape_int32, header_len = self.build_weight(linear, quant_bit, quant_block, quant_sym) if is_lm and self.lm_weight is None: self.lm_weight = [external, q_min, shape_int32, header_len] if is_lm and self.args.tie_word_embeddings: weight_offset = external[0] + header_len alpha_offset = external[0] + external[1] alpha_size = external[2] self.tie_embeddings_info = { "weight_offset": weight_offset, "alpha_offset": alpha_offset, "alpha_size": alpha_size, "quant_bit": quant_bit, "quant_block": quant_block, "alpha_dtype": "fp16" if self.args.scale_bit == 16 else "fp32", } origin_input = op['inputIndexes'] origin_output = op['outputIndexes'] # build new tensor pre_reshape_name = f'{name}/pre_reshape' pre_convert_name = f'{name}/pre_convert' conv_name = name post_convert_name = f'{name}/post_convert' post_reshape_name = f'{name}/post_reshape' pre_reshape_output = self.build_tensor(graph, pre_reshape_name) pre_convert_output = self.build_tensor(graph, pre_convert_name) conv_output = self.build_tensor(graph, conv_name) post_convert_output = self.build_tensor(graph, post_convert_name) # [batch, seq, hidden_size_i] -[Linear] -> [batch, seq, hidden_size_o] # [1, seq, hidden_size_i] ->[Reshape]-> [seq, hidden_size_i, 1, 1] # -[Convert]-[Convolution]-[Convert]-> [Reshape] -> [1, seq, hidden_size_o] pre_reshape = { "name": pre_reshape_name, "type": "Reshape", "inputIndexes": origin_input, "outputIndexes": pre_reshape_output, "main_type": "Reshape", "main": { "dims": [-1, ic, 1, 1], "dimType": "NCHW" }, "defaultDimentionFormat": "NHWC" } pre_convert = { "name": pre_convert_name, "inputIndexes": pre_reshape_output, "outputIndexes": pre_convert_output, "type": "ConvertTensor", "main_type": "TensorConvertInfo", "main": { "source": "NCHW", "dest": "NC4HW4" }, "defaultDimentionFormat": "NHWC" } if quant_bit == 16: quanParameter = { "type": 3 } else: if self.args.sym: aMin = 0 readType = 0 else: aMin = q_min readType = oc * (ic // block_size) quanParameter = { "quantScale": 1.0, "scaleIn": 0.0, "scaleOut": 0.0, "useInt32": False, "has_scaleInt": False, "shapeInt32": shape_int32, "type": 1, "aMaxOrBits": quant_bit, "aMin": aMin, "readType": readType, "weightSize": 0, "scaleStorage": "FP16" if self.args.scale_bit == 16 else "FP32", } conv_op = { "name": conv_name, "inputIndexes": pre_convert_output, "outputIndexes": conv_output, "type": "Convolution", "main_type": "Convolution2D", "main": { 'common': { 'dilateX': 1, 'dilateY': 1, 'strideX': 1, 'strideY': 1, 'kernelX': 1, 'kernelY': 1, 'padX': 0, 'padY': 0, 'group': 1, 'outputCount': oc, 'relu': False, 'padMode': 'CAFFE', 'relu6': False, 'inputCount': ic, 'hasOutputShape': False }, "quanParameter": quanParameter, "external": external }, "defaultDimentionFormat": "NHWC" } post_convert = { "name": post_convert_name, "inputIndexes": conv_output, "outputIndexes": post_convert_output, "type": "ConvertTensor", "main_type": "TensorConvertInfo", "main": { "source": "NC4HW4", "dest": "NCHW" }, "defaultDimentionFormat": "NHWC" } post_reshape = { "name": post_reshape_name, "type": "Reshape", "inputIndexes": post_convert_output, "outputIndexes": origin_output, "main_type": "Reshape", "main": { "dims": [1, -1, oc], "dimType": "NCHW" }, "defaultDimentionFormat": "NHWC" } if name.startswith('/expert/'): post_reshape['main']['dims'] = [-1, oc] return [pre_reshape, pre_convert, conv_op, post_convert, post_reshape]