import os import shutil from argparse import ArgumentParser from glob import glob from tqdm import tqdm, trange import torch import ctypes from safetensors.torch import safe_open, save_file from kernel import weight_dequant mapping = { "embed_tokens": ("embed", 0), "input_layernorm": ("attn_norm", None), "post_attention_layernorm": ("ffn_norm", None), "q_proj": ("wq", 0), "q_a_proj": ("wq_a", None), "q_a_layernorm": ("q_norm", None), "q_b_proj": ("wq_b", 0), "kv_a_proj_with_mqa": ("wkv_a", None), "kv_a_layernorm": ("kv_norm", None), "kv_b_proj": ("wkv_b", 0), "o_proj": ("wo", 1), "gate": ("gate", None), "gate_proj": ("w1", 0), "down_proj": ("w2", 1), "up_proj": ("w3", 0), "norm": ("norm", None), "lm_head": ("head", 0), "scale": ("scale", None), } EmbedsInOneFile = 256 EmbedsZKDir = "../zkdata/embeds/" wkv_b_1_rescales = [32, 34, 37, 36, 33, 32, 33, 33, 30, 32, 32, 30, 31, 30, 29, 30, 29, 30, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 29, 29, 30, 30, 30, 30, 29, 30, 30, 29, 30] wkv_b_2_rescales = [31, 32, 32, 31, 32, 30, 30, 30, 30, 30, 30, 30, 30, 29, 29, 29, 29, 30, 29, 29, 29, 29, 29, 29, 30, 30, 30, 29, 29, 29, 29, 29, 30, 29, 30, 29, 30, 29, 29, 29, 30, 29, 29, 29, 29, 30, 29, 30, 30, 30, 29, 29, 29, 30, 30, 29, 29, 29, 30, 30, 30] wo_rescales = [31, 32, 32, 32, 32, 31, 32, 31, 31, 31, 31, 31, 31, 31, 30, 31, 31, 32, 31, 31, 31, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 31, 30, 31, 30, 30, 31, 31, 31, 30, 31, 31, 31, 30, 31, 31, 31, 31, 32 ] gate_rescales = [0, 0, 0, 33, 32, 32, 32, 31, 32, 31, 30, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 31, 32, 31, 32, 32, 32, 32, 31, 32, 32, 31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 32, 32 ] w1_rescales = [32, 32, 32] w2_rescales = [31, 32, 31] w3_rescales = [32, 33, 32] shared_w1_rescales = [0, 0, 0, 30, 30, 29, 29, 29, 28, 29, 29, 28, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 29, 29, 30, 29, 29, 30, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29] shared_w2_rescales = [0, 0, 0, 30, 30, 30, 30, 30, 29, 29, 30, 29, 29, 29, 30, 30, 30, 30, 30, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 29, 29, 29, 29, 29, 29, 29, 29, 30, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30] shared_w3_rescales = [0, 0, 0, 30, 30, 30, 30, 30, 29, 29, 30, 29, 29, 29, 30, 30, 30, 29, 30, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 29, 29, 29, 29, 29, 29, 29, 30, 30, 29, 30, 29, 29, 29, 29, 30, 29, 29, 30, 30, 29, 30, 30, 30, 29, 29, 30, 30, 30, 29, 28] layer_state_dict0 = [{} for _ in range(61)] layer_state_dict = [{} for _ in range(61)] experts = [ [{} for _j in range(256)] for _i in range(61)] def getF32PrintStr(ele): v = int(ele.cpu().view(torch.uint32).item()) ex = str((v >> 23 & 0xFF) - 127) r = '(1+' + str(v & 0x7FFFFF) + '/8388608)' if v & 0x80000000: vstr = '-' + r + '*2^' + ex else: vstr = r + '*2^' + ex return vstr def getBF16PrintStr(ele): v = int(ele.cpu().view(torch.uint16).item()) ex = v >> 7 & 0xFF r = '(1+' + str(v & 0x7F) + '/128)' rraw = v & 0x7F if v & 0x8000: vstr = '-' + r + '*2^' + str(ex - 127) else: vstr = r + '*2^' + str(ex - 127) return vstr def getBF8PrintStr(ele): v = int(ele.cpu().view(torch.uint8).item()) ex = str((v >> 3 & 0xF) - 7) r = '(1+' + str(v & 0x7) + '/8)' if v & 0x80: vstr = '-' + r + '*2^' + ex else: vstr = r + '*2^' + ex if ex == -7 or ex == 8: print(vstr) return vstr def mem(i): a = torch.cuda.memory_allocated()/1024**2 r = torch.cuda.memory_reserved()/1024**2 m = torch.cuda.max_memory_allocated()/1024**2 print(f"{i} allocated={a:.1f}MB, reserved={r:.1f}MB, max={m:.1f}MB", flush=True) def handle_expert_w(layer_id, expert_id, idx, param_weight, weight_name, scale, typ, shape, experts_save_path): global layer_state_dict0 global experts scale_name = weight_name.replace('weight', 'scale') param_scale = layer_state_dict0[layer_id][scale_name] weight = weight_dequant(param_weight.cuda(), param_scale.cuda()) # scale = experts_w3_rescales[layer_id][expert_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) # layer_state_dict[layer_id][weight_name] = param_int.cpu() # layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) weight_name2 = f'w{idx}.weight' scale_name2 = f'w{idx}.scale' experts[layer_id][expert_id][weight_name2] = param_int experts[layer_id][expert_id][scale_name2] = torch.tensor(scale, dtype=torch.int32) if len(experts[layer_id][expert_id]) == 6: # w1, w2, w3 以及对应的 scale save_file(experts[layer_id][expert_id], os.path.join(experts_save_path, f"{expert_id}.safetensors")) experts[layer_id][expert_id] = {} print(f'layer {layer_id} expert {expert_id} w{idx} type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {scale_name}') def saveTensor(fileName, t): with open(fileName, "w", encoding="utf-8") as f: t = t.detach() if t.device.type != "cpu": t = t.cpu() t = t.contiguous() with open(fileName, "wb") as f: f.write(t.numpy().tobytes(order="C")) def main(hf_ckpt_path, save_path, n_experts, mp): """ Converts and saves model checkpoint files into a specified format. Args: hf_ckpt_path (str): Path to the directory containing the input checkpoint files. save_path (str): Path to the directory where the converted checkpoint files will be saved. n_experts (int): Total number of experts in the model. mp (int): Model parallelism factor. Returns: None """ torch.cuda.set_device(0) # 设置pytorch计算时的默认数据类型。这里使用的是BF16 torch.set_default_dtype(torch.bfloat16) # 限制 PyTorch 在 CPU 计算时最多使用 8 个线程，防止过多线程竞争资源： torch.set_num_threads(8) # 设定随机种子，保证不同进程初始化时随机数相同。 torch.manual_seed(965) # n_local_experts = n_experts // mp # state_dicts = [{} for _ in range(mp)] head_state_dict = {} norm_state_dict = {} embed_state_dict = {} experts_w1_rescales = [] experts_w2_rescales = [] experts_w3_rescales = [] with open("w1.txt", "r", encoding="utf-8") as f1: for line in f1: layer_line = line.strip().split() int_list = [int(s) for s in layer_line] experts_w1_rescales.append(int_list) with open("w2.txt", "r", encoding="utf-8") as f2: for line in f2: layer_line = line.strip().split() int_list = [int(s) for s in layer_line] experts_w2_rescales.append(int_list) with open("w3.txt", "r", encoding="utf-8") as f3: for line in f3: layer_line = line.strip().split() int_list = [int(s) for s in layer_line] experts_w3_rescales.append(int_list) # Tqdm 是一个快速，可扩展的Python进度条，可以在 Python 长循环中添加一个进度提示信息，用户只需要封装任意的迭代器 tqdm(iterator)。 # glob是python自己带的一个文件操作相关模块，用它可以查找符合自己目的的文件，类似于Windows下的文件搜索 for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))): with safe_open(file_path, framework="pt", device="cpu") as f: print('Opening ' + file_path, flush=True) for name in f.keys(): # print('name 1: ', name) if "model.layers.61" in name: continue param: torch.Tensor = f.get_tensor(name) if name.startswith("model."): name = name[len("model."):] name = name.replace("self_attn", "attn") name = name.replace("mlp", "ffn") name = name.replace("weight_scale_inv", "scale") name = name.replace("e_score_correction_bias", "bias") key = name.split(".")[-2] assert key in mapping, f"Key {key} not found in mapping" # print('key::: ' + key) new_key, dim = mapping[key] # print('dim::: ' + str(dim)) name = name.replace(key, new_key) ns = name.split(".") comp = ns[0] if comp == 'head': name2 = name[len('head.'):] print('head: ' + name2) param_int = (param.to(torch.float32) * (2 ** 43)).round().to(torch.int64) head_state_dict[name2] = param_int elif comp == 'norm': name2 = name[len('norm.'):] print('norm: ' + name2) param_int = (param.to(torch.float32) * (2 ** 15)).round().to(torch.int64) norm_state_dict[name2] = param_int elif comp == 'embed': name2 = name[len('embed.'):] print('embed: ' + name2) param_int = (param.to(torch.float32) * (2 ** 31)).round().to(torch.int64) embed_state_dict[name2] = param_int os.makedirs(EmbedsZKDir, exist_ok=True) fileCount = param_int.shape[0] // EmbedsInOneFile for i in range(0, fileCount): saveTensor(EmbedsZKDir + str(i) + '.bin', param_int[i * EmbedsInOneFile : (i+1) * EmbedsInOneFile].cpu()) elif comp == 'layers': layer_id = int(ns[1]) name2 = '.'.join(ns[2:]) layer_state_dict0[layer_id][name2] = param print('Finish loading state dict from disk! ++++++++++') # for layer_id, states in enumerate(layer_state_dict0): for layer_id in range(len(layer_state_dict0)): os.makedirs(f'{save_path}/experts-{layer_id}', exist_ok=True) states = layer_state_dict0[layer_id] for name, param in states.items(): ns = name.split(".") typ = param.type() shape = param.shape if ns[0] == 'attn_norm': print(f'layer {layer_id} {name}, type: {typ}', flush=True) if ns[1] == 'weight': param_int = (param.to(torch.float32) * (2 ** 21)).round().to(torch.int32) layer_state_dict[layer_id][name] = param_int elif ns[0] == 'ffn_norm': print(f'layer {layer_id} {name}, type: {typ}', flush=True) if ns[1] == 'weight': param_int2 = (param.to(torch.float32) * (2 ** 23)).round().to(torch.int32) layer_state_dict[layer_id][name] = param_int2 elif ns[0] == 'ffn': if len(ns) == 3: if ns[1] == 'w1' and ns[2] == 'scale': continue elif ns[1] == 'w1' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = w1_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} w1 weight, type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {name}', flush=True) elif ns[1] == 'w2' and ns[2] == 'scale': continue elif ns[1] == 'w2' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = w2_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} w2 weight, type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {name}', flush=True) elif ns[1] == 'w3' and ns[2] == 'scale': continue elif ns[1] == 'w3' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = w3_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} w3 weight, type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {name}', flush=True) elif ns[1] == 'gate' and ns[2] == 'weight': gate_rescale = 2 ** gate_rescales[layer_id] gate_int = (param.to(torch.float32) * gate_rescale).round().to(torch.int32) layer_state_dict[layer_id][name] = gate_int.cpu() rescale_name = name.replace('weight', 'scale') layer_state_dict[layer_id][rescale_name] = torch.tensor(gate_rescales[layer_id], dtype=torch.int32) print(f'layer {layer_id}: gate_weight_name: {name}, gate_scale_name: {rescale_name}') elif ns[1] == 'gate' and ns[2] == 'bias': bias_int = (param.to(torch.float32) * (2 ** 23)).round().to(torch.int32) layer_state_dict[layer_id][name] = bias_int.cpu() print(f'layer {layer_id} bias: {name}') else: layer_state_dict[layer_id][name] = param elif len(ns) == 4: if ns[1] == 'shared_experts': if (ns[2] == 'w1' or ns[2] == 'w2' or ns[2] == 'w3') and ns[3] == 'scale': continue elif ns[2] == 'w1' and ns[3] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = shared_w1_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} shared_expert w1 type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {scale_name}') elif ns[2] == 'w2' and ns[3] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = shared_w2_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} shared_expert w2 type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {scale_name}') elif ns[2] == 'w3' and ns[3] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = shared_w3_rescales[layer_id] rescale = 2 ** scale param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} shared_expert w3 type: {typ}, shape: {shape}, weight_name: {weight_name}, scale_name: {scale_name}') else: layer_state_dict[layer_id][name] = param else: layer_state_dict[layer_id][name] = param elif len(ns) == 5: if ns[1] == 'experts': expert_id = int(ns[2]) if (ns[3] == 'w1' or ns[3] == 'w2' or ns[3] == 'w3') and ns[4] == 'scale': continue elif ns[3] == 'w1' and ns[4] == 'weight': scale = experts_w1_rescales[layer_id][expert_id] handle_expert_w(layer_id, expert_id, 1, param, name, scale, typ, shape, f'{save_path}/experts-{layer_id}') elif ns[3] == 'w2' and ns[4] == 'weight': scale = experts_w2_rescales[layer_id][expert_id] handle_expert_w(layer_id, expert_id, 2, param, name, scale, typ, shape, f'{save_path}/experts-{layer_id}') elif ns[3] == 'w3' and ns[4] == 'weight': scale = experts_w3_rescales[layer_id][expert_id] handle_expert_w(layer_id, expert_id, 3, param, name, scale, typ, shape, f'{save_path}/experts-{layer_id}') else: layer_state_dict[layer_id][name] = param else: layer_state_dict[layer_id][name] = param elif ns[0] == 'attn': if len(ns) == 3: if ns[1] == 'wq_a' and ns[2] == 'scale': continue elif ns[1] == 'wq_a' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) weight_int = (weight.to(torch.float32) * (2 ** 30)).round().to(torch.int32) layer_state_dict[layer_id][weight_name] = weight_int.cpu() print(f'layer {layer_id} wq_a weight, type: {typ}, shape: {shape}', flush=True) elif ns[1] == 'q_norm': print(f'layer {layer_id} q_norm, type: {typ}, shape: {shape}', flush=True) param_int3 = (param.to(torch.float32) * (2 ** 19)).round().to(torch.int32) layer_state_dict[layer_id][name] = param_int3 elif ns[1] == 'kv_norm': print(f'layer {layer_id} kv_norm, type: {typ}, shape: {shape}', flush=True) param_int4 = (param.to(torch.float32) * (2 ** 23)).round().to(torch.int32) layer_state_dict[layer_id][name] = param_int4 elif ns[1] == 'wq_b' and ns[2] == 'scale': continue elif ns[1] == 'wq_b' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) weight_int = (weight.to(torch.float32) * (2 ** 30)).round().to(torch.int32) weight_int = weight_int.view(128, 192, 1536) wq_b1, wq_b2 = torch.split(weight_int, [128, 64], dim=-2) print(f'layer {layer_id} wq_b1 weight, shape: {wq_b1.shape}, wq_b2 weight, shape: {wq_b2.shape}', flush=True) wq_b1 = wq_b1.reshape(128 * 128, 1536) wq_b2 = wq_b2.reshape(128 * 64, 1536) wq_b1_name = weight_name.replace('wq_b', 'wq_b1') wq_b2_name = weight_name.replace('wq_b', 'wq_b2') # layer_state_dict[layer_id][weight_name] = weight_int.cpu() layer_state_dict[layer_id][wq_b1_name] = wq_b1.cpu() layer_state_dict[layer_id][wq_b2_name] = wq_b2.cpu() print(f'layer {layer_id} wq_b weight, type: {typ}, shape: {shape}', flush=True) elif ns[1] == 'wkv_a' and ns[2] == 'scale': continue elif ns[1] == 'wkv_a' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) weight_int = (weight.to(torch.float32) * (2 ** 29)).round().to(torch.int32) # layer_state_dict[layer_id][weight_name] = weight_int.cpu() weight_int = weight_int.view(576, 7168) wkv_a1, wkv_a2 = torch.split(weight_int, [512, 64], dim=-2) print(f'layer {layer_id} wkv_a1 weight, shape: {wkv_a1.shape}, wkv_a2 weight, shape: {wkv_a2.shape}', flush=True) wkv_a1_name = weight_name.replace('wkv_a', 'wkv_a1') wkv_a2_name = weight_name.replace('wkv_a', 'wkv_a2') # layer_state_dict[layer_id][weight_name] = weight_int.cpu() layer_state_dict[layer_id][wkv_a1_name] = wkv_a1.cpu() layer_state_dict[layer_id][wkv_a2_name] = wkv_a2.cpu() print(f'layer {layer_id} wkv_a weight, type: {typ}, shape: {shape}', flush=True) elif ns[1] == 'wkv_b' and ns[2] == 'scale': continue elif ns[1] == 'wkv_b' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) wkv_b = weight.view(128, 256, 512) wkv_b_1 = wkv_b[:, :128] wkv_b_1 = wkv_b_1.reshape(128 * 128, 512) scale1 = wkv_b_1_rescales[layer_id] wkv_b_1_rescale = 2 ** scale1 wkv_b_1_int = torch.round(wkv_b_1.to(torch.float32) * wkv_b_1_rescale).to(torch.int32) wkv_b_2 = wkv_b[:, -128:] wkv_b_2 = wkv_b_2.reshape(128 * 128, 512) scale2 = wkv_b_2_rescales[layer_id] wkv_b_2_rescale = 2 ** scale2 wkv_b_2_int = torch.round(wkv_b_2.to(torch.float32) * wkv_b_2_rescale).to(torch.int32) wkv_b_1_name = weight_name.replace("wkv_b", "wkv_b_1") wkv_b_1_scale_name = scale_name.replace("wkv_b", "wkv_b_1") layer_state_dict[layer_id][wkv_b_1_name] = wkv_b_1_int.cpu() layer_state_dict[layer_id][wkv_b_1_scale_name] = torch.tensor(scale1, dtype=torch.int32) wkv_b_2_name = weight_name.replace("wkv_b", "wkv_b_2") wkv_b_2_scale_name = scale_name.replace("wkv_b", "wkv_b_2") layer_state_dict[layer_id][wkv_b_2_name] = wkv_b_2_int.cpu() layer_state_dict[layer_id][wkv_b_2_scale_name] = torch.tensor(scale2, dtype=torch.int32) print(f'layer {layer_id} wkv_b, type: {typ}, shape: {shape}, wkv_b_1 weight: {wkv_b_1_name}, wkv_b_1 scale: {wkv_b_1_scale_name}, wkv_b_2 weight: {wkv_b_2_name}, wkv_b_2 scale: {wkv_b_2_scale_name}', flush=True) elif ns[1] == 'wo' and ns[2] == 'scale': continue elif ns[1] == 'wo' and ns[2] == 'weight': param_weight = param.cuda() weight_name = name scale_name = name.replace('weight', 'scale') param_scale = states[scale_name] weight = weight_dequant(param_weight, param_scale.cuda()) scale = wo_rescales[layer_id] rescale = 2 ** scale if layer_id != 58: param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) else: wo_abs = weight.abs().cpu() maxpos = wo_abs.argmax() row, col = divmod(maxpos.item(), weight.size(1)) print(f'maxpos: {maxpos}, {row} {col}', flush=True) vstr = getBF16PrintStr(weight[row][col]) print(f'weight[{row}][{col}]: {vstr}', flush=True) weight[row][col] = 0 param_int = (weight.to(torch.float32) * rescale).round().to(torch.int32) param_int[row][col] = -(2 ** 31) layer_state_dict[layer_id][weight_name] = param_int.cpu() layer_state_dict[layer_id][scale_name] = torch.tensor(scale, dtype=torch.int32) print(f'layer {layer_id} wo weight, type: {typ}, shape: {shape}, weight: {weight_name}, scale: {scale_name}', flush=True) else: layer_state_dict[layer_id][name] = param else: layer_state_dict[layer_id][name] = param else: layer_state_dict[layer_id][name] = param save_file(layer_state_dict[layer_id], os.path.join(save_path, f"layer-{layer_id}.safetensors")) print(f'Finish saving layer {layer_id}', flush=True) layer_state_dict0[layer_id] = {} layer_state_dict[layer_id] = {} print('Finish opening') os.makedirs(save_path, exist_ok=True) print(layer_state_dict) print(experts) save_file(head_state_dict, os.path.join(save_path, f"head_int.safetensors")) save_file(norm_state_dict, os.path.join(save_path, f"norm_int.safetensors")) save_file(embed_state_dict, os.path.join(save_path, f"embed_int.safetensors")) # for i, st in enumerate(layer_state_dict): # # print(f'{i} {st['attn_norm.weight']}', flush=True) # # print(f'{i} {st['ffn_norm.weight']}', flush=True) # save_file(st, os.path.join(save_path, f"layer-{i}.safetensors")) # print(f'Finish saving layer {i}', flush=True) # for i in trange(mp): # save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors")) # print('Finish saving files') for file_path in glob(os.path.join(hf_ckpt_path, "*token*")): new_file_path = os.path.join(save_path, os.path.basename(file_path)) shutil.copyfile(file_path, new_file_path) if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("--hf-ckpt-path", type=str, required=True) parser.add_argument("--save-path", type=str, required=True) parser.add_argument("--n-experts", type=int, required=True) parser.add_argument("--model-parallel", type=int, required=True) args = parser.parse_args() assert args.n_experts % args.model_parallel == 0, "Number of experts must be divisible by model parallelism" main(args.hf_ckpt_path, args.save_path, args.n_experts, args.model_parallel)