Files
2024-07-25 16:49:18 +00:00

9.9 KiB

用我们搭建的模型尝试读取官方权重并预测

In [1]:
import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from configuration_minicpm import MiniCPMConfig
from MiniCPM import MiniCPMForCausalLM
import logging
import gc

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
/home/jeeves/.local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

加载模型 config

In [2]:
config_json = json.load(open("/data/workspace/llms-from-scratch-cn/Model_Architecture_Discussions/MiniCPM/config.json"))
config = MiniCPMConfig(**config_json)

按照 config 初始化模型,并查看模型结构

In [3]:

try:
    logging.info("初始化模型")
    model = MiniCPMForCausalLM(config=config).to('cuda')
    logging.info("模型:\n: %s", model)
except Exception as e:
    logging.error(f"初始化模型时发生错误: {e}")
    raise
2024-07-25 15:57:28,490 - INFO - 初始化模型
2024-07-25 15:57:50,064 - INFO - 模型:
: MiniCPMForCausalLM(
  (model): MiniCPMModel(
    (embed_tokens): Embedding(122753, 2304)
    (layers): ModuleList(
      (0-39): 40 x MiniCPMDecoderLayer(
        (self_attn): MiniCPMAttention(
          (q_proj): Linear(in_features=2304, out_features=2304, bias=False)
          (k_proj): Linear(in_features=2304, out_features=2304, bias=False)
          (v_proj): Linear(in_features=2304, out_features=2304, bias=False)
          (o_proj): Linear(in_features=2304, out_features=2304, bias=False)
          (rotary_emb): MiniCPMRotaryEmbedding()
        )
        (mlp): MiniCPMMLP(
          (gate_proj): Linear(in_features=2304, out_features=5760, bias=False)
          (up_proj): Linear(in_features=2304, out_features=5760, bias=False)
          (down_proj): Linear(in_features=5760, out_features=2304, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): MiniCPMRMSNorm()
        (post_attention_layernorm): MiniCPMRMSNorm()
      )
    )
    (norm): MiniCPMRMSNorm()
  )
  (lm_head): Linear(in_features=2304, out_features=122753, bias=False)
)

读取模型权重

In [4]:

path = "/data/model/OpenBMB/MiniCPM-2B-dpo-bf16"

try:
    logging.info("加载模型权重")
    params = torch.load(
        f=path + "/pytorch_model.bin",
        map_location=torch.device('cuda'),
        weights_only=True,  # 设置为True表示仅加载模型的权重。这通常用于加载预训练权重进行微调或预测,而不需要完整的模型结构
        mmap=True  # 使用内存映射方式加载模型文件,这可以提高加载大型模型文件的效率,特别是在有限的内存资源下
    )
    # 打印出模型参数和params中不一致的参数名
    missing_keys, unexpected_keys = model.load_state_dict(params, strict=False)
    # 打印缺失的参数名
    if missing_keys:
        print("缺失的参数名:", missing_keys)

    # 打印多余的参数名
    if unexpected_keys:
        print("多余的参数名:", unexpected_keys)
    # modelV1 = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16, device_map='cuda', trust_remote_code=True)
    # 手动实现 tie embedding 即输入输出共享一个 Embedding
    model.get_output_embeddings().weight = model.get_input_embeddings().weight
    del params
    gc.collect()
    logging.info("加载模型权重完成。")
except Exception as e:
    logging.error(f"加载模型权重时发生错误: {e}")
    raise
2024-07-25 15:57:50,086 - INFO - 加载模型权重
2024-07-25 15:57:52,515 - INFO - 加载模型权重完成。
缺失的参数名: ['lm_head.weight']

MiniCPM 采用了 tie-Embedding 的方式,即词嵌入层和输出层共享参数。这种方式可以减少模型的参数量,提高模型的训练效率。所以需要有获取和设置输入输出词嵌入层的方法。 我们可以看到在加载权重时缺失 lm_head.weight 的参数,这里我们通过手动设置 model.get_output_embeddings().weight = model.get_input_embeddings().weight 来共享参数。

使用默认的 tokenizer 分词

In [11]:
logging.info("初始化分词器")
tokenizer = AutoTokenizer.from_pretrained("/data/model/OpenBMB/MiniCPM-2B-dpo-bf16/")

logging.info("生成文本")
input_texts = ["北京最高的山是哪座山?", "山东省最长的山是哪座山?" ]

tokenizer.pad_token_id=tokenizer.eos_token_id

inputs = tokenizer(input_texts, padding=True, return_tensors="pt").to('cuda')
2024-07-25 16:01:12,360 - INFO - 初始化分词器
2024-07-25 16:01:12,557 - INFO - 生成文本

可以看出 MiniCPM 采用 tokenizer 为 LlamaTokenizerFast, 词表大小为 122753 个 token。

In [12]:
print(tokenizer)
LlamaTokenizerFast(name_or_path='/data/model/OpenBMB/MiniCPM-2B-dpo-bf16/', vocab_size=122753, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '</s>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={
	0: AddedToken("<unk>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	1: AddedToken("<s>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	2: AddedToken("</s>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
}

我们让模型输出结果看看

In [14]:
generate_input = {
    "input_ids": inputs.input_ids,
    "attention_mask": inputs.attention_mask,
    "max_new_tokens": 32,
    "temperature": 1,
    "tokenizer": tokenizer,
}
model.eval()
outputs = model.generate(**generate_input)
for output in outputs:
    result = tokenizer.decode(output, skip_special_tokens=True)
    logging.info(f"生成结果: {result}")
2024-07-25 16:01:36,687 - INFO - 生成结果: 北京最高的山是哪座山?
 北京最高的山是香山。香山位于北京市海淀区,距离北京市中心约25公里,海拔572米。香山是北京市内最高峰
2024-07-25 16:01:36,687 - INFO - 生成结果: 山东省最长的山是哪座山?
 目前,山东省最长的山是泰山。泰山,位于山东省中部,是五岳之一,也是中国著名的山脉之一。泰山是中国著名的山脉之一

可以看出输出的结果还可以