文章目录
- 1. 概要
- 2. Llama-7B权重文件提取
- 3. Llama-8B权重文件提取
- 4. 主要代码功能解析
1. 概要
Llama 系列模型（Meta 发布的大语言模型）在开源社区广受欢迎，不同版本（前文已经介绍过7B和8B的区别，详情请点击链接）在应用场景和硬件需求上各有不同，其权重文件的提取方式也略有差异。本文将通过代码讲解如何获取和提取 Llama 7B 和 8B 的权重参数文件。
2. Llama-7B权重文件提取
from transformers import AutoTokenizer, AutoModelForCausalLMdef save_weight_int(int_weight: torch.Tensor, path):if path[-4:] != '.bin':raise ValueError('Path must end with .bin')int_weight.cpu().detach().numpy().astype(np.int32).tofile(path)if __name__ == '__main__':tokenizer = AutoTokenizer.from_pretrained(model_card, local_files_only = True, cache_dir = "./model-storage")model = AutoModelForSequenceClassification.from_pretrained(model_card, local_files_only = True, cache_dir = "./model-storage")for (i, w) in model.model.layers[0].named_parameters():if len(w.shape) == 2:pp_size = w.shape[0]pp_size <<= args.log_off_factor # 位移操作elif len(w.shape) == 1:(pp_size,) = w.shapeelse:raise ValueError(f"Unexpected shape {w.shape} for parameter {i}")print(f"Layer {i}, Parameter {j}, Shape: {w_orig.shape}")save_weight_int(w_orig, f"./zkllm-workdir/Llama-2-{args.model_size}b/layer-{i}-{j}-int.bin")3. Llama-8B权重文件提取
from transformers import AutoTokenizer, AutoModelForCausalLMdef save_weight_int(int_weight: torch.Tensor, path):if path[-4:] != '.bin':raise ValueError('Path must end with .bin')int_weight.cpu().detach().numpy().astype(np.int32).tofile(path)if __name__ == '__main__':for i, layer in enumerate(model.model.layers):for j, w in layer.named_parameters():# 中间层参数的处理if len(w.shape) == 2:w_orig = w.float().Telse:w_orig = w.float()print(f"Layer {i}, Parameter {j}, Shape: {w_orig.shape}")save_weight_int(w_orig, f"./zkllm-workdir/Llama-2-{args.model_size}b/layer-{i}-{j}-int.bin")# 处理顶层参数（如输出层的 score.weight）for name, param in model.named_parameters():if "score.weight" in name: # 仅处理输出权重if len(param.shape) == 2:w_orig = param.float().Telse:w_orig = param.float()print(f"Processing Output Layer Parameter {name}, Shape: {w_orig.shape}")save_weight_int(w_orig, f"./zkllm-workdir/Llama-2-{args.model_size}b/{name.replace('.', '-')}-int.bin")
4. 主要代码功能解析
-
save_weight_int(int_weight: torch.Tensor, path) 函数
作用：将权重量化为 int32 数据，并以 .bin 格式保存到指定路径。 -
遍历 model.model.layers 的所有参数
for i, layer in enumerate(model.model.layers):for j, w in layer.named_parameters():
- 遍历模型的每一层（model.model.layers），i是层索引，layer 是每一层的模块。
- 使用 named_parameters() 遍历每层中的所有参数（权重和偏置）。
- j 是参数名称（例如 self_attn.q_proj.weight）。
- w 是参数张量
- 中间参数处理（可以去掉）
if len(param.shape) == 2:w_orig = param.float().T
else:w_orig = param.float()
