diff --git a/Model_Architecture_Discussions/ChatGLM4/chatglm4-guide.ipynb b/Model_Architecture_Discussions/ChatGLM4/chatglm4-guide.ipynb new file mode 100644 index 0000000..0b1e5e0 --- /dev/null +++ b/Model_Architecture_Discussions/ChatGLM4/chatglm4-guide.ipynb @@ -0,0 +1,3735 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "a5d6ad19-adff-423b-8177-30a0d4f6ceed", + "metadata": {}, + "outputs": [], + "source": [ + "# import accelerate" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "02e86302-68c0-455a-9457-6a4618b95cba", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "env: HF_ENDPOINT=https://hf-mirror.com\n" + ] + } + ], + "source": [ + "%env HF_ENDPOINT=https://hf-mirror.com\n", + "import os\n", + "os.environ['HF_HOME'] = '/data1/ckw'\n", + "os.environ['HF_ENDPOINT']='https://hf-mirror.com'" + ] + }, + { + "cell_type": "markdown", + "id": "a45076d6-d882-4450-8596-424d364ba65e", + "metadata": {}, + "source": [ + "首先，我们构造Tokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "4f4f30f9-d442-4079-9f0b-04874e9de217", + "metadata": {}, + "outputs": [], + "source": [ + "import regex as re\n", + "import base64\n", + "import os\n", + "import json\n", + "import tiktoken\n", + "from torch import TensorType\n", + "from typing import List, Optional, Union, Dict, Any\n", + "from transformers import PreTrainedTokenizer\n", + "from transformers.utils import logging, PaddingStrategy\n", + "from transformers.tokenization_utils_base import EncodedInput, BatchEncoding" + ] + }, + { + "cell_type": "markdown", + "id": "ff7537fb-dcc2-47fc-a4da-a33e8a68f1ce", + "metadata": {}, + "source": [ + "`ChatGLM4Tokenizer` 类是一个自定义的 `PreTrainedTokenizer` 类，用于处理特定的 GLM-4 模型的分词需求。这个类主要有以下结构和成员函数：\n", + "\n", + "### 类结构\n", + "- **属性**：\n", + " - `vocab_files_names`：包含词汇文件名的字典。\n", + " - `model_input_names`：包含模型输入名的列表。\n", + " - `vocab_file`：词汇文件路径。\n", + " - `name`：分词器的名称。\n", + " - `pat_str`：正则表达式字符串，用于分词。\n", + " - `encode_special_tokens`：是否编码特殊字符。\n", + " - `mergeable_ranks`：可合并的词汇排名。\n", + " - `tokenizer`：基于 `tiktoken` 库的编码器。\n", + " - `decoder`：解码器，映射词汇排名到词汇。\n", + " - `n_words`：词汇表大小。\n", + "\n", + "### 成员函数\n", + "\n", + "- **初始化函数**：\n", + " ```python\n", + " def __init__(self, vocab_file, padding_side=\"left\", clean_up_tokenization_spaces=False, encode_special_tokens=False, **kwargs)\n", + " ```\n", + " 初始化 `ChatGLM4Tokenizer` 类，加载词汇文件，设置正则表达式和分词器。\n", + "\n", + "- **词汇表大小属性**：\n", + " ```python\n", + " @property\n", + " def vocab_size(self)\n", + " ```\n", + " 返回词汇表大小。\n", + "\n", + "- **获取词汇表**：\n", + " ```python\n", + " def get_vocab(self)\n", + " ```\n", + " 返回词汇表字典。\n", + "\n", + "- **将 tokens 转换为字符串**：\n", + " ```python\n", + " def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str\n", + " ```\n", + " 将 tokens 序列转换为字符串。\n", + "\n", + "- **分词**：\n", + " ```python\n", + " def _tokenize(self, text, **kwargs)\n", + " ```\n", + " 将文本分词为 token 列表。\n", + "\n", + "- **将 token 转换为 ID**：\n", + " ```python\n", + " def _convert_token_to_id(self, token)\n", + " ```\n", + " 将 token（字符串）转换为 ID。\n", + "\n", + "- **将 ID 转换为 token**：\n", + " ```python\n", + " def _convert_id_to_token(self, index)\n", + " ```\n", + " 将 ID（整数）转换为 token（字符串）。\n", + "\n", + "- **保存词汇表**：\n", + " ```python\n", + " def save_vocabulary(self, save_directory, filename_prefix=None)\n", + " ```\n", + " 保存词汇表到指定目录。\n", + "\n", + "- **获取前缀 tokens**：\n", + " ```python\n", + " def get_prefix_tokens(self)\n", + " ```\n", + " 返回前缀 tokens 列表。\n", + "\n", + "- **构建单条消息**：\n", + " ```python\n", + " def build_single_message(self, role, metadata, message, tokenize=True)\n", + " ```\n", + " 构建单条消息，包含角色、元数据和消息内容。\n", + "\n", + "- **应用聊天模板**：\n", + " ```python\n", + " def apply_chat_template(self, conversation, add_generation_prompt=False, tokenize=True, padding=False, truncation=False, max_length=None, return_tensors=None, return_dict=False, tokenizer_kwargs=None, add_special_tokens=True, **kwargs)\n", + " ```\n", + " 将会话数据应用到聊天模板中。\n", + "\n", + "- **构建带有特殊 token 的输入**：\n", + " ```python\n", + " def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None)\n", + " ```\n", + " 构建带有特殊 token 的输入。\n", + "\n", + "- **填充输入**：\n", + " ```python\n", + " def _pad(self, encoded_inputs, max_length=None, padding_strategy=PaddingStrategy.DO_NOT_PAD, pad_to_multiple_of=None, return_attention_mask=None)\n", + " ```\n", + " 填充编码后的输入，确保输入长度一致。\n", + "\n", + "### 类中的主要函数和方法概述\n", + "\n", + "- **初始化**：`__init__` 函数用于初始化分词器，包括加载词汇表、正则表达式的编译和编码器的初始化。\n", + "- **获取词汇表和大小**：`get_vocab` 和 `vocab_size` 属性用于获取词汇表及其大小。\n", + "- **分词及转换**：`_tokenize`、`convert_tokens_to_string`、`_convert_token_to_id` 和 `_convert_id_to_token` 函数用于实现文本的分词及 tokens 和 IDs 之间的转换。\n", + "- **保存和加载**：`save_vocabulary` 函数用于保存词汇表到指定目录。\n", + "- **会话处理**：`get_prefix_tokens`、`build_single_message` 和 `apply_chat_template` 函数用于处理会话数据和模板应用。\n", + "- **输入处理**：`build_inputs_with_special_tokens` 和 `_pad` 函数用于处理模型输入，确保输入格式和长度符合要求。\n", + "\n", + "通过以上结构和成员函数，`ChatGLM4Tokenizer` 类实现了一个完整的分词器功能，能够处理特定的 GLM-4 模型的分词需求。" + ] + }, + { + "cell_type": "code", + "execution_count": 61, + "id": "4a119f91-e04b-4f8b-aa2a-02ca20a88a0a", + "metadata": {}, + "outputs": [], + "source": [ + "class ChatGLM4Tokenizer(PreTrainedTokenizer):\n", + " # 定义词汇文件名和模型输入名\n", + " vocab_files_names = {\"vocab_file\": \"tokenizer.model\"}\n", + " model_input_names = [\"input_ids\", \"attention_mask\", \"position_ids\"]\n", + "\n", + " def __init__(\n", + " self,\n", + " vocab_file,\n", + " padding_side=\"left\",\n", + " clean_up_tokenization_spaces=False,\n", + " encode_special_tokens=False,\n", + " **kwargs\n", + " ):\n", + " # 初始化一些基础属性\n", + " self.name = \"GLM4Tokenizer\"\n", + " self.vocab_file = vocab_file\n", + " \n", + " # 正则表达式模式字符串，用于分词\n", + " pat_str = \"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\\\r\\\\n\\\\p{L}\\\\p{N}]?\\\\p{L}+|\\\\p{N}{1,3}| ?[^\\\\s\\\\p{L}\\\\p{N}]+[\\\\r\\\\n]*|\\\\s*[\\\\r\\\\n]+|\\\\s+(?!\\\\S)|\\\\s+\"\n", + " # 编译正则表达式\n", + " self.pat_str = regex.compile(pat_str)\n", + " # self.pat_str = re.compile(pat_str)\n", + " # 是否编码特殊字符\n", + " self.encode_special_tokens = encode_special_tokens\n", + "\n", + " # 用于存储可合并的词汇排名\n", + " mergeable_ranks = {}\n", + " # 读取词汇文件\n", + " with open(vocab_file) as f:\n", + " for line in f:\n", + " token, rank = line.strip().split() # 读取每一行，获取词汇和其对应的排名\n", + " rank = int(rank) # 将排名转换为整数\n", + " token = base64.b64decode(token) # 解码词汇\n", + " mergeable_ranks[token] = rank # 存储到词汇排名字典中\n", + "\n", + " self.mergeable_ranks = mergeable_ranks\n", + "\n", + " # 初始化编码器，使用 tiktoken 库\n", + " self.tokenizer = tiktoken.Encoding(\n", + " name=\"my_tokenizer\",\n", + " pat_str=pat_str,\n", + " mergeable_ranks=mergeable_ranks,\n", + " special_tokens={}\n", + " )\n", + " # 解码器，映射排名到词汇\n", + " self.decoder = {rank: token for token, rank in mergeable_ranks.items()}\n", + " # 词汇表大小\n", + " self.n_words = len(self.decoder)\n", + "\n", + " # 调用父类的初始化方法\n", + " super().__init__(\n", + " padding_side=padding_side,\n", + " clean_up_tokenization_spaces=clean_up_tokenization_spaces,\n", + " **kwargs\n", + " )\n", + "\n", + " @property\n", + " def vocab_size(self):\n", + " # 返回词汇表大小\n", + " return self.n_words\n", + "\n", + " def get_vocab(self):\n", + " \"\"\" 返回词汇表字典 \"\"\"\n", + " vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}\n", + " vocab.update(self.added_tokens_encoder)\n", + " return vocab\n", + "\n", + " def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:\n", + " \"\"\"\n", + " 将 tokens 序列转换为字符串。\n", + " \"\"\"\n", + " text = \"\"\n", + " temp = b\"\"\n", + " for t in tokens:\n", + " if isinstance(t, str):\n", + " if temp:\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " temp = b\"\"\n", + " text += t\n", + " elif isinstance(t, bytes):\n", + " temp += t\n", + " else:\n", + " raise TypeError(\"token should only be of type bytes or str\")\n", + " if temp:\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " return text\n", + "\n", + " def _tokenize(self, text, **kwargs):\n", + " # 使用正则表达式和编码器对文本进行分词\n", + " tokens = []\n", + " ids = self.tokenizer.encode(text)\n", + " for t in ids:\n", + " tokens.append(self.decoder[t])\n", + " return tokens\n", + "\n", + " def _convert_token_to_id(self, token):\n", + " \"\"\" 将 token (字符串) 转换为 id (整数) \"\"\"\n", + " return self.mergeable_ranks[token]\n", + "\n", + " def _convert_id_to_token(self, index):\n", + " \"\"\" 将 id (整数) 转换为 token (字符串) \"\"\"\n", + " return self.decoder.get(index, \"\")\n", + "\n", + " def save_vocabulary(self, save_directory, filename_prefix=None):\n", + " \"\"\"\n", + " 将词汇表和特殊字符文件保存到指定目录。\n", + "\n", + " Args:\n", + " save_directory (`str`): 保存词汇表的目录。\n", + " filename_prefix (`str`, *optional*): 保存文件名的前缀。\n", + "\n", + " Returns:\n", + " `Tuple(str)`: 保存的文件路径。\n", + " \"\"\"\n", + " if os.path.isdir(save_directory):\n", + " vocab_file = os.path.join(\n", + " save_directory, self.vocab_files_names[\"vocab_file\"]\n", + " )\n", + " else:\n", + " vocab_file = save_directory\n", + "\n", + " with open(self.vocab_file, 'rb') as fin:\n", + " proto_str = fin.read()\n", + "\n", + " with open(vocab_file, \"wb\") as writer:\n", + " writer.write(proto_str)\n", + "\n", + " return (vocab_file,)\n", + "\n", + " def get_prefix_tokens(self):\n", + " # 返回前缀 tokens 列表\n", + " prefix_tokens = [self.convert_tokens_to_ids(\"[gMASK]\"), self.convert_tokens_to_ids(\"\")]\n", + " return prefix_tokens\n", + "\n", + " def build_single_message(self, role, metadata, message, tokenize=True):\n", + " # 构建单条消息，包含角色、元数据和消息内容\n", + " assert role in [\"system\", \"user\", \"assistant\", \"observation\"], role\n", + " if tokenize:\n", + " role_tokens = [self.convert_tokens_to_ids(f\"<|{role}|>\")] + self.tokenizer.encode(f\"{metadata}\\n\",\n", + " disallowed_special=())\n", + " message_tokens = self.tokenizer.encode(message, disallowed_special=())\n", + " tokens = role_tokens + message_tokens\n", + " return tokens\n", + " else:\n", + " return str(f\"<|{role}|>{metadata}\\n{message}\")\n", + "\n", + " def apply_chat_template(\n", + " self,\n", + " conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], \"Conversation\"],\n", + " add_generation_prompt: bool = False,\n", + " tokenize: bool = True,\n", + " padding: bool = False,\n", + " truncation: bool = False,\n", + " max_length: Optional[int] = None,\n", + " return_tensors: Optional[Union[str, TensorType]] = None,\n", + " return_dict: bool = False,\n", + " tokenizer_kwargs: Optional[Dict[str, Any]] = None,\n", + " add_special_tokens: bool = True,\n", + " **kwargs,\n", + " ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:\n", + " \n", + " if return_dict and not tokenize:\n", + " raise ValueError(\n", + " \"`return_dict=True` is incompatible with `tokenize=False`, because there is no dict \"\n", + " \"of tokenizer outputs to return.\"\n", + " )\n", + " \n", + " def handle_single_conversation(conversation):\n", + " input_ids = self.get_prefix_tokens() if add_special_tokens else []\n", + " input_message = \"[gMASK]\" if add_special_tokens else \"\"\n", + " for item in conversation:\n", + " if item.get(\"tools\"):\n", + " tools = item[\"tools\"]\n", + " content = \"你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\"\n", + " for tool in tools:\n", + " if tool[\"type\"] == \"function\":\n", + " function = tool[\"function\"]\n", + " content += f\"\\n\\n## {function['name']}\\n\\n{json.dumps(function, ensure_ascii=False, indent=4)}\"\n", + " content += \"\\n在调用上述函数时，请使用 Json 格式表示调用的参数。\"\n", + " elif tool[\"type\"] == \"python\":\n", + " content += \"\\n\\n## python\\n\\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。\"\n", + " elif tool[\"type\"] == \"simple_browser\":\n", + " content += \"\\n\\n## simple_browser\\n\\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\\n`open_url(url: str)`：打开指定的 URL。\\n\\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\\n\\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。\"\n", + " elif tool[\"type\"] == \"cogview\":\n", + " content += \"\\n\\n## cogview\\n\\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。\"\n", + " else:\n", + " raise NotImplementedError(f\"Unknown tool type {tool['type']}\")\n", + " input = self.build_single_message(\"system\", \"\", content, tokenize=tokenize)\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if item[\"content\"]:\n", + " input = self.build_single_message(\n", + " item[\"role\"],\n", + " item.get(\"metadata\", \"\"),\n", + " item[\"content\"],\n", + " tokenize=tokenize\n", + " )\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if add_generation_prompt:\n", + " if tokenize:\n", + " input_ids.extend([self.convert_tokens_to_ids(\"<|assistant|>\")])\n", + " else:\n", + " input_message += \"<|assistant|>\"\n", + " # if tokenize:\n", + " # input_ids.extend([self.convert_tokens_to_ids(\"[gMASK]\")]) # 使用特殊标记代替空字符串\n", + " # else:\n", + " # input_message += \"[gMASK]\"\n", + " \n", + " return input_ids if tokenize else input_message\n", + " \n", + " # 处理不同会话格式的主逻辑\n", + " if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):\n", + " result = handle_single_conversation(conversation)\n", + " elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):\n", + " result = [handle_single_conversation(c) for c in conversation]\n", + " elif hasattr(conversation, \"messages\"):\n", + " result = handle_single_conversation(conversation.messages)\n", + " else:\n", + " raise ValueError(\"Invalid conversation format\")\n", + " \n", + " if tokenize:\n", + " output = self.batch_encode_plus(\n", + " [result] if isinstance(result[0], int) else result,\n", + " padding=padding,\n", + " truncation=truncation,\n", + " max_length=max_length,\n", + " return_tensors=return_tensors,\n", + " is_split_into_words=True,\n", + " add_special_tokens=False\n", + " )\n", + " if return_dict:\n", + " return output\n", + " else:\n", + " return output[\"input_ids\"]\n", + " else:\n", + " return result\n", + "\n", + "\n", + "\n", + " def build_inputs_with_special_tokens(\n", + " self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None\n", + " ) -> List[int]:\n", + " \"\"\"\n", + " 构建模型输入，适用于序列分类任务，通过连接和添加特殊 tokens。\n", + " BERT 序列格式：\n", + " - 单序列: `[CLS] X [SEP]`\n", + " - 序列对: `[CLS] A [SEP] B [SEP]`\n", + "\n", + " Args:\n", + " token_ids_0 (`List[int]`): 要添加特殊 tokens 的 ID 列表。\n", + " token_ids_1 (`List[int]`, *optional*): 可选的第二个 ID 列表，表示序列对。\n", + "\n", + " Returns:\n", + " `List[int]`: 添加了适当特殊 tokens 的输入 ID 列表。\n", + " \"\"\"\n", + " prefix_tokens = self.get_prefix_tokens()\n", + " token_ids_0 = prefix_tokens + token_ids_0\n", + " if token_ids_1 is not None:\n", + " token_ids_0 = token_ids_0 + token_ids_1 + [self.convert_tokens_to_ids(\"\")]\n", + " return token_ids_0\n", + "\n", + " def _pad(\n", + " self,\n", + " encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],\n", + " max_length: Optional[int] = None,\n", + " padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,\n", + " pad_to_multiple_of: Optional[int] = None,\n", + " return_attention_mask: Optional[bool] = None,\n", + " ) -> dict:\n", + " \"\"\"\n", + " 填充编码后的输入（左右填充，并预定义长度或批处理中的最大长度）\n", + "\n", + " Args:\n", + " encoded_inputs: 包含编码后的输入 (`List[int]`) 或批处理的输入 (`List[List[int]]`) 的字典。\n", + " max_length: 返回列表的最大长度，并可选的填充长度。\n", + " padding_strategy: 用于填充的策略。\n", + " - PaddingStrategy.LONGEST: 填充到批处理中的最长序列\n", + " - PaddingStrategy.MAX_LENGTH: 填充到最大长度（默认）\n", + " - PaddingStrategy.DO_NOT_PAD: 不填充\n", + " 填充策略由 self.padding_side 定义：\n", + " - 'left': 在序列左侧填充\n", + " - 'right': 在序列右侧填充\n", + " pad_to_multiple_of: (可选) 整数，如果设置将填充序列为该值的倍数。\n", + " return_attention_mask: (可选) 设置为 False 以避免返回注意力掩码（默认：根据模型具体情况设置）\n", + " \"\"\"\n", + " # 确保填充侧为 'left'\n", + " assert self.padding_side == \"left\"\n", + "\n", + " required_input = encoded_inputs[self.model_input_names[0]]\n", + " seq_length = len(required_input)\n", + "\n", + " if padding_strategy == PaddingStrategy.LONGEST:\n", + " max_length = len(required_input)\n", + "\n", + " if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):\n", + " max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of\n", + "\n", + " needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length\n", + "\n", + " # 如果没有注意力掩码，初始化\n", + " if \"attention_mask\" not in encoded_inputs:\n", + " encoded_inputs[\"attention_mask\"] = [1] * seq_length\n", + "\n", + " if \"position_ids\" not in encoded_inputs:\n", + " encoded_inputs[\"position_ids\"] = list(range(seq_length))\n", + "\n", + " if needs_to_be_padded:\n", + " difference = max_length - len(required_input)\n", + "\n", + " if \"attention_mask\" in encoded_inputs:\n", + " encoded_inputs[\"attention_mask\"] = [0] * difference + encoded_inputs[\"attention_mask\"]\n", + " if \"position_ids\" in encoded_inputs:\n", + " encoded_inputs[\"position_ids\"] = [0] * difference + encoded_inputs[\"position_ids\"]\n", + " encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input\n", + "\n", + " return encoded_inputs\n" + ] + }, + { + "cell_type": "markdown", + "id": "3bd7c0c2-83eb-4f01-a9c6-0b00ef9d05b2", + "metadata": {}, + "source": [ + "### 关键点解释\n", + "\n", + "1. **`mergeable_ranks`**:\n", + " - 存储可合并词汇的排名信息。读取词汇文件时，每个词汇都有一个对应的排名（整数），这些词汇被解码为原始字符串，然后存储在 `mergeable_ranks` 字典中。\n", + "\n", + "2. **`pat_str`**:\n", + " - 正则表达式字符串，用于定义分词模式。这段正则表达式主要用于匹配常见的英语缩写、单词、数字和其他非字母数字字符。编译后的正则表达式存储在 `self.pat_str` 中，用于分词器。\n", + "\n", + "3. **`self.tokenizer`**:\n", + " - 使用 `tiktoken` 库创建的编码器。通过提供的正则表达式模式和可合并的词汇排名，初始化一个自定义的分词器。\n", + "\n", + "4. **`self.decoder`**:\n", + " - 解码器，映射词汇排名到词汇。使用 `mergeable_ranks` 字典创建的逆向字典，用于从 ID 转换回原始词汇。\n", + "\n", + "5. **`self.n_words`**:\n", + " - 词汇表大小，即词汇数量。通过计算 `self.decoder` 的长度获得。\n", + "\n", + "6. **初始化过程**：\n", + " - 读取词汇文件，解析每一行以获取词汇及其排名，然后解码词汇并存储到 `mergeable_ranks` 字典中。接着，使用 `tiktoken` 库初始化编码器，并创建解码器和词汇表大小。\n", + "\n", + "其中有一点值得关注：\n", + "```python\n", + " def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:\n", + " \"\"\"\n", + " 将 tokens 序列转换为字符串。\n", + " \"\"\"\n", + " text = \"\"\n", + " temp = b\"\"\n", + " for t in tokens:\n", + " if isinstance(t, str):\n", + " if temp:\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " temp = b\"\"\n", + " text += t\n", + " elif isinstance(t, bytes):\n", + " temp += t\n", + " else:\n", + " raise TypeError(\"token should only be of type bytes or str\")\n", + " if temp:\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " return text\n", + "```\n", + "\n", + "这个函数之所以要将 `token` 解码为 `bytes` 是因为在实际应用中，token 可能是字符串或者字节序列的一部分。在某些情况下，词汇表中的 token 可能以 `bytes` 的形式存储，而不是直接存储为字符串。以下是一些详细原因和场景：\n", + "\n", + "1. **支持多种编码格式**：\n", + " - 分词器可能处理多种数据源，其中一些数据源可能以 `bytes` 格式存储 token。例如，当数据被压缩、加密或使用特定编码时，token 可能会以字节形式表示。\n", + "\n", + "2. **数据处理的灵活性**：\n", + " - 通过支持 `bytes` 和 `str` 两种类型，分词器可以更灵活地处理不同来源和格式的输入数据。这在处理包含二进制数据的文本时特别有用，例如一些特殊的标记或字符。\n", + "\n", + "3. **确保数据一致性**：\n", + " - 在分词和解码过程中，可能会遇到混合类型的 token（即既有字符串也有字节序列）。为了确保数据一致性并正确拼接字符串，函数需要处理 `bytes` 和 `str` 两种类型。\n", + "\n", + "4. **兼容性考虑**：\n", + " - 某些 NLP 工具和库在处理文本时可能会返回字节形式的 token。为了与这些工具和库兼容，分词器需要能够处理和转换这些字节 token。\n", + "\n", + "具体来看这个函数的工作流程：\n", + "\n", + "1. **初始化空字符串和字节序列**：\n", + " - `text = \"\"` 初始化一个空字符串，用于存储最终的结果。\n", + " - `temp = b\"\"` 初始化一个空字节序列，用于临时存储字节 token。\n", + "\n", + "2. **遍历 token 列表**：\n", + " - 对于每个 token，检查其类型。\n", + " - 如果 token 是字符串类型（`str`），检查 `temp` 是否为空。如果 `temp` 非空，将 `temp` 解码为字符串并添加到 `text`，然后清空 `temp`。之后，将当前字符串 token 添加到 `text`。\n", + " - 如果 token 是字节类型（`bytes`），将其添加到 `temp`，以便后续解码。\n", + " - 如果 token 既不是字符串也不是字节类型，抛出类型错误。\n", + "\n", + "3. **处理剩余的字节序列**：\n", + " - 在循环结束后，如果 `temp` 非空，将其解码为字符串并添加到 `text`。\n", + "\n", + "通过上述步骤，函数能够正确处理混合类型的 token 列表，并将其转换为一个完整的字符串。以下是代码中的注释，以更好地解释这些步骤：\n", + "\n", + "```python\n", + "def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:\n", + " \"\"\"\n", + " 将 tokens 序列转换为字符串。\n", + " \"\"\"\n", + " text = \"\" # 初始化一个空字符串用于存储结果\n", + " temp = b\"\" # 初始化一个空字节序列用于临时存储字节 token\n", + " for t in tokens:\n", + " if isinstance(t, str):\n", + " if temp:\n", + " # 如果 temp 非空，将其解码为字符串并添加到 text 中\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " temp = b\"\" # 清空 temp\n", + " text += t # 将字符串 token 添加到 text 中\n", + " elif isinstance(t, bytes):\n", + " temp += t # 将字节 token 添加到 temp 中\n", + " else:\n", + " raise TypeError(\"token should only be of type bytes or str\") # 抛出类型错误\n", + " if temp:\n", + " # 处理剩余的字节序列，将其解码为字符串并添加到 text 中\n", + " text += temp.decode(\"utf-8\", errors=\"replace\")\n", + " return text # 返回最终的字符串\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "2c439bb2-5250-4891-973d-f973bd14aae3", + "metadata": {}, + "source": [ + "接下来设置config，简单设置相关属性即可。" + ] + }, + { + "cell_type": "code", + "execution_count": 62, + "id": "ac8bdb0a-b279-4365-b6bb-224b0fe43d99", + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import PretrainedConfig\n", + "\n", + "\n", + "class ChatGLMConfig(PretrainedConfig):\n", + " model_type = \"chatglm\"\n", + "\n", + " def __init__(\n", + " self,\n", + " num_layers=28,\n", + " padded_vocab_size=65024,\n", + " hidden_size=4096,\n", + " ffn_hidden_size=13696,\n", + " kv_channels=128,\n", + " num_attention_heads=32,\n", + " seq_length=2048,\n", + " hidden_dropout=0.0,\n", + " classifier_dropout=None,\n", + " attention_dropout=0.0,\n", + " layernorm_epsilon=1e-5,\n", + " rmsnorm=True,\n", + " apply_residual_connection_post_layernorm=False,\n", + " post_layer_norm=True,\n", + " add_bias_linear=False,\n", + " add_qkv_bias=False,\n", + " bias_dropout_fusion=True,\n", + " multi_query_attention=False,\n", + " multi_query_group_num=1,\n", + " rope_ratio=1,\n", + " apply_query_key_layer_scaling=True,\n", + " attention_softmax_in_fp32=True,\n", + " fp32_residual_connection=False,\n", + " **kwargs\n", + " ):\n", + " self.num_layers = num_layers\n", + " self.vocab_size = padded_vocab_size\n", + " self.padded_vocab_size = padded_vocab_size\n", + " self.hidden_size = hidden_size\n", + " self.ffn_hidden_size = ffn_hidden_size\n", + " self.kv_channels = kv_channels\n", + " self.num_attention_heads = num_attention_heads\n", + " self.seq_length = seq_length\n", + " self.hidden_dropout = hidden_dropout\n", + " self.classifier_dropout = classifier_dropout\n", + " self.attention_dropout = attention_dropout\n", + " self.layernorm_epsilon = layernorm_epsilon\n", + " self.rmsnorm = rmsnorm\n", + " self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm\n", + " self.post_layer_norm = post_layer_norm\n", + " self.add_bias_linear = add_bias_linear\n", + " self.add_qkv_bias = add_qkv_bias\n", + " self.bias_dropout_fusion = bias_dropout_fusion\n", + " self.multi_query_attention = multi_query_attention\n", + " self.multi_query_group_num = multi_query_group_num\n", + " self.rope_ratio = rope_ratio\n", + " self.apply_query_key_layer_scaling = apply_query_key_layer_scaling\n", + " self.attention_softmax_in_fp32 = attention_softmax_in_fp32\n", + " self.fp32_residual_connection = fp32_residual_connection\n", + " super().__init__(**kwargs)\n" + ] + }, + { + "cell_type": "markdown", + "id": "7936d603-79a0-41ae-8777-185c5562e0bb", + "metadata": {}, + "source": [ + "然后就可以开始构建我们的模型了。" + ] + }, + { + "cell_type": "code", + "execution_count": 63, + "id": "851a39f0-1cad-48c2-a0ce-b781ed98c07f", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import math\n", + "import copy\n", + "import warnings\n", + "import re\n", + "import sys\n", + "\n", + "import torch\n", + "import torch.utils.checkpoint\n", + "import torch.nn.functional as F\n", + "from torch import nn\n", + "from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss\n", + "from torch.nn.utils import skip_init\n", + "from typing import Optional, Tuple, Union, List, Callable, Dict, Any\n", + "from copy import deepcopy\n", + "\n", + "from transformers.modeling_outputs import (\n", + " BaseModelOutputWithPast,\n", + " CausalLMOutputWithPast,\n", + " SequenceClassifierOutputWithPast,\n", + ")\n", + "from transformers.modeling_utils import PreTrainedModel\n", + "from transformers.utils import logging\n", + "from transformers.generation.logits_process import LogitsProcessor\n", + "from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput" + ] + }, + { + "cell_type": "code", + "execution_count": 64, + "id": "8ff61078-3447-4b2d-a77a-034000ce38d2", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'linux'" + ] + }, + "execution_count": 64, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sys.platform" + ] + }, + { + "cell_type": "code", + "execution_count": 65, + "id": "eaa943a3-5a6a-4f5a-9d1c-a1d2619e090e", + "metadata": {}, + "outputs": [], + "source": [ + "# 如果系统平台不是 Darwin (即 macOS或 linux)，则进行以下设置以提高性能\n", + "if sys.platform != 'darwin':\n", + " # 关闭 JIT 的 profiling 模式\n", + " torch._C._jit_set_profiling_mode(False)\n", + " # 关闭 JIT 的 profiling 执行器\n", + " torch._C._jit_set_profiling_executor(False)\n", + " # 允许在 CPU 上进行张量融合\n", + " torch._C._jit_override_can_fuse_on_cpu(True)\n", + " # 允许在 GPU 上进行张量融合\n", + " torch._C._jit_override_can_fuse_on_gpu(True)\n", + "\n", + "# 获取日志记录器\n", + "logger = logging.get_logger(__name__)\n", + "\n", + "# 用于文档的检查点\n", + "_CHECKPOINT_FOR_DOC = \"THUDM/ChatGLM\"\n", + "# 用于文档的配置\n", + "_CONFIG_FOR_DOC = \"ChatGLMConfig\"\n", + "\n", + "# 默认初始化函数\n", + "def default_init(cls, *args, **kwargs):\n", + " return cls(*args, **kwargs)\n", + "\n", + "# 无效分数 Logits 处理器类\n", + "class InvalidScoreLogitsProcessor(LogitsProcessor):\n", + " # 处理输入的 logits\n", + " def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:\n", + " # 如果 logits 中存在 NaN 或无穷大，则将其重置\n", + " if torch.isnan(scores).any() or torch.isinf(scores).any():\n", + " scores.zero_() # 将所有分数重置为 0\n", + " scores[..., 198] = 5e4 # 将特定位置的分数设置为一个大值\n", + " return scores\n", + "\n", + "# 沿着最后一个维度分割张量的函数\n", + "def split_tensor_along_last_dim(\n", + " tensor: torch.Tensor,\n", + " num_partitions: int,\n", + " contiguous_split_chunks: bool = False,\n", + ") -> List[torch.Tensor]:\n", + " \"\"\"\n", + " 沿着最后一个维度分割张量。\n", + "\n", + " 参数:\n", + " tensor: 输入张量。\n", + " num_partitions: 要分割的部分数量。\n", + " contiguous_split_chunks: 如果为 True，则使每个分块在内存中是连续的。\n", + "\n", + " 返回:\n", + " 张量的列表\n", + " \"\"\"\n", + " # 获取最后一个维度的大小\n", + " last_dim = tensor.dim() - 1\n", + " last_dim_size = tensor.size()[last_dim] // num_partitions\n", + " # 分割张量\n", + " tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)\n", + " # 注意: torch.split 默认不会创建连续的张量\n", + " if contiguous_split_chunks:\n", + " return tuple(chunk.contiguous() for chunk in tensor_list)\n", + " return tensor_list\n" + ] + }, + { + "cell_type": "markdown", + "id": "52adb283-3a15-4a82-a3e5-b80f2066c891", + "metadata": {}, + "source": [ + "### 关键点解释\n", + "\n", + "1. **系统平台检查**：\n", + " - 如果系统平台不是 `Darwin`（macOS），则关闭 JIT 的 profiling 模式和执行器，并启用 CPU 和 GPU 上的张量融合。这些设置可以提高模型的性能。\n", + "\n", + "2. **日志记录器**：\n", + " - 获取一个日志记录器，用于记录模型构建过程中的信息和调试。\n", + "\n", + "3. **文档字符串**：\n", + " - `_CHECKPOINT_FOR_DOC` 和 `_CONFIG_FOR_DOC` 是用于文档生成的字符串常量，指示模型和配置的检查点和配置文件。\n", + "\n", + "4. **默认初始化函数**：\n", + " - `default_init` 函数是一个通用的初始化函数，用于实例化类。\n", + "\n", + "5. **`InvalidScoreLogitsProcessor` 类**：\n", + " - 这是一个 logits 处理器类，用于处理无效的 logits 分数。如果分数中存在 `NaN` 或无穷大值，将这些值重置为 0，并将特定位置（例如 198）的分数设置为一个很大的值（`5e4`），以确保模型输出有效的结果。\n", + "\n", + "6. **`split_tensor_along_last_dim` 函数**：\n", + " - 这个函数用于沿着张量的最后一个维度分割张量。参数包括输入张量、要分割的部分数量和一个布尔值（是否使每个分块在内存中是连续的）。如果 `contiguous_split_chunks` 为真，则每个分块在内存中是连续的；否则，直接返回分割后的张量列表。\n" + ] + }, + { + "cell_type": "code", + "execution_count": 66, + "id": "7f8005f7-1339-40b4-b103-6bf0396b371a", + "metadata": {}, + "outputs": [], + "source": [ + "# 旋转位置嵌入类\n", + "class RotaryEmbedding(nn.Module):\n", + " def __init__(self, dim, rope_ratio=1, original_impl=False, device=None, dtype=None):\n", + " super().__init__()\n", + " # 计算倒数频率\n", + " inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))\n", + " # 注册倒数频率为 buffer\n", + " self.register_buffer(\"inv_freq\", inv_freq)\n", + " self.dim = dim\n", + " self.original_impl = original_impl\n", + " self.rope_ratio = rope_ratio\n", + "\n", + " # 实现前向传播的具体方法\n", + " def forward_impl(\n", + " self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000\n", + " ):\n", + " \"\"\"\n", + " 增强的 Transformer 使用旋转位置嵌入。\n", + "\n", + " 参考自:\n", + " https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/\n", + " transformers/rope/__init__.py. MIT 许可证:\n", + " https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.\n", + " \"\"\"\n", + " # 计算 $\\Theta = {\\theta_i = 10000^{\\frac{2(i-1)}{d}}, i \\in [1, 2, ..., \\frac{d}{2}]}$\n", + " base = base * self.rope_ratio\n", + " theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem))\n", + "\n", + " # 创建位置索引 `[0, 1, ..., seq_len - 1]`\n", + " seq_idx = torch.arange(seq_len, dtype=torch.float, device=device)\n", + "\n", + " # 计算位置索引和 $\\theta_i$ 的乘积\n", + " idx_theta = torch.outer(seq_idx, theta).float()\n", + "\n", + " # 缓存计算的余弦和正弦值\n", + " cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)\n", + "\n", + " # 模拟 complex32 的行为，否则会得到不同的结果\n", + " if dtype in (torch.float16, torch.bfloat16, torch.int8):\n", + " cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()\n", + " return cache\n", + "\n", + " # 前向传播方法\n", + " def forward(self, max_seq_len, offset=0):\n", + " return self.forward_impl(\n", + " max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device\n", + " )\n", + "\n", + "# 使用 TorchScript JIT 编译器优化的旋转位置嵌入应用函数\n", + "@torch.jit.script\n", + "def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:\n", + " # x: [b, np, sq, hn]\n", + " b, np, sq, hn = x.size(0), x.size(1), x.size(2), x.size(3)\n", + " rot_dim = rope_cache.shape[-2] * 2\n", + " x, x_pass = x[..., :rot_dim], x[..., rot_dim:]\n", + " # 截断以支持可变大小\n", + " rope_cache = rope_cache[:, :sq]\n", + " xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)\n", + " rope_cache = rope_cache.view(-1, 1, sq, xshaped.size(3), 2)\n", + " x_out2 = torch.stack(\n", + " [\n", + " xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],\n", + " xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],\n", + " ],\n", + " -1,\n", + " )\n", + " x_out2 = x_out2.flatten(3)\n", + " return torch.cat((x_out2, x_pass), dim=-1)\n", + "\n", + "# RMS 归一化层类\n", + "class RMSNorm(torch.nn.Module):\n", + " def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):\n", + " super().__init__()\n", + " # 初始化权重参数\n", + " self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))\n", + " self.eps = eps\n", + "\n", + " def forward(self, hidden_states: torch.Tensor):\n", + " # 获取输入张量的数据类型\n", + " input_dtype = hidden_states.dtype\n", + " # 计算方差\n", + " variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)\n", + " # 进行 RMS 归一化\n", + " hidden_states = hidden_states * torch.rsqrt(variance + self.eps)\n", + " # 应用权重并返回与输入相同的数据类型\n", + " return (self.weight * hidden_states).to(input_dtype)" + ] + }, + { + "cell_type": "markdown", + "id": "c1f66720-11b0-4bc5-9468-083e952ce6de", + "metadata": {}, + "source": [ + "### 关键点解释\n", + "\n", + "1. **`RotaryEmbedding` 类**：\n", + " - 用于实现旋转位置嵌入的模块。\n", + " - `inv_freq` 是一个倒数频率的张量，用于位置嵌入的计算。\n", + " - `forward_impl` 方法根据序列长度和嵌入维度计算旋转位置嵌入。\n", + " - `forward` 方法调用 `forward_impl` 进行前向传播。\n", + "\n", + "2. **`apply_rotary_pos_emb` 函数**：\n", + " - 使用旋转位置嵌入更新输入张量。\n", + " - `x` 是输入张量，`rope_cache` 是预计算的旋转位置嵌入。\n", + " - 该函数首先对输入张量进行分割和重塑，然后将旋转位置嵌入应用于每个分块，最后将结果拼接回原始张量。\n", + "\n", + "3. **`RMSNorm` 类**：\n", + " - 实现 RMS 归一化的模块。\n", + " - `weight` 是归一化的权重参数。\n", + " - `forward` 方法计算输入张量的方差，并进行归一化处理，然后应用权重。\n", + "\n", + "这些注释和解释可以帮助理解每个部分的功能和实现细节，对于模型构建和调试非常有用。" + ] + }, + { + "cell_type": "code", + "execution_count": 67, + "id": "b76dc446-e37e-4c1f-94f1-96f209bd674e", + "metadata": {}, + "outputs": [], + "source": [ + "class CoreAttention(torch.nn.Module):\n", + " def __init__(self, config: ChatGLMConfig, layer_number):\n", + " super(CoreAttention, self).__init__()\n", + "\n", + " # 配置参数\n", + " self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling\n", + " self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32\n", + " if self.apply_query_key_layer_scaling:\n", + " self.attention_softmax_in_fp32 = True\n", + " self.layer_number = max(1, layer_number)\n", + "\n", + " # 计算投影大小\n", + " projection_size = config.kv_channels * config.num_attention_heads\n", + "\n", + " # 每个注意力头和每个分区的值\n", + " self.hidden_size_per_partition = projection_size\n", + " self.hidden_size_per_attention_head = projection_size // config.num_attention_heads\n", + " self.num_attention_heads_per_partition = config.num_attention_heads\n", + "\n", + " coeff = None\n", + " self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)\n", + " if self.apply_query_key_layer_scaling:\n", + " coeff = self.layer_number\n", + " self.norm_factor *= coeff\n", + " self.coeff = coeff\n", + "\n", + " # 注意力 dropout\n", + " self.attention_dropout = torch.nn.Dropout(config.attention_dropout)\n", + "\n", + " def forward(self, query_layer, key_layer, value_layer, attention_mask):\n", + " pytorch_major_version = int(torch.__version__.split('.')[0])\n", + " if pytorch_major_version >= 2:\n", + " # PyTorch 2.0 及以上版本\n", + " if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:\n", + " context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,\n", + " is_causal=True)\n", + " else:\n", + " if attention_mask is not None:\n", + " attention_mask = ~attention_mask\n", + " context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,\n", + " attention_mask)\n", + " context_layer = context_layer.transpose(1, 2).contiguous()\n", + " new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)\n", + " context_layer = context_layer.reshape(*new_context_layer_shape)\n", + " else:\n", + " # 处理 PyTorch 2.0 以下版本\n", + "\n", + " # 原始注意力得分\n", + " # [b, np, sq, sk]\n", + " output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))\n", + "\n", + " # 重新调整视图 [b, np, sq, hn] -> [b * np, sq, hn]\n", + " query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)\n", + " # 重新调整视图 [b, np, sk, hn] -> [b * np, sk, hn]\n", + " key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)\n", + "\n", + " # 预分配输入张量: [b * np, sq, sk]\n", + " matmul_input_buffer = torch.empty(\n", + " output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,\n", + " device=query_layer.device\n", + " )\n", + "\n", + " # 计算原始注意力得分. [b * np, sq, sk]\n", + " matmul_result = torch.baddbmm(\n", + " matmul_input_buffer,\n", + " query_layer, # [b * np, sq, hn]\n", + " key_layer.transpose(1, 2), # [b * np, hn, sk]\n", + " beta=0.0,\n", + " alpha=(1.0 / self.norm_factor),\n", + " )\n", + "\n", + " # 改变视图到 [b, np, sq, sk]\n", + " attention_scores = matmul_result.view(*output_size)\n", + "\n", + " # 处理注意力得分和 dropout\n", + " if self.attention_softmax_in_fp32:\n", + " attention_scores = attention_scores.float()\n", + " if self.coeff is not None:\n", + " attention_scores = attention_scores * self.coeff\n", + " if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:\n", + " attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],\n", + " device=attention_scores.device, dtype=torch.bool)\n", + " attention_mask.tril_()\n", + " attention_mask = ~attention_mask\n", + " if attention_mask is not None:\n", + " attention_scores = attention_scores.masked_fill(attention_mask, float(\"-inf\"))\n", + " attention_probs = F.softmax(attention_scores, dim=-1)\n", + " attention_probs = attention_probs.type_as(value_layer)\n", + "\n", + " # 丢弃整个 token 的注意力，这源自原始的 Transformer 论文\n", + " attention_probs = self.attention_dropout(attention_probs)\n", + "\n", + " # 重新调整视图 [b * np, sq, hn]\n", + " value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)\n", + " # 重新调整视图 [b * np, sq, sk]\n", + " attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)\n", + " # 计算上下文层 [b * np, sq, hn]\n", + " context_layer = torch.bmm(attention_probs, value_layer)\n", + " # 重新调整视图 [b, np, sq, hn]\n", + " context_layer = context_layer.view(*output_size)\n", + " # [b, np, sq, hn] --> [b, sq, np, hn]\n", + " context_layer = context_layer.transpose(1, 2).contiguous()\n", + " # [b, sq, np, hn] --> [b, sq, hp]\n", + " new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)\n", + " context_layer = context_layer.reshape(*new_context_layer_shape)\n", + "\n", + " return context_layer" + ] + }, + { + "cell_type": "markdown", + "id": "1a6f04a8-0373-4a78-baa9-52e665fe604d", + "metadata": {}, + "source": [ + "### CoreAttention 详解及公式\n", + "\n", + "在神经网络模型中，特别是 Transformer 架构中，注意力机制起着至关重要的作用。CoreAttention 实现了自注意力机制中的核心部分，具体来说包括以下步骤：计算注意力得分、应用注意力掩码、计算注意力权重以及计算上下文向量。\n", + "\n", + "#### 1. 计算注意力得分 (Attention Scores)\n", + "\n", + "注意力得分的计算可以表示为矩阵乘法。对于查询 (query) 向量 $Q$ 和键 (key) 向量 $K$，计算注意力得分矩阵 $A$ 的公式为：\n", + "\\begin{align*} A = \\frac{QK^T}{\\sqrt{d_k}} \\end{align*}\n", + "其中 $d_k$ 是键向量的维度，这里的 $\\sqrt{d_k}$ 是一个缩放因子，防止得分值过大。\n", + "\n", + "在代码中，通过矩阵乘法实现：\n", + "```python\n", + "matmul_result = torch.baddbmm(\n", + " matmul_input_buffer,\n", + " query_layer, # [b * np, sq, hn]\n", + " key_layer.transpose(1, 2), # [b * np, hn, sk]\n", + " beta=0.0,\n", + " alpha=(1.0 / self.norm_factor),\n", + ")\n", + "```\n", + "其中 `self.norm_factor` 是 $\\sqrt{d_k}$。\n", + "\n", + "#### 2. 应用注意力掩码 (Attention Mask)\n", + "\n", + "为了避免模型关注不必要的部分，使用注意力掩码来屏蔽某些位置。注意力掩码通过将相应位置的注意力得分设为负无穷大来实现：\n", + "```python\n", + "if attention_mask is not None:\n", + " attention_scores = attention_scores.masked_fill(attention_mask, float(\"-inf\"))\n", + "```\n", + "\n", + "#### 3. 计算注意力权重 (Attention Weights)\n", + "\n", + "应用 softmax 函数将注意力得分转换为概率分布，表示每个查询向量对键向量的关注程度：\n", + "\\begin{align*} \\text{AttentionWeights} = \\text{softmax}(A) \\end{align*}\n", + "代码中实现为：\n", + "```python\n", + "attention_probs = F.softmax(attention_scores, dim=-1)\n", + "attention_probs = attention_probs.type_as(value_layer)\n", + "```\n", + "\n", + "#### 4. 计算上下文向量 (Context Vectors)\n", + "\n", + "上下文向量通过将注意力权重与值 (value) 向量相乘并求和得到：\n", + "\\begin{align*} \\text{Context} = \\text{AttentionWeights} \\cdot V \\end{align*}\n", + "其中 $V$ 是值向量。\n", + "\n", + "在代码中实现为：\n", + "```python\n", + "context_layer = torch.bmm(attention_probs, value_layer)\n", + "```\n", + "\n", + "#### 详细原理解释\n", + "\n", + "1. **初始化和设置**：\n", + " - 初始化类时，会根据配置参数设置注意力缩放、层数等信息。\n", + " - `projection_size` 定义了投影大小，`hidden_size_per_attention_head` 和 `num_attention_heads_per_partition` 定义了每个注意力头和每个分区的隐藏层大小。\n", + "\n", + "2. **前向传播步骤**：\n", + " - **查询、键和值的计算**：计算查询、键和值向量。\n", + " - **计算注意力得分**：通过矩阵乘法计算注意力得分矩阵。\n", + " - **应用注意力掩码**：将需要屏蔽的位置设置为负无穷大，避免影响后续计算。\n", + " - **计算注意力权重**：应用 softmax 函数，得到注意力权重。\n", + " - **计算上下文向量**：将注意力权重与值向量相乘，得到上下文向量。\n", + "\n", + "3. **具体实现细节**：\n", + " - 根据 PyTorch 版本，选择合适的注意力计算方式。\n", + " - 对不同维度的张量进行变换，确保形状匹配。\n", + " - 使用 dropout 防止过拟合。\n", + "\n", + "### 公式和代码对应关系\n", + "\n", + "- **注意力得分**：\n", + " \\begin{align*}\n", + " A = \\frac{QK^T}{\\sqrt{d_k}}\n", + " \\end{align*}\n", + " 对应代码：\n", + " ```python\n", + " matmul_result = torch.baddbmm(\n", + " matmul_input_buffer,\n", + " query_layer,\n", + " key_layer.transpose(1, 2),\n", + " beta=0.0,\n", + " alpha=(1.0 / self.norm_factor),\n", + " )\n", + " ```\n", + "\n", + "- **应用注意力掩码**：\n", + " \\begin{align*}\n", + " A'_{ij} = \\begin{cases} \n", + " A_{ij} & \\text{if } \\text{mask}_{ij} = 1 \\\\\n", + " -\\infty & \\text{if } \\text{mask}_{ij} = 0 \n", + " \\end{cases}\n", + " \\end{align*}\n", + " 对应代码：\n", + " ```python\n", + " if attention_mask is not None:\n", + " attention_scores = attention_scores.masked_fill(attention_mask, float(\"-inf\"))\n", + " ```\n", + "\n", + "- **注意力权重**：\n", + " \\begin{align*}\n", + " \\text{AttentionWeights}_{ij} = \\frac{\\exp(A'_{ij})}{\\sum_k \\exp(A'_{ik})}\n", + " \\end{align*}\n", + " 对应代码：\n", + " ```python\n", + " attention_probs = F.softmax(attention_scores, dim=-1)\n", + " ```\n", + "\n", + "- **上下文向量**：\n", + " \\begin{align*}\n", + " \\text{Context}_{i} = \\sum_j \\text{AttentionWeights}_{ij} V_j\n", + " \\end{align*}\n", + " 对应代码：\n", + " ```python\n", + " context_layer = torch.bmm(attention_probs, value_layer)\n", + " ```\n" + ] + }, + { + "cell_type": "markdown", + "id": "47fcf489-71b3-4077-af03-990979abc001", + "metadata": {}, + "source": [ + "接下来，我们再重温一下注意力机制\n", + "\n", + "### 注意力机制原理及公式解释\n", + "\n", + "#### 1. 注意力机制是什么？\n", + "注意力机制（Attention Mechanism）是深度学习中特别是自然语言处理领域中的一种技术，它使模型能够在处理输入序列时动态地关注不同的部分。简单来说，注意力机制让模型在处理某个元素时，可以有选择地关注输入序列的其他部分，而不是全部一视同仁。\n", + "\n", + "#### 2. 注意力机制的核心公式\n", + "我们来看注意力得分的公式：\n", + "\\begin{align*} A = \\frac{QK^T}{\\sqrt{d_k}} \\end{align*}\n", + "\n", + "这里的符号解释如下：\n", + "- \\( Q \\) 是查询（Query）向量。\n", + "- \\( K \\) 是键（Key）向量。\n", + "- \\( d_k \\) 是键向量的维度。\n", + "- \\( A \\) 是注意力得分矩阵。\n", + "\n", + "#### 3. 为什么使用这个公式计算注意力得分？\n", + "\n", + "注意力得分公式的核心思想是通过计算查询向量和键向量的点积，来衡量查询与每个键之间的相关性。点积结果越大，表示查询与该键的相关性越强，模型应该对该键对应的值（Value）向量给予更多的关注。\n", + "\n", + "公式中的缩放因子 \\(\\sqrt{d_k}\\) 是为了避免点积结果过大，因为如果不缩放，较大的向量维度会导致点积结果非常大，进而导致 softmax 函数输出接近于零的梯度，影响模型训练效果。\n", + "\n", + "#### 4. 计算注意力权重\n", + "得到注意力得分矩阵 \\( A \\) 后，通过 softmax 函数将其转换为注意力权重矩阵：\n", + "\\begin{align*} \\text{AttentionWeights}_{ij} = \\frac{\\exp(A'_{ij})}{\\sum_k \\exp(A'_{ik})} \\end{align*}\n", + "其中 \\( A' \\) 是应用掩码后的注意力得分矩阵。\n", + "\n", + "#### 5. 计算上下文向量\n", + "最后，使用注意力权重矩阵对值向量进行加权求和，得到上下文向量：\n", + "\\begin{align*} \\text{Context}_{i} = \\sum_j \\text{AttentionWeights}_{ij} V_j \\end{align*}\n", + "这里 \\( V \\) 是值向量。\n", + "\n", + "### 注意力机制与人类注意力的关系\n", + "\n", + "注意力机制与人类的注意力有一定的相似之处，但也有显著的区别。\n", + "\n", + "- **相似之处**：\n", + " - **选择性关注**：就像人类在阅读一篇文章时会选择性地关注某些重要段落，忽略其他部分，注意力机制也让模型在处理一个序列时可以选择性地关注不同的部分。\n", + " - **动态调整**：人类的注意力是动态的，会根据上下文调整关注点。注意力机制也是动态的，可以根据输入的变化调整注意力权重。\n", + "\n", + "- **区别**：\n", + " - **机制不同**：人类注意力是通过大脑的复杂神经网络实现的，包括视觉、听觉等多种感官信息的综合处理。而注意力机制是一种数学方法，通过点积、softmax 等操作实现。\n", + " - **目的不同**：人类注意力用于理解和互动，而注意力机制主要用于提高模型在处理长序列数据时的性能和效率。\n", + "\n", + "### 具体代码实现中的细节\n", + "\n", + "在 `CoreAttention` 类中，注意力得分的计算和应用通过以下代码片段实现：\n", + "```python\n", + "matmul_result = torch.baddbmm(\n", + " matmul_input_buffer,\n", + " query_layer, # [b * np, sq, hn]\n", + " key_layer.transpose(1, 2), # [b * np, hn, sk]\n", + " beta=0.0,\n", + " alpha=(1.0 / self.norm_factor),\n", + ")\n", + "```\n", + "这里 `torch.baddbmm` 函数执行的是批量矩阵乘法，计算公式中的 \\(QK^T\\) 部分，并除以 \\(\\sqrt{d_k}\\) 进行缩放。\n", + "\n", + "应用 softmax 得到注意力权重：\n", + "```python\n", + "attention_probs = F.softmax(attention_scores, dim=-1)\n", + "```\n", + "\n", + "最后，计算上下文向量：\n", + "```python\n", + "context_layer = torch.bmm(attention_probs, value_layer)\n", + "```\n", + "\n", + "这一步将注意力权重与值向量相乘并求和，得到最终的上下文表示。\n", + "\n", + "通过上述解释，可以更好地理解注意力机制的原理和实现，以及它在模型中的重要作用。\n", + "\n", + "完成注意力机制的核心实现后，我们构建自注意力" + ] + }, + { + "cell_type": "markdown", + "id": "7f0b003c-7e42-4e0a-96fe-a052826ea982", + "metadata": {}, + "source": [ + "### SelfAttention 类的原理及公式解释\n", + "\n", + "#### SelfAttention 类的功能\n", + "\n", + "SelfAttention 类实现了自注意力机制（Self-Attention Mechanism），这是 Transformer 模型的核心部分。自注意力机制的目标是让每个位置的表示能够动态地关注输入序列的其他位置，从而捕捉全局信息。\n", + "\n", + "#### 自注意力机制的步骤及公式\n", + "\n", + "自注意力机制包括以下几个关键步骤：\n", + "\n", + "1. **线性变换**：\n", + " 输入的隐藏状态 \\(X\\) 通过线性层分别映射到查询 \\(Q\\)、键 \\(K\\) 和值 \\(V\\) 三个空间：\n", + " \\begin{align*}\n", + " Q = XW_Q, \\quad K = XW_K, \\quad V = XW_V\n", + " \\end{align*}\n", + " 其中，\\(W_Q\\)、\\(W_K\\) 和 \\(W_V\\) 是学习到的权重矩阵。\n", + "\n", + "2. **计算注意力得分**：\n", + " 通过计算查询 \\(Q\\) 和键 \\(K\\) 的点积并除以缩放因子 \\(\\sqrt{d_k}\\) 得到注意力得分矩阵 \\(A\\)：\n", + " \\begin{align*}\n", + " A = \\frac{QK^T}{\\sqrt{d_k}}\n", + " \\end{align*}\n", + "\n", + "3. **应用注意力掩码**：\n", + " 对于自注意力机制，如果使用掩码（例如在解码阶段），将不需要关注的位置设为负无穷大以屏蔽：\n", + " \\begin{align*}\n", + " A'_{ij} = \\begin{cases} \n", + " A_{ij} & \\text{if } \\text{mask}_{ij} = 1 \\\\\n", + " -\\infty & \\text{if } \\text{mask}_{ij} = 0 \n", + " \\end{cases}\n", + " \\end{align*}\n", + "\n", + "4. **计算注意力权重**：\n", + " 对注意力得分矩阵 \\(A'\\) 应用 softmax 函数，得到注意力权重矩阵：\n", + " \\begin{align*}\n", + " \\text{AttentionWeights}_{ij} = \\frac{\\exp(A'_{ij})}{\\sum_k \\exp(A'_{ik})}\n", + " \\end{align*}\n", + "\n", + "5. **计算上下文向量**：\n", + " 使用注意力权重对值 \\(V\\) 进行加权求和，得到上下文向量：\n", + " \\begin{align*}\n", + " \\text{Context} = \\text{AttentionWeights} \\cdot V\n", + " \\end{align*}\n", + "\n", + "### SelfAttention 类的具体实现\n", + "\n", + "以下是 SelfAttention 类中的核心步骤和公式的实现细节：\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 68, + "id": "691ee97d-c7a7-4a41-846c-eb5f4a33443a", + "metadata": {}, + "outputs": [], + "source": [ + "class SelfAttention(torch.nn.Module):\n", + " \"\"\"并行自注意力层的抽象类。\n", + "\n", + " 自注意力层接受形状为 [s, b, h] 的输入，并返回相同形状的输出。\n", + " \"\"\"\n", + "\n", + " def __init__(self, config: ChatGLMConfig, layer_number, device=None):\n", + " super(SelfAttention, self).__init__()\n", + " self.layer_number = max(1, layer_number)\n", + "\n", + " self.projection_size = config.kv_channels * config.num_attention_heads\n", + "\n", + " # 每个注意力头和每个分区的值\n", + " self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads\n", + " self.num_attention_heads_per_partition = config.num_attention_heads\n", + "\n", + " self.multi_query_attention = config.multi_query_attention\n", + " self.qkv_hidden_size = 3 * self.projection_size\n", + " if self.multi_query_attention:\n", + " self.num_multi_query_groups_per_partition = config.multi_query_group_num\n", + " self.qkv_hidden_size = (\n", + " self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num\n", + " )\n", + " self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,\n", + " bias=config.add_bias_linear or config.add_qkv_bias,\n", + " device=device, **_config_to_kwargs(config)\n", + " )\n", + "\n", + " self.core_attention = CoreAttention(config, self.layer_number)\n", + "\n", + " # 输出层\n", + " self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,\n", + " device=device, **_config_to_kwargs(config)\n", + " )\n", + "\n", + " def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True):\n", + " # hidden_states: [b, sq, h]\n", + "\n", + " # =====================\n", + " # Query, Key 和 Value\n", + " # =====================\n", + "\n", + " # 注意力头 [b, sq, h] --> [b, sq, (np * 3 * hn)]\n", + " mixed_x_layer = self.query_key_value(hidden_states)\n", + "\n", + " if self.multi_query_attention:\n", + " (query_layer, key_layer, value_layer) = mixed_x_layer.split(\n", + " [\n", + " self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,\n", + " self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,\n", + " self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,\n", + " ],\n", + " dim=-1,\n", + " )\n", + " query_layer = query_layer.view(\n", + " query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)\n", + " )\n", + " key_layer = key_layer.view(\n", + " key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)\n", + " )\n", + " value_layer = value_layer.view(\n", + " value_layer.size()[:-1]\n", + " + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)\n", + " )\n", + " else:\n", + " new_tensor_shape = mixed_x_layer.size()[:-1] + \\\n", + " (self.num_attention_heads_per_partition,\n", + " 3 * self.hidden_size_per_attention_head)\n", + " mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)\n", + "\n", + " # [b, sq, np, 3 * hn] --> 3 [b, sq, np, hn]\n", + " (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)\n", + "\n", + " # [b, sq, np, hn] -> [b, np, sq, hn]\n", + " query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]\n", + "\n", + " # 应用相对位置编码（旋转嵌入）\n", + " if rotary_pos_emb is not None:\n", + " query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)\n", + " key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)\n", + "\n", + " # 调整 key 和 value 以用于推理\n", + " if kv_cache is not None:\n", + " cache_k, cache_v = kv_cache\n", + " key_layer = torch.cat((cache_k, key_layer), dim=2)\n", + " value_layer = torch.cat((cache_v, value_layer), dim=2)\n", + " if use_cache:\n", + " if kv_cache is None:\n", + " kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)\n", + " else:\n", + " kv_cache = (key_layer, value_layer)\n", + " else:\n", + " kv_cache = None\n", + "\n", + " if self.multi_query_attention:\n", + " key_layer = key_layer.unsqueeze(2)\n", + " key_layer = key_layer.expand(\n", + " -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1\n", + " )\n", + " key_layer = key_layer.contiguous().view(\n", + " key_layer.size()[:1] + (self.num_attention_heads_per_partition,) + key_layer.size()[3:]\n", + " )\n", + " value_layer = value_layer.unsqueeze(2)\n", + " value_layer = value_layer.expand(\n", + " -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1\n", + " )\n", + " value_layer = value_layer.contiguous().view(\n", + " value_layer.size()[:1] + (self.num_attention_heads_per_partition,) + value_layer.size()[3:]\n", + " )\n", + "\n", + " # 核心注意力计算\n", + " context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)\n", + "\n", + " # 输出. [sq, b, h]\n", + " output = self.dense(context_layer)\n", + "\n", + " return output, kv_cache" + ] + }, + { + "cell_type": "markdown", + "id": "345e5da0-10a8-4c65-83ce-072f59c3b376", + "metadata": {}, + "source": [ + "### 详细原理解释\n", + "\n", + "1. **线性变换**：\n", + " ```python\n", + " mixed_x_layer = self.query_key_value(hidden_states)\n", + " ```\n", + " 输入隐藏状态 `hidden_states` 通过线性层映射到查询、键和值的空间。\n", + "\n", + "2. **拆分查询、键和值**：\n", + " ```python\n", + " (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)\n", + " ```\n", + " 将线性变换后的结果拆分成查询、键和值。\n", + "\n", + "3. **形状变换**：\n", + " ```python\n", + " query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]\n", + " ```\n", + " 将查询、键和值的形状从 `[b, sq, np, hn]` 转换为 `[b, np, sq, hn]`。\n", + "\n", + "4. **应用旋转位置编码**：\n", + " ```python\n", + " if rotary_pos_emb is not None:\n", + " query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)\n", + " key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)\n", + " ```\n", + " 如果存在旋转位置编码，则对查询和键应用位置编码。\n", + "\n", + "5. **调整键和值用于推理**：\n", + " ```python\n", + " if kv_cache is not None:\n", + " cache_k, cache_v = kv_cache\n", + " \n", + "\n", + " key_layer = torch.cat((cache_k, key_layer), dim=2)\n", + " value_layer = torch.cat((cache_v, value_layer), dim=2)\n", + " if use_cache:\n", + " if kv_cache is None:\n", + " kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)\n", + " else:\n", + " kv_cache = (key_layer, value_layer)\n", + " else:\n", + " kv_cache = None\n", + " ```\n", + "\n", + "6. **核心注意力计算**：\n", + " ```python\n", + " context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)\n", + " ```\n", + "\n", + "7. **输出层**：\n", + " ```python\n", + " output = self.dense(context_layer)\n", + " ```\n", + "\n", + "通过上述详细的原理解释和公式，可以更好地理解 SelfAttention 类的实现以及其在 Transformer 模型中的作用。自注意力机制通过动态调整不同位置之间的权重，使得模型能够更有效地捕捉全局信息，从而提高模型的性能和泛化能力。" + ] + }, + { + "cell_type": "code", + "execution_count": 69, + "id": "00c11e28-47e7-437b-aab8-1da8b736443a", + "metadata": {}, + "outputs": [], + "source": [ + "def _config_to_kwargs(args):\n", + " common_kwargs = {\n", + " \"dtype\": args.torch_dtype,\n", + " }\n", + " return common_kwargs\n", + "\n", + "\n", + "class MLP(torch.nn.Module):\n", + " \"\"\"多层感知机（MLP）。\n", + "\n", + " MLP 将接受隐藏状态为 h 的输入，将其投影到 4*h 的隐藏维度，执行非线性变换，然后将状态投影回 h 的隐藏维度。\n", + " \"\"\"\n", + "\n", + " def __init__(self, config: ChatGLMConfig, device=None):\n", + " super(MLP, self).__init__()\n", + "\n", + " self.add_bias = config.add_bias_linear\n", + "\n", + " # 投影到 4h。如果使用 swiglu 则将输出宽度加倍，详见 https://arxiv.org/pdf/2002.05202.pdf\n", + " self.dense_h_to_4h = nn.Linear(\n", + " config.hidden_size,\n", + " config.ffn_hidden_size * 2,\n", + " bias=self.add_bias,\n", + " device=device,\n", + " **_config_to_kwargs(config)\n", + " )\n", + "\n", + " def swiglu(x):\n", + " x = torch.chunk(x, 2, dim=-1)\n", + " return F.silu(x[0]) * x[1]\n", + "\n", + " self.activation_func = swiglu\n", + "\n", + " # 投影回 h.\n", + " self.dense_4h_to_h = nn.Linear(\n", + " config.ffn_hidden_size,\n", + " config.hidden_size,\n", + " bias=self.add_bias,\n", + " device=device,\n", + " **_config_to_kwargs(config)\n", + " )\n", + "\n", + " def forward(self, hidden_states):\n", + " # [s, b, 4hp]\n", + " intermediate_parallel = self.dense_h_to_4h(hidden_states)\n", + " intermediate_parallel = self.activation_func(intermediate_parallel)\n", + " # [s, b, h]\n", + " output = self.dense_4h_to_h(intermediate_parallel)\n", + " return output\n", + "\n", + "\n", + "class GLMBlock(torch.nn.Module):\n", + " \"\"\"一个 transformer 层。\n", + "\n", + " Transformer 层接受形状为 [s, b, h] 的输入，并返回相同形状的输出。\n", + " \"\"\"\n", + "\n", + " def __init__(self, config: ChatGLMConfig, layer_number, device=None):\n", + " super(GLMBlock, self).__init__()\n", + " self.layer_number = layer_number\n", + "\n", + " self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm\n", + " self.fp32_residual_connection = config.fp32_residual_connection\n", + "\n", + " LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm\n", + " # 输入数据上的层归一化\n", + " self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,\n", + " dtype=config.torch_dtype)\n", + "\n", + " # 自注意力层\n", + " self.self_attention = SelfAttention(config, layer_number, device=device)\n", + " self.hidden_dropout = config.hidden_dropout\n", + "\n", + " # 注意力输出上的层归一化\n", + " self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,\n", + " dtype=config.torch_dtype)\n", + "\n", + " # 多层感知机（MLP）\n", + " self.mlp = MLP(config, device=device)\n", + "\n", + " def forward(\n", + " self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,\n", + " ):\n", + " # hidden_states: [s, b, h]\n", + "\n", + " # 在 transformer 层开始时的层归一化\n", + " layernorm_output = self.input_layernorm(hidden_states)\n", + " # 自注意力层\n", + " attention_output, kv_cache = self.self_attention(\n", + " layernorm_output,\n", + " attention_mask,\n", + " rotary_pos_emb,\n", + " kv_cache=kv_cache,\n", + " use_cache=use_cache\n", + " )\n", + "\n", + " # 残差连接\n", + " if self.apply_residual_connection_post_layernorm:\n", + " residual = layernorm_output\n", + " else:\n", + " residual = hidden_states\n", + "\n", + " layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)\n", + " layernorm_input = residual + layernorm_input\n", + "\n", + " # 自注意力后的层归一化\n", + " layernorm_output = self.post_attention_layernorm(layernorm_input)\n", + "\n", + " # 多层感知机（MLP）\n", + " mlp_output = self.mlp(layernorm_output)\n", + "\n", + " # 第二个残差连接\n", + " if self.apply_residual_connection_post_layernorm:\n", + " residual = layernorm_output\n", + " else:\n", + " residual = layernorm_input\n", + "\n", + " output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)\n", + " output = residual + output\n", + "\n", + " return output, kv_cache\n", + "\n", + "\n", + "class GLMTransformer(torch.nn.Module):\n", + " \"\"\"Transformer 类。\"\"\"\n", + "\n", + " def __init__(self, config: ChatGLMConfig, device=None):\n", + " super(GLMTransformer, self).__init__()\n", + "\n", + " self.fp32_residual_connection = config.fp32_residual_connection\n", + " self.post_layer_norm = config.post_layer_norm\n", + "\n", + " # 层数\n", + " self.num_layers = config.num_layers\n", + "\n", + " # Transformer 层\n", + " def build_layer(layer_number):\n", + " return GLMBlock(config, layer_number, device=device)\n", + "\n", + " self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])\n", + "\n", + " if self.post_layer_norm:\n", + " LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm\n", + " # 输出前的最终层归一化\n", + " self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,\n", + " dtype=config.torch_dtype)\n", + "\n", + " self.gradient_checkpointing = False\n", + "\n", + " def _get_layer(self, layer_number):\n", + " return self.layers[layer_number]\n", + "\n", + " def forward(\n", + " self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,\n", + " use_cache: Optional[bool] = True,\n", + " output_hidden_states: Optional[bool] = False,\n", + " ):\n", + " if not kv_caches:\n", + " kv_caches = [None for _ in range(self.num_layers)]\n", + " presents = () if use_cache else None\n", + " if self.gradient_checkpointing and self.training:\n", + " if use_cache:\n", + " logger.warning_once(\n", + " \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n", + " )\n", + " use_cache = False\n", + "\n", + " all_self_attentions = None\n", + " all_hidden_states = () if output_hidden_states else None\n", + " for index in range(self.num_layers):\n", + " if output_hidden_states:\n", + " all_hidden_states = all_hidden_states + (hidden_states,)\n", + "\n", + " layer = self._get_layer(index)\n", + " if self.gradient_checkpointing and self.training:\n", + " layer_ret = torch.utils.checkpoint.checkpoint(\n", + " layer,\n", + " hidden_states,\n", + " attention_mask,\n", + " rotary_pos_emb,\n", + " kv_caches[index],\n", + " use_cache,\n", + " use_reentrant=False\n", + " )\n", + " else:\n", + " layer_ret = layer(\n", + " hidden_states,\n", + " attention_mask,\n", + " rotary_pos_emb,\n", + " kv_cache=kv_caches[index],\n", + " use_cache=use_cache\n", + " )\n", + " hidden_states, kv_cache = layer_ret\n", + " if use_cache:\n", + " # token by token 解码，使用元组格式\n", + " if kv_caches[0] is not None:\n", + " presents = presents + (kv_cache,)\n", + " # 预填充解码，使用张量格式以节省 CUDA 内存\n", + " else:\n", + " if len(presents) == 0:\n", + " presents = kv_cache\n", + " else:\n", + " presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)\n", + "\n", + " if output_hidden_states:\n", + " all_hidden_states = all_hidden_states + (hidden_states,)\n", + "\n", + " # 最终层归一化\n", + " if self.post_layer_norm:\n", + " hidden_states = self.final_layernorm(hidden_states)\n", + "\n", + " return hidden_states, presents, all_hidden_states, all_self_attentions\n" + ] + }, + { + "cell_type": "markdown", + "id": "5ad537c5-ceb0-41a7-b889-fc0fe633bab6", + "metadata": {}, + "source": [ + "## 张量形状变化及其意义。\n", + "\n", + "### `CoreAttention` 类中的张量形状\n", + "\n", + "1. **`query_layer`, `key_layer`, `value_layer` 形状变化**:\n", + " ```python\n", + " # 输入形状 [b, np, sq, hn] 和 [b, np, sk, hn]\n", + " # [b, np, sq, hn] -> [b * np, sq, hn]\n", + " query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)\n", + " # [b, np, sk, hn] -> [b * np, sk, hn]\n", + " key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)\n", + " ```\n", + "\n", + "2. **原始注意力得分的计算**:\n", + " ```python\n", + " # 输入形状 [b * np, sq, hn] 和 [b * np, hn, sk]\n", + " # 结果形状 [b * np, sq, sk]\n", + " matmul_result = torch.baddbmm(matmul_input_buffer, query_layer, key_layer.transpose(1, 2), beta=0.0, alpha=(1.0 / self.norm_factor))\n", + " ```\n", + "\n", + "3. **调整视图到原始形状**:\n", + " ```python\n", + " # [b * np, sq, sk] -> [b, np, sq, sk]\n", + " attention_scores = matmul_result.view(*output_size)\n", + " ```\n", + "\n", + "4. **注意力概率形状变化**:\n", + " ```python\n", + " # 输入形状 [b, np, sq, sk]\n", + " # 调整后的形状 [b * np, sq, sk]\n", + " attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)\n", + " ```\n", + "\n", + "5. **计算上下文层**:\n", + " ```python\n", + " # 输入形状 [b * np, sq, sk] 和 [b * np, sk, hn]\n", + " # 结果形状 [b * np, sq, hn]\n", + " context_layer = torch.bmm(attention_probs, value_layer)\n", + " # 调整视图 [b * np, sq, hn] -> [b, np, sq, hn]\n", + " context_layer = context_layer.view(*output_size)\n", + " # [b, np, sq, hn] -> [b, sq, np, hn]\n", + " context_layer = context_layer.transpose(1, 2).contiguous()\n", + " # [b, sq, np, hn] -> [b, sq, hp]\n", + " new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)\n", + " context_layer = context_layer.reshape(*new_context_layer_shape)\n", + " ```\n", + "\n", + "### `SelfAttention` 类中的张量形状\n", + "\n", + "1. **`mixed_x_layer` 形状变化**:\n", + " ```python\n", + " # 输入形状 [b, sq, h]\n", + " # 结果形状 [b, sq, (np * 3 * hn)]\n", + " mixed_x_layer = self.query_key_value(hidden_states)\n", + " ```\n", + "\n", + "2. **多查询注意力的形状变化**:\n", + " ```python\n", + " # [b, sq, (np * 3 * hn)] -> [b, sq, np, hn] 和 [b, sq, np, hn]\n", + " (query_layer, key_layer, value_layer) = mixed_x_layer.split([...], dim=-1)\n", + " # 调整视图\n", + " query_layer = query_layer.view(query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head))\n", + " key_layer = key_layer.view(key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head))\n", + " value_layer = value_layer.view(value_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head))\n", + " ```\n", + "\n", + "3. **普通注意力的形状变化**:\n", + " ```python\n", + " # [b, sq, (np * 3 * hn)] -> [b, sq, np, 3 * hn]\n", + " new_tensor_shape = mixed_x_layer.size()[:-1] + (self.num_attention_heads_per_partition, 3 * self.hidden_size_per_attention_head)\n", + " mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)\n", + " # [b, sq, np, 3 * hn] -> 3 [b, sq, np, hn]\n", + " (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)\n", + " ```\n", + "\n", + "4. **转置操作**:\n", + " ```python\n", + " # [b, sq, np, hn] -> [b, np, sq, hn]\n", + " query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]\n", + " ```\n", + "\n", + "### `MLP` 类中的张量形状\n", + "\n", + "1. **前向传播**:\n", + " ```python\n", + " # 输入形状 [s, b, h]\n", + " # 经过 dense_h_to_4h 线性层后形状 [s, b, 4hp]\n", + " intermediate_parallel = self.dense_h_to_4h(hidden_states)\n", + " intermediate_parallel = self.activation_func(intermediate_parallel)\n", + " # 经过 dense_4h_to_h 线性层后形状 [s, b, h]\n", + " output = self.dense_4h_to_h(intermediate_parallel)\n", + " ```\n", + "\n", + "### `GLMBlock` 类中的张量形状\n", + "\n", + "1. **自注意力层的输入输出形状**:\n", + " ```python\n", + " # 输入形状 [s, b, h]\n", + " # 自注意力层输出形状 [s, b, h]\n", + " attention_output, kv_cache = self.self_attention(layernorm_output, attention_mask, rotary_pos_emb, kv_cache=kv_cache, use_cache=use_cache)\n", + " ```\n", + "\n", + "2. **残差连接**:\n", + " ```python\n", + " # 残差连接，形状保持不变 [s, b, h]\n", + " layernorm_input = residual + layernorm_input\n", + " ```\n", + "\n", + "3. **多层感知机（MLP）的输入输出形状**:\n", + " ```python\n", + " # MLP 输出形状 [s, b, h]\n", + " mlp_output = self.mlp(layernorm_output)\n", + " ```\n", + "\n", + "### `GLMTransformer` 类中的张量形状\n", + "\n", + "1. **逐层处理**:\n", + " ```python\n", + " # 每层的输入输出形状 [s, b, h]\n", + " for index in range(self.num_layers):\n", + " layer = self._get_layer(index)\n", + " hidden_states, kv_cache = layer(hidden_states, attention_mask, rotary_pos_emb, kv_cache=kv_caches[index], use_cache=use_cache)\n", + " ```\n", + "\n", + "2. **最终层归一化**:\n", + " ```python\n", + " # 最终层归一化，形状 [s, b, h]\n", + " if self.post_layer_norm:\n", + " hidden_states = self.final_layernorm(hidden_states)\n", + " ```\n", + "\n", + "总结这些形状变化，有助于理解每个层的输入输出如何传递和处理，确保模型在每个步骤中保持正确的张量形状。" + ] + }, + { + "cell_type": "markdown", + "id": "8b59c956-a31d-459c-a22d-ef3920644ab5", + "metadata": {}, + "source": [ + "### 关键点解释\n", + "\n", + "1. **`CoreAttention` 类**：\n", + " - 实现了核心注意力机制。\n", + " - 支持 PyTorch 2.0 及以上版本的 `scaled_dot_product_attention`，以及旧版本的自定义注意力计算。\n", + "\n", + "2. **`SelfAttention` 类**：\n", + " - 实现了自注意力层。\n", + " - 包括 query、key、value 的计算以及核心注意力机制的应用。\n", + " - 支持多查询注意力机制。\n", + "\n", + "3. **`MLP` 类**：\n", + " - 多层感知机（MLP），包括两个线性层和一个非线性激活函数。\n", + "\n", + "4. **`GLMBlock` 类**：\n", + " - 实现了一个 Transformer 层，包括自注意力层和 MLP 层。\n", + " - 包括层归一化和残差连接。\n", + "\n", + "5. **`GLMTransformer` 类**：\n", + " - 实现了 Transformer 模型，包括多个 Transformer 层。\n", + " - 支持梯度检查点和最终层归一化。" + ] + }, + { + "cell_type": "markdown", + "id": "48414b7e-3bb1-4a06-b411-5492035dabcf", + "metadata": {}, + "source": [ + "### 自注意力机制提取输入隐藏状态的步骤\n", + "\n", + "在 SelfAttention 类中，自注意力机制通过一系列的步骤来提取和处理输入的隐藏状态（hidden states），最终生成上下文向量。这些步骤包括线性变换、计算注意力得分、应用注意力掩码、计算注意力权重和生成上下文向量。以下是详细的步骤和解释：\n", + "\n", + "#### 1. 输入隐藏状态\n", + "\n", + "输入隐藏状态 \\( \\text{hidden_states} \\) 的形状通常为 \\([b, sq, h]\\)，其中：\n", + "- \\( b \\) 是批次大小（batch size）。\n", + "- \\( sq \\) 是序列长度（sequence length）。\n", + "- \\( h \\) 是隐藏层维度（hidden size）。\n", + "\n", + "#### 2. 线性变换\n", + "\n", + "将输入隐藏状态 \\( \\text{hidden_states} \\) 通过线性层投影到查询 \\( Q \\)、键 \\( K \\) 和值 \\( V \\) 空间。这一步的目的是将输入映射到不同的子空间，以便进行注意力计算。公式如下：\n", + "\\begin{align*} Q = \\text{hidden_states} \\cdot W_Q \\end{align*}\n", + "\\begin{align*} K = \\text{hidden_states} \\cdot W_K \\end{align*}\n", + "\\begin{align*} V = \\text{hidden_states} \\cdot W_V \\end{align*}\n", + "\n", + "在代码中，通过以下方式实现：\n", + "```python\n", + "mixed_x_layer = self.query_key_value(hidden_states)\n", + "```\n", + "这里 `self.query_key_value` 是一个线性层，它将输入隐藏状态投影到一个更大的空间，结果的形状为 \\([b, sq, (3 \\times \\text{hidden_size})]\\)。\n", + "\n", + "#### 3. 拆分查询、键和值\n", + "\n", + "将上述结果拆分成查询、键和值向量。拆分后，每个向量的形状为 \\([b, sq, \\text{hidden_size}]\\)。\n", + "```python\n", + "(query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)\n", + "```\n", + "`split_tensor_along_last_dim` 函数按最后一个维度将张量分割成三部分。\n", + "\n", + "#### 4. 形状变换\n", + "\n", + "为了适应后续的矩阵乘法操作，需要调整查询、键和值的形状。将它们从 \\([b, sq, np, hn]\\) 转换为 \\([b, np, sq, hn]\\)，其中 \\( np \\) 是注意力头的数量，\\( hn \\) 是每个注意力头的维度。\n", + "```python\n", + "query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]\n", + "```\n", + "\n", + "#### 5. 应用旋转位置编码\n", + "\n", + "如果存在旋转位置编码（rotary position embedding），则应用到查询和键向量上。旋转位置编码可以帮助模型更好地捕捉位置信息。\n", + "```python\n", + "if rotary_pos_emb is not None:\n", + " query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)\n", + " key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)\n", + "```\n", + "\n", + "#### 6. 调整键和值用于推理\n", + "\n", + "在推理阶段，可能需要缓存键和值，以便在下一个时间步中重复使用，从而提高效率。将缓存的键和值与当前时间步的键和值进行拼接。\n", + "```python\n", + "if kv_cache is not None:\n", + " cache_k, cache_v = kv_cache\n", + " key_layer = torch.cat((cache_k, key_layer), dim=2)\n", + " value_layer = torch.cat((cache_v, value_layer), dim=2)\n", + "if use_cache:\n", + " if kv_cache is None:\n", + " kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)\n", + " else:\n", + " kv_cache = (key_layer, value_layer)\n", + "else:\n", + " kv_cache = None\n", + "```\n", + "\n", + "#### 7. 核心注意力计算\n", + "\n", + "使用 CoreAttention 类来计算注意力得分、注意力权重和上下文向量。\n", + "```python\n", + "context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)\n", + "```\n", + "\n", + "#### 8. 输出层\n", + "\n", + "最后，将上下文向量通过线性层投影回原始的隐藏层维度。这一步是将计算后的结果映射回输入的形状，以便进行后续处理。\n", + "```python\n", + "output = self.dense(context_layer)\n", + "```\n", + "\n", + "### 总结\n", + "\n", + "通过上述步骤，自注意力机制实现了从输入隐藏状态中提取和处理信息的过程。每一步的详细解释如下：\n", + "\n", + "1. **输入隐藏状态**：输入序列的隐藏表示。\n", + "2. **线性变换**：将隐藏表示投影到查询、键和值的空间。\n", + "3. **拆分查询、键和值**：将投影后的结果拆分成查询、键和值向量。\n", + "4. **形状变换**：调整查询、键和值的形状，以适应后续操作。\n", + "5. **应用旋转位置编码**：增强查询和键向量的位置信息。\n", + "6. **调整键和值用于推理**：在推理阶段缓存键和值，以提高效率。\n", + "7. **核心注意力计算**：通过计算注意力得分、权重和上下文向量，完成注意力机制的核心部分。\n", + "8. **输出层**：将上下文向量投影回原始的隐藏层维度。\n", + "\n", + "这些步骤共同构成了自注意力机制，从而使模型能够动态地关注输入序列中的不同部分，捕捉全局信息并生成更丰富的表示。" + ] + }, + { + "cell_type": "markdown", + "id": "3d16f0c2-db97-4512-9a09-1e59502d6e19", + "metadata": {}, + "source": [ + "### ChatGLMPreTrainedModel 类" + ] + }, + { + "cell_type": "code", + "execution_count": 70, + "id": "9d1e5174-cc11-4ba0-ad18-b58110bf084b", + "metadata": {}, + "outputs": [], + "source": [ + "class ChatGLMPreTrainedModel(PreTrainedModel):\n", + " \"\"\"\n", + " 处理权重初始化和下载及加载预训练模型的简单接口的抽象类。\n", + " \"\"\"\n", + " is_parallelizable = False # 是否可并行化\n", + " supports_gradient_checkpointing = True # 是否支持梯度检查点\n", + " config_class = ChatGLMConfig # 配置类\n", + " base_model_prefix = \"transformer\" # 基础模型前缀\n", + " _no_split_modules = [\"GLMBlock\"] # 不拆分的模块列表\n", + "\n", + " def _init_weights(self, module: nn.Module):\n", + " \"\"\"初始化权重\"\"\"\n", + " return\n", + "\n", + " def get_masks(self, input_ids, past_key_values, padding_mask=None):\n", + " \"\"\"\n", + " 获取注意力掩码\n", + " \"\"\"\n", + " batch_size, seq_length = input_ids.shape\n", + " # 创建下三角矩阵作为全注意力掩码\n", + " full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)\n", + " full_attention_mask.tril_()\n", + " past_length = 0\n", + " if past_key_values:\n", + " past_length = past_key_values[0][0].shape[2]\n", + " if past_length:\n", + " full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length, device=input_ids.device), full_attention_mask), dim=-1)\n", + " if padding_mask is not None:\n", + " full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)\n", + " if not past_length and padding_mask is not None:\n", + " full_attention_mask -= padding_mask.unsqueeze(-1) - 1\n", + " full_attention_mask = (full_attention_mask < 0.5).bool()\n", + " full_attention_mask.unsqueeze_(1)\n", + " return full_attention_mask\n", + "\n", + " def get_position_ids(self, input_ids, device):\n", + " \"\"\"\n", + " 获取位置ID\n", + " \"\"\"\n", + " batch_size, seq_length = input_ids.shape\n", + " position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)\n", + " return position_ids\n", + "\n", + " def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):\n", + " if not self.supports_gradient_checkpointing:\n", + " raise ValueError(f\"{self.__class__.__name__} does not support gradient checkpointing.\")" + ] + }, + { + "cell_type": "markdown", + "id": "16d19827-5201-4981-9618-520275a14b6c", + "metadata": {}, + "source": [ + "### Embedding 类" + ] + }, + { + "cell_type": "code", + "execution_count": 71, + "id": "ede82883-45bb-41f8-b679-a97b2d04c87e", + "metadata": {}, + "outputs": [], + "source": [ + "class Embedding(torch.nn.Module):\n", + " \"\"\"语言模型嵌入层。\"\"\"\n", + "\n", + " def __init__(self, config: ChatGLMConfig, device=None):\n", + " super(Embedding, self).__init__()\n", + "\n", + " self.hidden_size = config.hidden_size\n", + " # 词嵌入层（并行）\n", + " self.word_embeddings = nn.Embedding(\n", + " config.padded_vocab_size,\n", + " self.hidden_size,\n", + " dtype=config.torch_dtype,\n", + " device=device\n", + " )\n", + " self.fp32_residual_connection = config.fp32_residual_connection\n", + "\n", + " def forward(self, input_ids):\n", + " # 获取词嵌入\n", + " words_embeddings = self.word_embeddings(input_ids)\n", + " embeddings = words_embeddings\n", + " # 如果设置了fp32残差连接，则转换为浮点数\n", + " if self.fp32_residual_connection:\n", + " embeddings = embeddings.float()\n", + " return embeddings" + ] + }, + { + "cell_type": "markdown", + "id": "20e7568a-a42f-4053-9939-152a656181e4", + "metadata": {}, + "source": [ + "### ChatGLMModel 类" + ] + }, + { + "cell_type": "code", + "execution_count": 72, + "id": "21ed53de-5687-4093-9662-efbc224093b4", + "metadata": {}, + "outputs": [], + "source": [ + "class ChatGLMModel(ChatGLMPreTrainedModel):\n", + " def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):\n", + " super().__init__(config)\n", + " if empty_init:\n", + " init_method = skip_init # 跳过初始化\n", + " else:\n", + " init_method = default_init # 使用默认初始化方法\n", + " init_kwargs = {}\n", + " if device is not None:\n", + " init_kwargs[\"device\"] = device\n", + " self.embedding = init_method(Embedding, config, **init_kwargs) # 使用 Embedding 类\n", + " self.num_layers = config.num_layers\n", + " self.multi_query_group_num = config.multi_query_group_num\n", + " self.kv_channels = config.kv_channels\n", + "\n", + " # 旋转位置嵌入\n", + " self.seq_length = config.seq_length\n", + " rotary_dim = (\n", + " config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels\n", + " )\n", + "\n", + " self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio, original_impl=config.original_rope, \n", + " device=device, dtype=config.torch_dtype) # 使用 RotaryEmbedding 类\n", + " self.encoder = init_method(GLMTransformer, config, **init_kwargs) # 使用 GLMTransformer 类\n", + " self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,\n", + " dtype=config.torch_dtype, **init_kwargs) # 使用 nn.Linear 类\n", + "\n", + " def get_input_embeddings(self):\n", + " return self.embedding.word_embeddings\n", + "\n", + " def set_input_embeddings(self, value):\n", + " self.embedding.word_embeddings = value\n", + "\n", + " def forward(\n", + " self,\n", + " input_ids,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.BoolTensor] = None,\n", + " full_attention_mask: Optional[torch.BoolTensor] = None,\n", + " past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,\n", + " inputs_embeds: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " output_hidden_states: Optional[bool] = None,\n", + " return_dict: Optional[bool] = None,\n", + " ):\n", + " output_hidden_states = (\n", + " output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n", + " )\n", + " use_cache = use_cache if use_cache is not None else self.config.use_cache\n", + " return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n", + "\n", + " batch_size, seq_length = input_ids.shape\n", + "\n", + " if inputs_embeds is None:\n", + " inputs_embeds = self.embedding(input_ids) # 使用 Embedding 类\n", + "\n", + " if full_attention_mask is None:\n", + " if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):\n", + " full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)\n", + "\n", + " # 旋转位置嵌入\n", + " rotary_pos_emb = self.rotary_pos_emb(self.seq_length) # 使用 RotaryEmbedding 类\n", + " if position_ids is not None:\n", + " rotary_pos_emb = rotary_pos_emb[position_ids]\n", + " else:\n", + " rotary_pos_emb = rotary_pos_emb[None, :seq_length]\n", + "\n", + " # 运行编码器\n", + " hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(\n", + " inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,\n", + " kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states\n", + " ) # 使用 GLMTransformer 类\n", + " if presents is not None and type(presents) is torch.Tensor:\n", + " presents = presents.split(1, dim=0)\n", + " presents = list(presents)\n", + " presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]\n", + " presents = [tuple([x.squeeze(0) for x in y]) for y in presents]\n", + " presents = tuple(presents)\n", + "\n", + " if not return_dict:\n", + " return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)\n", + "\n", + " return BaseModelOutputWithPast(\n", + " last_hidden_state=hidden_states,\n", + " past_key_values=presents,\n", + " hidden_states=all_hidden_states,\n", + " attentions=all_self_attentions,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "48b86c96-a071-49aa-8676-f1d8be23a31f", + "metadata": {}, + "source": [ + "### 总结\n", + "\n", + "1. **ChatGLMPreTrainedModel**:\n", + " - **get_masks**：生成注意力掩码。\n", + " - **get_position_ids**：生成位置 ID。\n", + "\n", + "2. **Embedding**:\n", + " - **word_embeddings**：实现词嵌入。\n", + "\n", + "3. **ChatGLMModel**:\n", + " - **RotaryEmbedding**：用于位置编码。\n", + " - **GLMTransformer**：实现 Transformer 编码器。\n", + " - **forward**：执行前向传播，集成所有模块。\n", + "\n", + "4. **使用的先前构建的模块**：\n", + " - **RotaryEmbedding**：用于生成旋转位置嵌入。\n", + " - **GLMTransformer**：用于编码器部分。\n", + " - **Embedding**：用于生成词嵌入。\n", + " - **CoreAttention**、**SelfAttention**：间接通过 GLMTransformer 使用。" + ] + }, + { + "cell_type": "markdown", + "id": "4af49b32-bd82-4b62-9d95-494ae2d78dfb", + "metadata": {}, + "source": [ + "### ChatGLMForConditionalGeneration 类" + ] + }, + { + "cell_type": "code", + "execution_count": 73, + "id": "4725aadc-c0e6-4d42-bebb-a4194dd695f6", + "metadata": {}, + "outputs": [], + "source": [ + "class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):\n", + " def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):\n", + " super().__init__(config)\n", + "\n", + " self.max_sequence_length = config.max_length # 最大序列长度\n", + " self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) # 使用 ChatGLMModel 类\n", + " self.config = config\n", + "\n", + " def _update_model_kwargs_for_generation(\n", + " self,\n", + " outputs: ModelOutput,\n", + " model_kwargs: Dict[str, Any],\n", + " is_encoder_decoder: bool = False,\n", + " standardize_cache_format: bool = False,\n", + " ) -> Dict[str, Any]:\n", + " # 更新 past_key_values\n", + " model_kwargs[\"past_key_values\"] = self._extract_past_from_model_output(\n", + " outputs, standardize_cache_format=standardize_cache_format\n", + " )\n", + "\n", + " # 更新注意力掩码\n", + " if \"attention_mask\" in model_kwargs:\n", + " attention_mask = model_kwargs[\"attention_mask\"]\n", + " model_kwargs[\"attention_mask\"] = torch.cat(\n", + " [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1\n", + " )\n", + "\n", + " # 更新位置 ids\n", + " if \"position_ids\" in model_kwargs:\n", + " position_ids = model_kwargs[\"position_ids\"]\n", + " new_position_id = position_ids[..., -1:].clone()\n", + " new_position_id += 1\n", + " model_kwargs[\"position_ids\"] = torch.cat(\n", + " [position_ids, new_position_id], dim=-1\n", + " )\n", + "\n", + " model_kwargs[\"is_first_forward\"] = False\n", + " return model_kwargs\n", + "\n", + " def prepare_inputs_for_generation(\n", + " self,\n", + " input_ids: torch.LongTensor,\n", + " past_key_values: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " is_first_forward: bool = True,\n", + " **kwargs\n", + " ) -> dict:\n", + " # 如果 past_key_values 不为空，只取 input_ids 的最后一个 token\n", + " if position_ids is None:\n", + " position_ids = self.get_position_ids(input_ids, device=input_ids.device)\n", + " if not is_first_forward:\n", + " if past_key_values is not None:\n", + " position_ids = position_ids[..., -1:]\n", + " input_ids = input_ids[:, -1:]\n", + " return {\n", + " \"input_ids\": input_ids,\n", + " \"past_key_values\": past_key_values,\n", + " \"position_ids\": position_ids,\n", + " \"attention_mask\": attention_mask,\n", + " \"return_last_logit\": True,\n", + " \"use_cache\": use_cache\n", + " }\n", + "\n", + " def forward(\n", + " self,\n", + " input_ids: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " past_key_values: Optional[Tuple[torch.FloatTensor]] = None,\n", + " inputs_embeds: Optional[torch.Tensor] = None,\n", + " labels: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " output_attentions: Optional[bool] = None,\n", + " output_hidden_states: Optional[bool] = None,\n", + " return_dict: Optional[bool] = None,\n", + " return_last_logit: Optional[bool] = False,\n", + " ):\n", + " use_cache = use_cache if use_cache is not None else self.config.use_cache\n", + " return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n", + "\n", + " transformer_outputs = self.transformer(\n", + " input_ids=input_ids,\n", + " position_ids=position_ids,\n", + " attention_mask=attention_mask,\n", + " past_key_values=past_key_values,\n", + " inputs_embeds=inputs_embeds,\n", + " use_cache=use_cache,\n", + " output_hidden_states=output_hidden_states,\n", + " return_dict=return_dict,\n", + " ) # 使用 ChatGLMModel 类\n", + "\n", + " hidden_states = transformer_outputs[0]\n", + " if return_last_logit:\n", + " hidden_states = hidden_states[:, -1:]\n", + " lm_logits = self.transformer.output_layer(hidden_states)\n", + "\n", + " loss = None\n", + " if labels is not None:\n", + " lm_logits = lm_logits.to(torch.float32)\n", + "\n", + " # Shift so that tokens < n predict n\n", + " shift_logits = lm_logits[..., :-1, :].contiguous()\n", + " shift_labels = labels[..., 1:].contiguous()\n", + " # Flatten the tokens\n", + " loss_fct = CrossEntropyLoss(ignore_index=-100)\n", + " loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n", + "\n", + " lm_logits = lm_logits.to(hidden_states.dtype)\n", + " loss = loss.to(hidden_states.dtype)\n", + "\n", + " if not return_dict:\n", + " output = (lm_logits,) + transformer_outputs[1:]\n", + " return ((loss,) + output) if loss is not None else output\n", + "\n", + " return CausalLMOutputWithPast(\n", + " loss=loss,\n", + " logits=lm_logits,\n", + " past_key_values=transformer_outputs.past_key_values,\n", + " hidden_states=transformer_outputs.hidden_states,\n", + " attentions=transformer_outputs.attentions,\n", + " )\n", + "\n", + " @staticmethod\n", + " def _reorder_cache(\n", + " past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor\n", + " ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:\n", + " \"\"\"\n", + " 重新排序 `past_key_values` 缓存以匹配每个生成步骤中的 `beam_idx`。\n", + " \"\"\"\n", + " return tuple(\n", + " (\n", + " layer_past[0].index_select(0, beam_idx.to(layer_past[0].device)),\n", + " layer_past[1].index_select(0, beam_idx.to(layer_past[1].device)),\n", + " )\n", + " for layer_past in past\n", + " )\n", + "\n", + " def process_response(self, output, history):\n", + " content = \"\"\n", + " history = deepcopy(history)\n", + " for response in output.split(\"\"):\n", + " if \"\\n\" in response:\n", + " metadata, content = response.split(\"\\n\", maxsplit=1)\n", + " else:\n", + " metadata, content = \"\", response\n", + " if not metadata.strip():\n", + " content = content.strip()\n", + " history.append({\"role\": \"assistant\", \"metadata\": metadata, \"content\": content})\n", + " content = content.replace(\"[[训练时间]]\", \"2023年\")\n", + " else:\n", + " history.append({\"role\": \"assistant\", \"metadata\": metadata, \"content\": content})\n", + " if history[0][\"role\"] == \"system\" and \"tools\" in history[0]:\n", + " parameters = json.loads(content)\n", + " content = {\"name\": metadata.strip(), \"parameters\": parameters}\n", + " else:\n", + " content = {\"name\": metadata.strip(), \"content\": content}\n", + " return content, history\n", + "\n", + " @torch.inference_mode()\n", + " def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = \"user\",\n", + " max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,\n", + " **kwargs):\n", + " if history is None:\n", + " history = []\n", + " if logits_processor is None:\n", + " logits_processor = LogitsProcessorList()\n", + " logits_processor.append(InvalidScoreLogitsProcessor())\n", + " gen_kwargs = {\"max_length\": max_length, \"num_beams\": num_beams, \"do_sample\": do_sample, \"top_p\": top_p,\n", + " \"temperature\": temperature, \"logits_processor\": logits_processor, **kwargs}\n", + " history.append({\"role\": role, \"content\": query})\n", + " inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,\n", + " return_tensors=\"pt\", return_dict=True)\n", + " inputs = inputs.to(self.device)\n", + " eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(\"\"),\n", + " tokenizer.convert_tokens_to_ids(\"\")]\n", + " outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)\n", + " outputs = outputs.tolist()[0][len(inputs[\"input_ids\"][0]):-1]\n", + " response = tokenizer.decode(outputs)\n", + " response, history = self.process_response(response, history)\n", + " return response, history\n", + "\n", + " @torch.inference_mode()\n", + " def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = \"user\",\n", + " past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,\n", + " logits_processor=None, return_past_key_values=False, **kwargs):\n", + " if history is None:\n", + " history = []\n", + " if logits_processor is None:\n", + " logits_processor = LogitsProcessorList()\n", + " logits_processor.append(InvalidScoreLogitsProcessor())\n", + " eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(\"\"),\n", + " tokenizer.convert_tokens_to_ids(\"\")]\n", + " gen_kwargs = {\"max_length\": max_length, \"do_sample\": do_sample, \"top_p\": top_p,\n", + " \"temperature\": temperature, \"logits_processor\": logits_processor, **kwargs}\n", + " if past_key_values is None:\n", + " inputs = tokenizer.apply_chat_template(history + [{\"role\": role, \"content\": query}],\n", + " add_generation_prompt=True, tokenize=True, return_tensors=\"pt\",\n", + " return_dict=True)\n", + " else:\n", + " inputs = tokenizer.apply_chat_template([{\"role\": role, \"content\": query}], add_special_tokens=False,\n", + " add_generation_prompt=True, tokenize=True, return_tensors=\"pt\",\n", + " return_dict=True)\n", + " inputs = inputs.to(self.device)\n", + " if past_key_values is not None:\n", + " past_length = past_key_values[0][0].shape[2]\n", + " inputs.position_ids += past_length\n", + "\n", + "\n", + " attention_mask = inputs.attention_mask\n", + " attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)\n", + " inputs['attention_mask'] = attention_mask\n", + " history.append({\"role\": role, \"content\": query})\n", + " for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,\n", + " eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,\n", + " **gen_kwargs):\n", + " if return_past_key_values:\n", + " outputs, past_key_values = outputs\n", + " outputs = outputs.tolist()[0][len(inputs[\"input_ids\"][0]):-1]\n", + " response = tokenizer.decode(outputs)\n", + " if response and response[-1] != \"�\":\n", + " response, new_history = self.process_response(response, history)\n", + " if return_past_key_values:\n", + " yield response, new_history, past_key_values\n", + " else:\n", + " yield response, new_history\n", + "\n", + " @torch.inference_mode()\n", + " def stream_generate(\n", + " self,\n", + " input_ids,\n", + " generation_config: Optional[GenerationConfig] = None,\n", + " logits_processor: Optional[LogitsProcessorList] = None,\n", + " stopping_criteria: Optional[StoppingCriteriaList] = None,\n", + " prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,\n", + " return_past_key_values=False,\n", + " **kwargs,\n", + " ):\n", + " batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]\n", + "\n", + " if generation_config is None:\n", + " generation_config = self.generation_config\n", + " generation_config = copy.deepcopy(generation_config)\n", + " model_kwargs = generation_config.update(**kwargs)\n", + " model_kwargs[\"use_cache\"] = generation_config.use_cache\n", + " bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id\n", + "\n", + " if isinstance(eos_token_id, int):\n", + " eos_token_id = [eos_token_id]\n", + " eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None\n", + "\n", + " has_default_max_length = kwargs.get(\"max_length\") is None and generation_config.max_length is not None\n", + " if has_default_max_length and generation_config.max_new_tokens is None:\n", + " warnings.warn(\n", + " f\"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. \"\n", + " \"This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we\"\n", + " \" recommend using `max_new_tokens` to control the maximum length of the generation.\",\n", + " UserWarning,\n", + " )\n", + " elif generation_config.max_new_tokens is not None:\n", + " generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length\n", + " if not has_default_max_length:\n", + " logger.warn(\n", + " f\"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=\"\n", + " f\"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. \"\n", + " \"Please refer to the documentation for more information. \"\n", + " \"(https://hf-mirror.com/docs/transformers/main/en/main_classes/text_generation)\",\n", + " UserWarning,\n", + " )\n", + "\n", + " if input_ids_seq_length >= generation_config.max_length:\n", + " input_ids_string = \"decoder_input_ids\" if self.config.is_encoder_decoder else \"input_ids\"\n", + " logger.warning(\n", + " f\"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to\"\n", + " f\" {generation_config.max_length}. This can lead to unexpected behavior. You should consider\"\n", + " \" increasing `max_new_tokens`.\"\n", + " )\n", + "\n", + " # 2. Set generation parameters if not already defined\n", + " logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()\n", + " stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()\n", + "\n", + " logits_processor = self._get_logits_processor(\n", + " generation_config=generation_config,\n", + " input_ids_seq_length=input_ids_seq_length,\n", + " encoder_input_ids=input_ids,\n", + " prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,\n", + " logits_processor=logits_processor,\n", + " )\n", + "\n", + " stopping_criteria = self._get_stopping_criteria(\n", + " generation_config=generation_config, stopping_criteria=stopping_criteria\n", + " )\n", + " logits_warper = self._get_logits_warper(generation_config)\n", + "\n", + " unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)\n", + " scores = None\n", + " while True:\n", + " model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)\n", + " # 前向传递获取下一个 token\n", + " outputs = self(\n", + " **model_inputs,\n", + " return_dict=True,\n", + " output_attentions=False,\n", + " output_hidden_states=False,\n", + " )\n", + "\n", + " next_token_logits = outputs.logits[:, -1, :]\n", + "\n", + " # 预处理分布\n", + " next_token_scores = logits_processor(input_ids, next_token_logits)\n", + " next_token_scores = logits_warper(input_ids, next_token_scores)\n", + "\n", + " # 采样\n", + " probs = nn.functional.softmax(next_token_scores, dim=-1)\n", + " if generation_config.do_sample:\n", + " next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)\n", + " else:\n", + " next_tokens = torch.argmax(probs, dim=-1)\n", + " # 更新生成的 ids、模型输入和下一个步骤的长度\n", + " input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)\n", + " model_kwargs = self._update_model_kwargs_for_generation(\n", + " outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder\n", + " )\n", + " unfinished_sequences = unfinished_sequences.mul(\n", + " next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)\n", + " )\n", + " if return_past_key_values:\n", + " yield input_ids, outputs.past_key_values\n", + " else:\n", + " yield input_ids\n", + " # 当每个句子完成时或超出最大长度时停止\n", + " if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):\n", + " break" + ] + }, + { + "cell_type": "markdown", + "id": "20a7db6f-8889-4e73-b882-edf7e4c72ca7", + "metadata": {}, + "source": [ + "### ChatGLMForSequenceClassification 类" + ] + }, + { + "cell_type": "code", + "execution_count": 74, + "id": "832d1a09-3f8c-4035-ad73-ec3433a71f50", + "metadata": {}, + "outputs": [], + "source": [ + "class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):\n", + " def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):\n", + " super().__init__(config)\n", + "\n", + " self.num_labels = config.num_labels # 标签数量\n", + " self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) # 使用 ChatGLMModel 类\n", + "\n", + " self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half)\n", + " if config.classifier_dropout is not None:\n", + " self.dropout = nn.Dropout(config.classifier_dropout)\n", + " else:\n", + " self.dropout = None\n", + " self.config = config\n", + "\n", + " def forward(\n", + " self,\n", + " input_ids: Optional[torch.LongTensor] = None,\n", + " position_ids: Optional[torch.LongTensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " full_attention_mask: Optional[torch.Tensor] = None,\n", + " past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,\n", + " inputs_embeds: Optional[torch.LongTensor] = None,\n", + " labels: Optional[torch.LongTensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " output_hidden_states: Optional[bool] = None,\n", + " return_dict: Optional[bool] = None,\n", + " ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:\n", + " return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n", + "\n", + " transformer_outputs = self.transformer(\n", + " input_ids=input_ids,\n", + " position_ids=position_ids,\n", + " attention_mask=attention_mask,\n", + " full_attention_mask=full_attention_mask,\n", + " past_key_values=past_key_values,\n", + " inputs_embeds=inputs_embeds,\n", + " use_cache=use_cache,\n", + " output_hidden_states=output_hidden_states,\n", + " return_dict=return_dict,\n", + " ) # 使用 ChatGLMModel 类\n", + "\n", + " hidden_states = transformer_outputs[0]\n", + " pooled_hidden_states = hidden_states[:, -1]\n", + " if self.dropout is not None:\n", + " pooled_hidden_states = self.dropout(pooled_hidden_states)\n", + " logits = self.classifier_head(pooled_hidden_states)\n", + "\n", + " loss = None\n", + " if labels is not None:\n", + " if self.config.problem_type is None:\n", + " if self.num_labels == 1:\n", + " self.config.problem_type = \"regression\"\n", + " elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):\n", + " self.config.problem_type = \"single_label_classification\"\n", + " else:\n", + " self.config.problem_type = \"multi_label_classification\"\n", + "\n", + " if self.config.problem_type == \"regression\":\n", + " loss_fct = MSELoss()\n", + " if self.num_labels == 1:\n", + " loss = loss_fct(logits.squeeze().float(), labels.squeeze())\n", + " else:\n", + " loss = loss_fct(logits.float(), labels)\n", + " elif self.config.problem_type == \"single_label_classification\":\n", + " loss_fct = CrossEntropyLoss()\n", + " loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1))\n", + " elif self.config.problem_type == \"multi_label_classification\":\n", + " loss_fct = BCEWithLogitsLoss()\n", + " loss = loss_fct(logits.float(), labels.view(-1, self.num_labels))\n", + "\n", + " if not return_dict:\n", + " output = (logits,) + transformer_outputs[1:]\n", + " return ((loss,) + output) if loss is not None else output\n", + "\n", + " return SequenceClassifierOutputWithPast(\n", + " loss\n", + "\n", + "=loss,\n", + " logits=logits,\n", + " past_key_values=transformer_outputs.past_key_values,\n", + " hidden_states=transformer_outputs.hidden_states,\n", + " attentions=transformer_outputs.attentions,\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "id": "415020b6-a65c-4738-aaac-c25b00c84996", + "metadata": {}, + "source": [ + "### 使用的先前构建的模块\n", + "\n", + "在这段代码中，多个模块和方法是基于之前构建的类和函数的：\n", + "\n", + "1. **ChatGLMModel**：\n", + " - **用于 ChatGLMForConditionalGeneration 和 ChatGLMForSequenceClassification 中**，作为 Transformer 模型的核心部分。\n", + "\n", + "2. **ChatGLMPreTrainedModel**：\n", + " - **作为 ChatGLMForConditionalGeneration 和 ChatGLMForSequenceClassification 的基类**，提供权重初始化和加载预训练模型的接口。\n", + "\n", + "3. **RotaryEmbedding**：\n", + " - **在 ChatGLMModel 中用于位置编码**。\n", + "\n", + "4. **CoreAttention 和 SelfAttention**：\n", + " - **在 GLMTransformer 中使用**，实现了注意力机制的核心部分。\n", + "\n", + "通过详细注释和说明，可以更好地理解代码的构建和实现原理。这些模块共同构成了 ChatGLM 模型的整体架构，实现了条件生成和序列分类的功能。" + ] + }, + { + "cell_type": "markdown", + "id": "d43e2c5d-40de-4c9d-a8b7-e49099306fa7", + "metadata": {}, + "source": [ + "### `past_key_values` 变量的含义\n", + "\n", + "在 Transformer 模型中，特别是用于生成任务的模型，如 GPT 类模型中，`past_key_values` 是一个非常重要的变量。它用于缓存模型在前一个时间步计算得到的键（key）和值（value）向量。这些缓存的数据可以在后续的时间步中重复使用，从而提高计算效率，尤其是在长序列生成任务中。\n", + "\n", + "### `past_key_values` 的作用\n", + "\n", + "1. **缓存先前计算结果**：\n", + " 在生成文本的过程中，每一步生成一个新的词，这时需要将当前时间步的查询向量（query）与所有先前时间步的键和值向量进行计算。如果每次都重新计算所有的键和值，将会非常低效。`past_key_values` 缓存了这些先前时间步的结果，避免了重复计算。\n", + "\n", + "2. **加速生成过程**：\n", + " 在长序列生成中，通过缓存先前时间步的键和值向量，只需要对当前时间步进行计算并与缓存结果结合，大大加速了生成过程。\n", + "\n", + "### 具体实现中的 `past_key_values`\n", + "\n", + "#### 在 `ChatGLMForConditionalGeneration` 类中的使用\n", + "\n", + "```python\n", + "def prepare_inputs_for_generation(\n", + " self,\n", + " input_ids: torch.LongTensor,\n", + " past_key_values: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " is_first_forward: bool = True,\n", + " **kwargs\n", + ") -> dict:\n", + " # 如果 past_key_values 不为空，只取 input_ids 的最后一个 token\n", + " if position_ids is None:\n", + " position_ids = self.get_position_ids(input_ids, device=input_ids.device)\n", + " if not is_first_forward:\n", + " if past_key_values is not None:\n", + " position_ids = position_ids[..., -1:]\n", + " input_ids = input_ids[:, -1:]\n", + " return {\n", + " \"input_ids\": input_ids,\n", + " \"past_key_values\": past_key_values,\n", + " \"position_ids\": position_ids,\n", + " \"attention_mask\": attention_mask,\n", + " \"return_last_logit\": True,\n", + " \"use_cache\": use_cache\n", + " }\n", + "```\n", + "\n", + "在 `prepare_inputs_for_generation` 方法中，如果 `past_key_values` 不为空，只会取 `input_ids` 的最后一个 token。这样做的目的是为了在生成新 token 时，只计算当前时间步的数据，而不需要重新计算整个序列。\n", + "\n", + "#### 在 `ChatGLMForConditionalGeneration` 类的 `forward` 方法中\n", + "\n", + "```python\n", + "def forward(\n", + " self,\n", + " input_ids: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " past_key_values: Optional[Tuple[torch.FloatTensor]] = None,\n", + " inputs_embeds: Optional[torch.Tensor] = None,\n", + " labels: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " output_attentions: Optional[bool] = None,\n", + " output_hidden_states: Optional[bool] = None,\n", + " return_dict: Optional[bool] = None,\n", + " return_last_logit: Optional[bool] = False,\n", + "):\n", + " use_cache = use_cache if use_cache is not None else self.config.use_cache\n", + " return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n", + "\n", + " transformer_outputs = self.transformer(\n", + " input_ids=input_ids,\n", + " position_ids=position_ids,\n", + " attention_mask=attention_mask,\n", + " past_key_values=past_key_values,\n", + " inputs_embeds=inputs_embeds,\n", + " use_cache=use_cache,\n", + " output_hidden_states=output_hidden_states,\n", + " return_dict=return_dict,\n", + " ) # 使用 ChatGLMModel 类\n", + "\n", + " hidden_states = transformer_outputs[0]\n", + " if return_last_logit:\n", + " hidden_states = hidden_states[:, -1:]\n", + " lm_logits = self.transformer.output_layer(hidden_states)\n", + "\n", + " loss = None\n", + " if labels is not None:\n", + " lm_logits = lm_logits.to(torch.float32)\n", + "\n", + " # Shift so that tokens < n predict n\n", + " shift_logits = lm_logits[..., :-1, :].contiguous()\n", + " shift_labels = labels[..., 1:].contiguous()\n", + " # Flatten the tokens\n", + " loss_fct = CrossEntropyLoss(ignore_index=-100)\n", + " loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n", + "\n", + " lm_logits = lm_logits.to(hidden_states.dtype)\n", + " loss = loss.to(hidden_states.dtype)\n", + "\n", + " if not return_dict:\n", + " output = (lm_logits,) + transformer_outputs[1:]\n", + " return ((loss,) + output) if loss is not None else output\n", + "\n", + " return CausalLMOutputWithPast(\n", + " loss=loss,\n", + " logits=lm_logits,\n", + " past_key_values=transformer_outputs.past_key_values,\n", + " hidden_states=transformer_outputs.hidden_states,\n", + " attentions=transformer_outputs.attentions,\n", + " )\n", + "```\n", + "\n", + "在 `forward` 方法中，`past_key_values` 作为参数传递给 `transformer` 模型。`transformer` 模型内部会使用这些缓存的键和值向量来加速计算。\n", + "`past_key_values` 是一种用于加速 Transformer 模型在生成任务中的缓存机制。它保存了前一个时间步计算得到的键和值向量，避免了在每个时间步中重复计算这些向量，从而提高了生成过程的效率。通过使用 `past_key_values`，模型可以更快地生成长序列数据，这在实际应用中是非常重要的。" + ] + }, + { + "cell_type": "markdown", + "id": "e0e361aa-a1b5-4eae-9178-306b888713f4", + "metadata": {}, + "source": [ + "### 函数分析及其区别和联系\n", + "\n", + "在 `ChatGLMForConditionalGeneration` 类中，有三个与生成有关的重要函数：`_update_model_kwargs_for_generation`、`prepare_inputs_for_generation` 和 `forward`。它们的作用、区别和联系如下：\n", + "\n", + "#### 1. `_update_model_kwargs_for_generation` 函数\n", + "\n", + "```python\n", + "def _update_model_kwargs_for_generation(\n", + " self,\n", + " outputs: ModelOutput,\n", + " model_kwargs: Dict[str, Any],\n", + " is_encoder_decoder: bool = False,\n", + " standardize_cache_format: bool = False,\n", + ") -> Dict[str, Any]:\n", + " # 更新 past_key_values\n", + " model_kwargs[\"past_key_values\"] = self._extract_past_from_model_output(\n", + " outputs, standardize_cache_format=standardize_cache_format\n", + " )\n", + "\n", + " # 更新注意力掩码\n", + " if \"attention_mask\" in model_kwargs:\n", + " attention_mask = model_kwargs[\"attention_mask\"]\n", + " model_kwargs[\"attention_mask\"] = torch.cat(\n", + " [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1\n", + " )\n", + "\n", + " # 更新位置 ids\n", + " if \"position_ids\" in model_kwargs:\n", + " position_ids = model_kwargs[\"position_ids\"]\n", + " new_position_id = position_ids[..., -1:].clone()\n", + " new_position_id += 1\n", + " model_kwargs[\"position_ids\"] = torch.cat(\n", + " [position_ids, new_position_id], dim=-1\n", + " )\n", + "\n", + " model_kwargs[\"is_first_forward\"] = False\n", + " return model_kwargs\n", + "```\n", + "\n", + "- **作用**：更新生成过程中所需的模型参数。具体包括：\n", + " - 更新 `past_key_values` 以缓存先前计算的键和值向量。\n", + " - 更新 `attention_mask` 以包括新的生成的 token。\n", + " - 更新 `position_ids` 以增加新的位置 ID。\n", + "- **区别**：该函数不直接进行前向传播，而是更新模型参数，为下一步的生成做准备。\n", + "- **联系**：该函数在每一步生成新 token 后调用，用于更新模型参数，为下一步的生成做准备。\n", + "\n", + "#### 2. `prepare_inputs_for_generation` 函数\n", + "\n", + "```python\n", + "def prepare_inputs_for_generation(\n", + " self,\n", + " input_ids: torch.LongTensor,\n", + " past_key_values: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " is_first_forward: bool = True,\n", + " **kwargs\n", + ") -> dict:\n", + " # 如果 past_key_values 不为空，只取 input_ids 的最后一个 token\n", + " if position_ids is None:\n", + " position_ids = self.get_position_ids(input_ids, device=input_ids.device)\n", + " if not is_first_forward:\n", + " if past_key_values is not None:\n", + " position_ids = position_ids[..., -1:]\n", + " input_ids = input_ids[:, -1:]\n", + " return {\n", + " \"input_ids\": input_ids,\n", + " \"past_key_values\": past_key_values,\n", + " \"position_ids\": position_ids,\n", + " \"attention_mask\": attention_mask,\n", + " \"return_last_logit\": True,\n", + " \"use_cache\": use_cache\n", + " }\n", + "```\n", + "\n", + "- **作用**：准备生成过程所需的输入。具体包括：\n", + " - 获取或更新 `position_ids`。\n", + " - 如果不是第一次前向传播，且存在 `past_key_values`，则只取 `input_ids` 和 `position_ids` 的最后一个 token。\n", + "- **区别**：该函数主要用于处理输入数据，确保输入数据的形状和内容适合当前生成步骤。\n", + "- **联系**：在每一步生成过程中，会调用该函数准备输入数据，尤其是处理 `past_key_values` 以提高生成效率。\n", + "\n", + "#### 3. `forward` 函数\n", + "\n", + "```python\n", + "def forward(\n", + " self,\n", + " input_ids: Optional[torch.Tensor] = None,\n", + " position_ids: Optional[torch.Tensor] = None,\n", + " attention_mask: Optional[torch.Tensor] = None,\n", + " past_key_values: Optional[Tuple[torch.FloatTensor]] = None,\n", + " inputs_embeds: Optional[torch.Tensor] = None,\n", + " labels: Optional[torch.Tensor] = None,\n", + " use_cache: Optional[bool] = None,\n", + " output_attentions: Optional[bool] = None,\n", + " output_hidden_states: Optional[bool] = None,\n", + " return_dict: Optional[bool] = None,\n", + " return_last_logit: Optional[bool] = False,\n", + "):\n", + " use_cache = use_cache if use_cache is not None else self.config.use_cache\n", + " return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n", + "\n", + " transformer_outputs = self.transformer(\n", + " input_ids=input_ids,\n", + " position_ids=position_ids,\n", + " attention_mask=attention_mask,\n", + " past_key_values=past_key_values,\n", + " inputs_embeds=inputs_embeds,\n", + " use_cache=use_cache,\n", + " output_hidden_states=output_hidden_states,\n", + " return_dict=return_dict,\n", + " ) # 使用 ChatGLMModel 类\n", + "\n", + " hidden_states = transformer_outputs[0]\n", + " if return_last_logit:\n", + " hidden_states = hidden_states[:, -1:]\n", + " lm_logits = self.transformer.output_layer(hidden_states)\n", + "\n", + " loss = None\n", + " if labels is not None:\n", + " lm_logits = lm_logits.to(torch.float32)\n", + "\n", + " # Shift so that tokens < n predict n\n", + " shift_logits = lm_logits[..., :-1, :].contiguous()\n", + " shift_labels = labels[..., 1:].contiguous()\n", + " # Flatten the tokens\n", + " loss_fct = CrossEntropyLoss(ignore_index=-100)\n", + " loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n", + "\n", + " lm_logits = lm_logits.to(hidden_states.dtype)\n", + " loss = loss.to(hidden_states.dtype)\n", + "\n", + " if not return_dict:\n", + " output = (lm_logits,) + transformer_outputs[1:]\n", + " return ((loss,) + output) if loss is not None else output\n", + "\n", + " return CausalLMOutputWithPast(\n", + " loss=loss,\n", + " logits=lm_logits,\n", + " past_key_values=transformer_outputs.past_key_values,\n", + " hidden_states=transformer_outputs.hidden_states,\n", + " attentions=transformer_outputs.attentions,\n", + " )\n", + "```\n", + "\n", + "- **作用**：执行前向传播，生成模型的输出。具体包括：\n", + " - 将输入数据传递给 `transformer`（`ChatGLMModel`），进行前向计算。\n", + " - 计算语言模型的 logits 和（如果有标签）计算损失。\n", + " - 返回模型输出，包括 logits、`past_key_values`、隐藏状态和注意力权重。\n", + "- **区别**：这是模型的核心前向传播逻辑，直接处理输入数据并生成输出。\n", + "- **联系**：`forward` 函数使用了 `ChatGLMModel` 类来进行实际的前向传播，并调用了之前的 `prepare_inputs_for_generation` 来准备输入。\n", + "\n", + "### 联系和流程\n", + "\n", + "1. **准备输入数据**：\n", + " - `prepare_inputs_for_generation` 函数用于处理输入数据，尤其是处理 `past_key_values` 以便只传递必要的最后一个 token。\n", + " \n", + "2. **执行前向传播**：\n", + " - `forward` 函数使用准备好的输入数据进行前向传播，生成输出。\n", + "\n", + "3. **更新模型参数**：\n", + " - `_update_model_kwargs_for_generation` 函数在每一步生成之后，更新模型的关键参数（如 `past_key_values`、`attention_mask` 和 `position_ids`），确保在下一步生成中使用最新的数据。\n", + "\n", + "通过这些函数的紧密配合，可以高效地实现生成任务中的前向传播和缓存管理，从而提高模型的生成效率和效果。" + ] + }, + { + "cell_type": "markdown", + "id": "1bfbd264-47bc-43bf-853a-4470462e31cd", + "metadata": {}, + "source": [ + "这三个函数是 `ChatGLMForConditionalGeneration` 类中的核心函数，分别用于处理不同的生成任务需求。它们之间有一定的联系，同时也有各自的用途和特点。以下是对它们的详细解释及其区别和联系：\n", + "\n", + "### 1. `chat` 函数\n", + "\n", + "```python\n", + "@torch.inference_mode()\n", + "def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = \"user\",\n", + " max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,\n", + " **kwargs):\n", + " if history is None:\n", + " history = []\n", + " if logits_processor is None:\n", + " logits_processor = LogitsProcessorList()\n", + " logits_processor.append(InvalidScoreLogitsProcessor())\n", + " gen_kwargs = {\"max_length\": max_length, \"num_beams\": num_beams, \"do_sample\": do_sample, \"top_p\": top_p,\n", + " \"temperature\": temperature, \"logits_processor\": logits_processor, **kwargs}\n", + " history.append({\"role\": role, \"content\": query})\n", + " inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,\n", + " return_tensors=\"pt\", return_dict=True)\n", + " inputs = inputs.to(self.device)\n", + " eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(\"\"), tokenizer.convert_tokens_to_ids(\"\")]\n", + " outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)\n", + " outputs = outputs.tolist()[0][len(inputs[\"input_ids\"][0]):-1]\n", + " response = tokenizer.decode(outputs)\n", + " response, history = self.process_response(response, history)\n", + " return response, history\n", + "```\n", + "\n", + "- **作用**：执行一次完整的聊天会话。将用户的查询和历史记录编码成模型输入，生成响应并更新历史记录。\n", + "- **区别**：这是一个高层次的接口，适用于一次性生成完整响应。适合用于需要立即获得完整回答的应用场景。\n", + "- **联系**：它依赖于 `generate` 函数来实际生成响应，并调用 `process_response` 函数来处理生成的输出。\n", + "\n", + "### 2. `stream_chat` 函数\n", + "\n", + "```python\n", + "@torch.inference_mode()\n", + "def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = \"user\",\n", + " past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,\n", + " logits_processor=None, return_past_key_values=False, **kwargs):\n", + " if history is None:\n", + " history = []\n", + " if logits_processor is None:\n", + " logits_processor = LogitsProcessorList()\n", + " logits_processor.append(InvalidScoreLogitsProcessor())\n", + " eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(\"\"), tokenizer.convert_tokens_to_ids(\"\")]\n", + " gen_kwargs = {\"max_length\": max_length, \"do_sample\": do_sample, \"top_p\": top_p,\n", + " \"temperature\": temperature, \"logits_processor\": logits_processor, **kwargs}\n", + " if past_key_values is None:\n", + " inputs = tokenizer.apply_chat_template(history + [{\"role\": role, \"content\": query}],\n", + " add_generation_prompt=True, tokenize=True, return_tensors=\"pt\",\n", + " return_dict=True)\n", + " else:\n", + " inputs = tokenizer.apply_chat_template([{\"role\": role, \"content\": query}], add_special_tokens=False,\n", + " add_generation_prompt=True, tokenize=True, return_tensors=\"pt\",\n", + " return_dict=True)\n", + " inputs = inputs.to(self.device)\n", + " if past_key_values is not None:\n", + " past_length = past_key_values[0][0].shape[2]\n", + " inputs.position_ids += past_length\n", + " attention_mask = inputs.attention_mask\n", + " attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)\n", + " inputs['attention_mask'] = attention_mask\n", + " history.append({\"role\": role, \"content\": query})\n", + " for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,\n", + " eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,\n", + " **gen_kwargs):\n", + " if return_past_key_values:\n", + " outputs, past_key_values = outputs\n", + " outputs = outputs.tolist()[0][len(inputs[\"input_ids\"][0]):-1]\n", + " response = tokenizer.decode(outputs)\n", + " if response and response[-1] != \"�\":\n", + " response, new_history = self.process_response(response, history)\n", + " if return_past_key_values:\n", + " yield response, new_history, past_key_values\n", + " else:\n", + " yield response, new_history\n", + "```\n", + "\n", + "- **作用**：实现流式聊天会话。与 `chat` 函数类似，但它通过生成器逐步返回响应，适用于流式生成应用场景。\n", + "- **区别**：支持逐步生成响应，使得可以在生成过程中动态处理和显示部分响应。\n", + "- **联系**：依赖于 `stream_generate` 函数来逐步生成响应，并在每次生成新的响应片段后调用 `process_response` 函数来处理和更新历史记录。\n", + "\n", + "### 3. `stream_generate` 函数\n", + "\n", + "```python\n", + "@torch.inference_mode()\n", + "def stream_generate(\n", + " self,\n", + " input_ids,\n", + " generation_config: Optional[GenerationConfig] = None,\n", + " logits_processor: Optional[LogitsProcessorList] = None,\n", + " stopping_criteria: Optional[StoppingCriteriaList] = None,\n", + " prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,\n", + " return_past_key_values=False,\n", + " **kwargs,\n", + "):\n", + " batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]\n", + "\n", + " if generation_config is None:\n", + " generation_config = self.generation_config\n", + " generation_config = copy.deepcopy(generation_config)\n", + " model_kwargs = generation_config.update(**kwargs)\n", + " model_kwargs[\"use_cache\"] = generation_config.use_cache\n", + " bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id\n", + "\n", + " if isinstance(eos_token_id, int):\n", + " eos_token_id = [eos_token_id]\n", + " eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None\n", + "\n", + " has_default_max_length = kwargs.get(\"max_length\") is None and generation_config.max_length is not None\n", + " if has_default_max_length and generation_config.max_new_tokens is None:\n", + " warnings.warn(\n", + " f\"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. \"\n", + " \"This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we\"\n", + " \" recommend using `max_new_tokens` to control the maximum length of the generation.\",\n", + " UserWarning,\n", + " )\n", + " elif generation_config.max_new_tokens is not None:\n", + " generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length\n", + " if not has_default_max_length:\n", + " logger.warn(\n", + " f\"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=\"\n", + " f\"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. \"\n", + " \"Please refer to the documentation for more information. \"\n", + " \"(https://hf-mirror.com/docs/transformers/main/en/main_classes/text_generation)\",\n", + " UserWarning,\n", + " )\n", + "\n", + " if input_ids_seq_length >= generation_config.max_length:\n", + " input_ids_string = \"decoder_input_ids\" if self.config.is_encoder_decoder else \"input_ids\"\n", + " logger.warning(\n", + " f\"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to\"\n", + " f\" {generation_config.max_length}. This can lead to unexpected behavior. You should consider\"\n", + " \" increasing `max_new_tokens`.\"\n", + " )\n", + "\n", + " # 2. Set generation parameters if not already defined\n", + " logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()\n", + " stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()\n", + "\n", + " logits_processor = self._get_logits_processor(\n", + " generation_config=generation_config,\n", + " input_ids_seq_length=input_ids_seq_length,\n", + " encoder_input_ids=input_ids,\n", + " prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,\n", + " logits_processor=logits_processor,\n", + " )\n", + "\n", + " stopping_criteria = self._get_stopping_criteria(\n", + " generation_config=generation_config, stopping_criteria=stopping_criteria\n", + " )\n", + " logits_warper = self._get_logits_warper(generation_config)\n", + "\n", + " unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)\n", + " scores = None\n", + " while True:\n", + " model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)\n", + " # 前向传递获取下一个 token\n", + " outputs = self(\n", + " **model_inputs,\n", + " return_dict=True,\n", + " output_attentions=False,\n", + " output_hidden_states=False,\n", + " )\n", + "\n", + " next_token_logits = outputs.logits[:, -1, :]\n", + "\n", + " # 预处理分布\n", + " next_token_scores = logits_processor(input_ids, next_token_logits)\n", + " next_token_scores = logits_warper(input_ids, next_token_scores)\n", + "\n", + " # 采样\n", + " probs = nn.functional.softmax(next_token_scores, dim=-1)\n", + " if generation_config.do_sample:\n", + " next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)\n", + " else:\n", + " next_tokens = torch.argmax(probs, dim=-1)\n", + " \n", + "\n", + " # 更新生成的 ids、模型输入和下一个步骤的长度\n", + " input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)\n", + " model_kwargs = self._update_model_kwargs_for_generation(\n", + " outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder\n", + " )\n", + " unfinished_sequences = unfinished_sequences.mul(\n", + " next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)\n", + " )\n", + " if return_past_key_values:\n", + " yield input_ids, outputs.past_key_values\n", + " else:\n", + " yield input_ids\n", + " # 当每个句子完成时或超出最大长度时停止\n", + " if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):\n", + " break\n", + "```\n", + "\n", + "- **作用**：流式生成模型输出，逐步返回生成的 token 以便实时处理。\n", + "- **区别**：该函数通过生成器实现流式生成，每生成一个 token 就返回一次结果，适用于需要逐步展示生成结果的应用场景。\n", + "- **联系**：`stream_chat` 函数依赖 `stream_generate` 来逐步生成响应，并在每次生成新的 token 后更新输入和模型参数。\n", + "\n", + "### 区别和联系总结\n", + "\n", + "1. **区别**：\n", + " - `chat` 函数：用于一次性生成完整响应，适合需要立即获得完整回答的应用场景。\n", + " - `stream_chat` 函数：用于流式生成响应，适合逐步展示生成结果的应用场景。\n", + " - `stream_generate` 函数：实现流式生成的核心逻辑，通过生成器逐步返回生成的 token。\n", + "\n", + "2. **联系**：\n", + " - `chat` 和 `stream_chat` 都是高层次接口，用户通过这些接口与模型交互。\n", + " - `chat` 函数调用 `generate` 函数实现一次性生成，而 `stream_chat` 函数调用 `stream_generate` 函数实现流式生成。\n", + " - `stream_generate` 函数使用了 `prepare_inputs_for_generation` 来处理输入数据，并通过 `_update_model_kwargs_for_generation` 更新模型参数，确保在生成过程中使用最新的数据。\n", + "\n", + "通过这三个函数的协调工作，模型能够实现高效且灵活的生成任务，满足不同应用场景的需求。" + ] + }, + { + "cell_type": "code", + "execution_count": 75, + "id": "62284871-67f4-47a2-941f-46befd2032b7", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n" + ] + } + ], + "source": [ + "device = \"cuda\"\n", + "\n", + "tokenizer = ChatGLM4Tokenizer.from_pretrained(\"THUDM/glm-4-9b-chat\", trust_remote_code=True)\n", + "\n", + "query = \"你好\"\n", + "\n", + "inputs = tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": query}],\n", + " add_generation_prompt=True,\n", + " tokenize=True,\n", + " return_tensors=\"pt\",\n", + " return_dict=True\n", + " )\n", + "\n", + "inputs = inputs.to(device)" + ] + }, + { + "cell_type": "code", + "execution_count": 76, + "id": "0801f0ee-2d71-4eda-b2f5-5e00209cd0fb", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n", + "The `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed in the future versions. Please, pass a `BitsAndBytesConfig` object in `quantization_config` argument instead.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "41bf96f6983b4492b2d95ae8f5eaa7ae", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/10 [00:00 Union[str, List[int], List[str], List[List[int]], BatchEncoding]:\n", + "\n", + " if return_dict and not tokenize:\n", + " raise ValueError(\n", + " \"`return_dict=True` is incompatible with `tokenize=False`, because there is no dict \"\n", + " \"of tokenizer outputs to return.\"\n", + " )\n", + "\n", + " def handle_single_conversation(conversation):\n", + " input_ids = self.get_prefix_tokens() if add_special_tokens else []\n", + " input_message = \"[gMASK]\" if add_special_tokens else \"\"\n", + " for item in conversation:\n", + " if item.get(\"tools\"):\n", + " tools = item[\"tools\"]\n", + " content = \"你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\"\n", + " for tool in tools:\n", + " if tool[\"type\"] == \"function\":\n", + " function = tool[\"function\"]\n", + " content += f\"\\n\\n## {function['name']}\\n\\n{json.dumps(function, ensure_ascii=False, indent=4)}\"\n", + " content += \"\\n在调用上述函数时，请使用 Json 格式表示调用的参数。\"\n", + " elif tool[\"type\"] == \"python\":\n", + " content += \"\\n\\n## python\\n\\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。\"\n", + " elif tool[\"type\"] == \"simple_browser\":\n", + " content += \"\\n\\n## simple_browser\\n\\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\\n`open_url(url: str)`：打开指定的 URL。\\n\\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\\n\\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。\"\n", + " elif tool[\"type\"] == \"cogview\":\n", + " content += \"\\n\\n## cogview\\n\\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。\"\n", + " else:\n", + " raise NotImplementedError(f\"Unknown tool type {tool['type']}\")\n", + " input = self.build_single_message(\"system\", \"\", content, tokenize=tokenize)\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if item[\"content\"]:\n", + " input = self.build_single_message(\n", + " item[\"role\"],\n", + " item.get(\"metadata\", \"\"),\n", + " item[\"content\"],\n", + " tokenize=tokenize\n", + " )\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if add_generation_prompt:\n", + " if tokenize:\n", + " input_ids.extend([self.convert_tokens_to_ids(\"<|assistant|>\")])\n", + " else:\n", + " input_message += \"<|assistant|>\"\n", + "\n", + " return input_ids if tokenize else input_message\n", + "\n", + " # Main logic to handle different conversation formats\n", + " if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):\n", + " result = handle_single_conversation(conversation)\n", + " elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):\n", + " result = [handle_single_conversation(c) for c in conversation]\n", + " elif hasattr(conversation, \"messages\"):\n", + " result = handle_single_conversation(conversation.messages)\n", + " else:\n", + " raise ValueError(\"Invalid conversation format\")\n", + "\n", + " if tokenize:\n", + " output = self.batch_encode_plus(\n", + " [result] if isinstance(result[0], int) else result,\n", + " padding=padding,\n", + " truncation=truncation,\n", + " max_length=max_length,\n", + " return_tensors=return_tensors,\n", + " is_split_into_words=True,\n", + " add_special_tokens=False\n", + " )\n", + " if return_dict:\n", + " return output\n", + " else:\n", + " return output[\"input_ids\"]\n", + " else:\n", + " return result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ece39da8-2b28-4415-b226-0919f2beb6de", + "metadata": {}, + "outputs": [], + "source": [ + "def handle_single_conversation(conversation):\n", + " input_ids = self.get_prefix_tokens() if add_special_tokens else []\n", + " input_message = \"[gMASK]\" if add_special_tokens else \"\"\n", + " for item in conversation:\n", + " if item.get(\"tools\"):\n", + " tools = item[\"tools\"]\n", + " content = \"你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\"\n", + " for tool in tools:\n", + " if tool[\"type\"] == \"function\":\n", + " function = tool[\"function\"]\n", + " content += f\"\\n\\n## {function['name']}\\n\\n{json.dumps(function, ensure_ascii=False, indent=4)}\"\n", + " content += \"\\n在调用上述函数时，请使用 Json 格式表示调用的参数。\"\n", + " elif tool[\"type\"] == \"python\":\n", + " content += \"\\n\\n## python\\n\\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。\"\n", + " elif tool[\"type\"] == \"simple_browser\":\n", + " content += \"\\n\\n## simple_browser\\n\\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\\n`open_url(url: str)`：打开指定的 URL。\\n\\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\\n\\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。\"\n", + " elif tool[\"type\"] == \"cogview\":\n", + " content += \"\\n\\n## cogview\\n\\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。\"\n", + " else:\n", + " raise NotImplementedError(f\"Unknown tool type {tool['type']}\")\n", + " input = self.build_single_message(\"system\", \"\", content, tokenize=tokenize)\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if item[\"content\"]:\n", + " input = self.build_single_message(\n", + " item[\"role\"],\n", + " item.get(\"metadata\", \"\"),\n", + " item[\"content\"],\n", + " tokenize=tokenize\n", + " )\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if add_generation_prompt:\n", + " if tokenize:\n", + " input_ids.extend([self.convert_tokens_to_ids(\"<|assistant|>\")])\n", + " else:\n", + " input_message += \"<|assistant|>\"\n", + "\n", + " return input_ids if tokenize else input_message" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69cc9551-ad44-4ecd-b7a6-00267fbcc186", + "metadata": {}, + "outputs": [], + "source": [ + "convert_tokens_to_ids(\"<|assistant|>\")" + ] + }, + { + "cell_type": "code", + "execution_count": 58, + "id": "ef2589e3-5f81-4c93-bb2d-a240b6041411", + "metadata": {}, + "outputs": [], + "source": [ + "def handle_single_conversation(conversation):\n", + " input_ids = self.get_prefix_tokens() if add_special_tokens else []\n", + " input_message = \"[gMASK]\" if add_special_tokens else \"\"\n", + " for item in conversation:\n", + " if item.get(\"tools\"):\n", + " tools = item[\"tools\"]\n", + " content = \"你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\"\n", + " for tool in tools:\n", + " if tool[\"type\"] == \"function\":\n", + " function = tool[\"function\"]\n", + " content += f\"\\n\\n## {function['name']}\\n\\n{json.dumps(function, ensure_ascii=False, indent=4)}\"\n", + " content += \"\\n在调用上述函数时，请使用 Json 格式表示调用的参数。\"\n", + " elif tool[\"type\"] == \"python\":\n", + " content += \"\\n\\n## python\\n\\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。\"\n", + " elif tool[\"type\"] == \"simple_browser\":\n", + " content += \"\\n\\n## simple_browser\\n\\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\\n`open_url(url: str)`：打开指定的 URL。\\n\\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\\n\\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。\"\n", + " elif tool[\"type\"] == \"cogview\":\n", + " content += \"\\n\\n## cogview\\n\\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。\"\n", + " else:\n", + " raise NotImplementedError(f\"Unknown tool type {tool['type']}\")\n", + " input = self.build_single_message(\"system\", \"\", content, tokenize=tokenize)\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if item[\"content\"]:\n", + " input = self.build_single_message(\n", + " item[\"role\"],\n", + " item.get(\"metadata\", \"\"),\n", + " item[\"content\"],\n", + " tokenize=tokenize\n", + " )\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if add_generation_prompt:\n", + " if tokenize:\n", + " input_ids.extend([self.convert_tokens_to_ids(\"<|assistant|>\")])\n", + " else:\n", + " input_message += \"<|assistant|>\"\n", + " # if tokenize:\n", + " # input_ids.extend([self.convert_tokens_to_ids(\"[gMASK]\")]) # 使用特殊标记代替空字符串\n", + " # else:\n", + " # input_message += \"[gMASK]\"\n", + "\n", + " return input_ids if tokenize else input_message" + ] + }, + { + "cell_type": "code", + "execution_count": 60, + "id": "a8400577-cfae-44ab-a1a9-831177ec24c2", + "metadata": {}, + "outputs": [], + "source": [ + "def apply_chat_template(\n", + " self,\n", + " conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], \"Conversation\"],\n", + " add_generation_prompt: bool = False,\n", + " tokenize: bool = True,\n", + " padding: bool = False,\n", + " truncation: bool = False,\n", + " max_length: Optional[int] = None,\n", + " return_tensors: Optional[Union[str, TensorType]] = None,\n", + " return_dict: bool = False,\n", + " tokenizer_kwargs: Optional[Dict[str, Any]] = None,\n", + " add_special_tokens: bool = True,\n", + " **kwargs,\n", + ") -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:\n", + "\n", + " if return_dict and not tokenize:\n", + " raise ValueError(\n", + " \"`return_dict=True` is incompatible with `tokenize=False`, because there is no dict \"\n", + " \"of tokenizer outputs to return.\"\n", + " )\n", + "\n", + " def handle_single_conversation(conversation):\n", + " input_ids = self.get_prefix_tokens() if add_special_tokens else []\n", + " input_message = \"[gMASK]\" if add_special_tokens else \"\"\n", + " for item in conversation:\n", + " if item.get(\"tools\"):\n", + " tools = item[\"tools\"]\n", + " content = \"你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\"\n", + " for tool in tools:\n", + " if tool[\"type\"] == \"function\":\n", + " function = tool[\"function\"]\n", + " content += f\"\\n\\n## {function['name']}\\n\\n{json.dumps(function, ensure_ascii=False, indent=4)}\"\n", + " content += \"\\n在调用上述函数时，请使用 Json 格式表示调用的参数。\"\n", + " elif tool[\"type\"] == \"python\":\n", + " content += \"\\n\\n## python\\n\\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。\"\n", + " elif tool[\"type\"] == \"simple_browser\":\n", + " content += \"\\n\\n## simple_browser\\n\\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\\n`open_url(url: str)`：打开指定的 URL。\\n\\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\\n\\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。\"\n", + " elif tool[\"type\"] == \"cogview\":\n", + " content += \"\\n\\n## cogview\\n\\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。\"\n", + " else:\n", + " raise NotImplementedError(f\"Unknown tool type {tool['type']}\")\n", + " input = self.build_single_message(\"system\", \"\", content, tokenize=tokenize)\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if item[\"content\"]:\n", + " input = self.build_single_message(\n", + " item[\"role\"],\n", + " item.get(\"metadata\", \"\"),\n", + " item[\"content\"],\n", + " tokenize=tokenize\n", + " )\n", + " if tokenize:\n", + " input_ids.extend(input)\n", + " else:\n", + " input_message += input\n", + " if add_generation_prompt:\n", + " if tokenize:\n", + " input_ids.extend([self.convert_tokens_to_ids(\"<|assistant|>\")])\n", + " else:\n", + " input_message += \"<|assistant|>\"\n", + " # if tokenize:\n", + " # input_ids.extend([self.convert_tokens_to_ids(\"[gMASK]\")]) # 使用特殊标记代替空字符串\n", + " # else:\n", + " # input_message += \"[gMASK]\"\n", + "\n", + " return input_ids if tokenize else input_message\n", + "\n", + " # 处理不同会话格式的主逻辑\n", + " if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):\n", + " result = handle_single_conversation(conversation)\n", + " elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):\n", + " result = [handle_single_conversation(c) for c in conversation]\n", + " elif hasattr(conversation, \"messages\"):\n", + " result = handle_single_conversation(conversation.messages)\n", + " else:\n", + " raise ValueError(\"Invalid conversation format\")\n", + "\n", + " if tokenize:\n", + " output = self.batch_encode_plus(\n", + " [result] if isinstance(result[0], int) else result,\n", + " padding=padding,\n", + " truncation=truncation,\n", + " max_length=max_length,\n", + " return_tensors=return_tensors,\n", + " is_split_into_words=True,\n", + " add_special_tokens=False\n", + " )\n", + " if return_dict:\n", + " return output\n", + " else:\n", + " return output[\"input_ids\"]\n", + " else:\n", + " return result\n" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "38eb42ba-0628-4f9d-bf29-410607552950", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple\n", + "Requirement already satisfied: regex in /data1/ckw/micromamba/envs/kewei-ai/lib/python3.11/site-packages (2023.10.3)\n", + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], + "source": [ + "%pip install regex" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "80613542-7de6-458e-9b93-8dc022fc2801", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "regex.Regex(\"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\\\r\\\\n\\\\p{L}\\\\p{N}]?\\\\p{L}+|\\\\p{N}{1,3}| ?[^\\\\s\\\\p{L}\\\\p{N}]+[\\\\r\\\\n]*|\\\\s*[\\\\r\\\\n]+|\\\\s+(?!\\\\S)|\\\\s+\", flags=regex.V0)" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pat_str = \"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\\\r\\\\n\\\\p{L}\\\\p{N}]?\\\\p{L}+|\\\\p{N}{1,3}| ?[^\\\\s\\\\p{L}\\\\p{N}]+[\\\\r\\\\n]*|\\\\s*[\\\\r\\\\n]+|\\\\s+(?!\\\\S)|\\\\s+\"\n", + "regex.compile(pat_str)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "e08863f0-9acd-458b-9af8-4ba37f1d7c21", + "metadata": {}, + "outputs": [], + "source": [ + "import regex\n", + "\n", + "pat_str = r\"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+\"\n", + "pattern = regex.compile(pat_str)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kewei-ai", + "language": "python", + "name": "kewei-ai" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/Model_Architecture_Discussions/rwkv-v4/rwkv-v4-guide.ipynb b/Model_Architecture_Discussions/rwkv-v4/rwkv-v4-guide.ipynb new file mode 100644 index 0000000..b9f1809 --- /dev/null +++ b/Model_Architecture_Discussions/rwkv-v4/rwkv-v4-guide.ipynb @@ -0,0 +1,528 @@ +{ + "cells": [ + { + "cell_type": "raw", + "id": "bcd88fb5-6a0f-4c4b-81fd-34be59ea7903", + "metadata": {}, + "source": [ + "模型下载链接：https://hf-mirror.com/BlinkDL/rwkv-4-pile-430m/resolve/main/RWKV-4-Pile-430M-20220808-8066.pth?download=true" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5b78b7ef-acc6-46cf-88c2-f90a2835e4b3", + "metadata": {}, + "outputs": [], + "source": [ + "########################################################################################################\n", + "# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM\n", + "########################################################################################################\n", + "\n", + "import numpy as np\n", + "np.set_printoptions(precision=4, suppress=True, linewidth=200)\n", + "import types, torch\n", + "from torch.nn import functional as F\n", + "from tokenizers import Tokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "deacc22b-2896-4b77-b595-3284b0c13544", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer = Tokenizer.from_file(\"20B_tokenizer.json\")\n", + "\n", + "args = types.SimpleNamespace()\n", + "args.MODEL_NAME = '/data1/ckw/RWKV-4-Pile-430M-20220808-8066'\n", + "args.n_layer = 24\n", + "args.n_embd = 1024\n", + "\n", + "context = \"\\nDataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence.\"\n", + "NUM_TRIALS = 3\n", + "LENGTH_PER_TRIAL = 100\n", + "TEMPERATURE = 1.0\n", + "TOP_P = 0.85\n", + "########################################################################################################" + ] + }, + { + "cell_type": "markdown", + "id": "3c85bca7-1342-4d8b-870c-baddcf2661d6", + "metadata": {}, + "source": [ + "### RWKV 的时间混合实现\n", + "\n", + "在 RWKV 模型中，时间混合（Time Mixing）是一个关键步骤，用于处理输入序列随时间的变化。以下是 `time_mixing` 函数的详细公式说明和代码注释。\n", + "\n", + "#### 公式说明\n", + "\n", + "时间混合的核心思想是通过时间混合系数将当前输入与先前的状态混合，以生成新的键、值和门控信号。这一过程涉及如下步骤：\n", + "\n", + "1. **混合输入**：\n", + " - 对当前输入 \\( x \\) 和前一状态进行加权平均：\n", + " $$ x_k = x \\cdot \\text{time\\_mix\\_k} + \\text{state}[5i+1] \\cdot (1 - \\text{time\\_mix\\_k}) $$\n", + " $$ x_v = x \\cdot \\text{time\\_mix\\_v} + \\text{state}[5i+1] \\cdot (1 - \\text{time\\_mix\\_v}) $$\n", + " $$ x_r = x \\cdot \\text{time\\_mix\\_r} + \\text{state}[5i+1] \\cdot (1 - \\text{time\\_mix\\_r}) $$\n", + "\n", + "2. **状态更新**：\n", + " - 更新状态：\n", + " $$ \\text{state}[5i+1] = x $$\n", + "\n", + "3. **计算门控信号**：\n", + " - 使用 sigmoid 激活函数计算门控信号 \\( r \\)：\n", + " $$ r = \\sigma(\\text{rw} @ x_r) $$\n", + "\n", + "4. **计算键和值**：\n", + " - 通过线性变换生成键 \\( k \\) 和值 \\( v \\)：\n", + " $$ k = \\text{kw} @ x_k $$\n", + " $$ v = \\text{vw} @ x_v $$\n", + "\n", + "5. **加权和计算**：\n", + " - 根据加权和公式计算加权和 \\( wkv \\)：\n", + " $$ a = e1 \\cdot aa + e2 \\cdot v $$\n", + " $$ b = e1 \\cdot bb + e2 $$\n", + " $$ \\text{wkv} = a / b $$\n", + "\n", + "代码如下：\n", + "\n", + "```python\n", + "@torch.jit.script_method\n", + "def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_first, time_decay, kw, vw, rw, ow):\n", + " # 混合当前输入和先前的状态\n", + " xk = x * time_mix_k + state[5*i+1] * (1 - time_mix_k)\n", + " xv = x * time_mix_v + state[5*i+1] * (1 - time_mix_v)\n", + " xr = x * time_mix_r + state[5*i+1] * (1 - time_mix_r)\n", + "\n", + " # 更新状态\n", + " state[5*i+1] = x\n", + "\n", + " # 计算门控信号\n", + " r = torch.sigmoid(rw @ xr)\n", + " \n", + " # 计算键和值\n", + " k = kw @ xk\n", + " v = vw @ xv\n", + "\n", + " # 从状态中读取先前的累积值\n", + " aa = state[5*i+2]\n", + " bb = state[5*i+3]\n", + " pp = state[5*i+4]\n", + "\n", + " # 计算加权和的第一部分\n", + " ww = time_first + k\n", + " qq = torch.maximum(pp, ww)\n", + " e1 = torch.exp(pp - qq)\n", + " e2 = torch.exp(ww - qq)\n", + " a = e1 * aa + e2 * v\n", + " b = e1 * bb + e2\n", + " wkv = a / b\n", + "\n", + " # 计算新的权重和状态\n", + " ww = pp + time_decay\n", + " qq = torch.maximum(ww, k)\n", + " e1 = torch.exp(ww - qq)\n", + " e2 = torch.exp(k - qq)\n", + " state[5*i+2] = e1 * aa + e2 * v\n", + " state[5*i+3] = e1 * bb + e2\n", + " state[5*i+4] = qq\n", + "\n", + " # 计算最终的输出\n", + " return ow @ (r * wkv)\n", + "```\n", + "\n", + "### 详细解释\n", + "\n", + "1. **混合输入**：\n", + " - `xk`, `xv`, `xr` 是输入 `x` 与状态 `state` 的加权混合，分别用于计算键、值和门控信号。\n", + "\n", + "2. **状态更新**：\n", + " - 将当前输入 `x` 存储在状态数组中，供下一步计算使用。\n", + "\n", + "3. **计算门控信号**：\n", + " - 使用 `torch.sigmoid` 计算门控信号 `r`，它决定了多少信息将被传递。\n", + "\n", + "4. **计算键和值**：\n", + " - 使用矩阵乘法计算键 `k` 和值 `v`。\n", + "\n", + "5. **加权和计算**：\n", + " - 通过指数加权平均计算加权和 `wkv`，这涉及处理数值稳定性问题（通过 `torch.maximum` 和指数运算）。\n", + "\n", + "6. **更新状态**：\n", + " - 更新状态数组中的累积值，以便后续时间步使用。\n", + "\n", + "7. **计算最终输出**：\n", + " - 使用门控信号 `r` 和加权和 `wkv` 计算最终输出。\n", + "\n", + "这样，通过逐步混合当前输入和先前的状态，RWKV 模型实现了时间序列数据的有效处理。" + ] + }, + { + "cell_type": "markdown", + "id": "1f0d47e4-1792-47e4-a506-3071f510526e", + "metadata": {}, + "source": [ + "### RWKV 的通道混合（Channel Mixing）实现与代码注释\n", + "\n", + "在 RWKV 模型中，通道混合（Channel Mixing）是另一个关键步骤，用于处理不同通道之间的信息交换。以下是 `channel_mixing` 函数的详细公式说明和代码注释。\n", + "\n", + "#### 公式说明\n", + "\n", + "通道混合的核心思想是通过通道混合系数将当前输入与先前的状态混合，以生成新的键和门控信号。这一过程涉及如下步骤：\n", + "\n", + "1. **混合输入**：\n", + " - 对当前输入 \\( x \\) 和前一状态进行加权平均：\n", + " $$ x_k = x \\cdot \\text{time\\_mix\\_k} + \\text{state}[5i+0] \\cdot (1 - \\text{time\\_mix\\_k}) $$\n", + " $$ x_r = x \\cdot \\text{time\\_mix\\_r} + \\text{state}[5i+0] \\cdot (1 - \\text{time\\_mix\\_r}) $$\n", + "\n", + "2. **状态更新**：\n", + " - 更新状态：\n", + " $$ \\text{state}[5i+0] = x $$\n", + "\n", + "3. **计算门控信号**：\n", + " - 使用 sigmoid 激活函数计算门控信号 \\( r \\)：\n", + " $$ r = \\sigma(\\text{rw} @ x_r) $$\n", + "\n", + "4. **计算键**：\n", + " - 通过 ReLU 和平方变换生成键 \\( k \\)：\n", + " $$ k = (\\text{ReLU}(\\text{kw} @ x_k))^2 $$\n", + "\n", + "5. **计算输出**：\n", + " - 使用门控信号和键计算最终的输出：\n", + " $$ \\text{output} = r \\cdot (\\text{vw} @ k) $$\n", + "\n", + "代码如下：\n", + "\n", + "```python\n", + "@torch.jit.script_method\n", + "def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", + " # 混合当前输入和先前的状态\n", + " xk = x * time_mix_k + state[5*i+0] * (1 - time_mix_k)\n", + " xr = x * time_mix_r + state[5*i+0] * (1 - time_mix_r)\n", + "\n", + " # 更新状态\n", + " state[5*i+0] = x\n", + "\n", + " # 计算门控信号\n", + " r = torch.sigmoid(rw @ xr)\n", + "\n", + " # 计算键，并通过ReLU和平方变换\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + "\n", + " # 计算最终的输出\n", + " return r * (vw @ k)\n", + "```\n", + "\n", + "\n", + "1. **混合输入**：\n", + " - `xk`, `xr` 是输入 `x` 与状态 `state` 的加权混合，分别用于计算键和门控信号。\n", + "\n", + "2. **状态更新**：\n", + " - 将当前输入 `x` 存储在状态数组中，供下一步计算使用。\n", + "\n", + "3. **计算门控信号**：\n", + " - 使用 `torch.sigmoid` 计算门控信号 `r`，它决定了多少信息将被传递。\n", + "\n", + "4. **计算键**：\n", + " - 使用 `torch.relu` 计算键 `k`，然后进行平方变换以增加非线性特性。\n", + "\n", + "5. **计算最终输出**：\n", + " - 使用门控信号 `r` 和键 `k` 计算最终输出。\n", + "\n", + "通过这些步骤，RWKV 模型实现了通道间的信息有效交换，增强了模型对输入数据的处理能力。" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "0f1b2e2b-9f0d-4db3-b9d9-d43e3e2537ee", + "metadata": {}, + "outputs": [], + "source": [ + "class RWKV_RNN(torch.jit.ScriptModule):\n", + " def __init__(self, args):\n", + " super().__init__()\n", + " self.args = args\n", + " self.eval() # set torch to inference mode\n", + " \n", + " w = torch.load(args.MODEL_NAME + '.pth', map_location='cpu')\n", + " for k in w.keys():\n", + " if '.time_' in k: w[k] = w[k].squeeze()\n", + " if '.time_decay' in k: w[k] = -torch.exp(w[k].float()) # the real time decay is like e^{-e^x}\n", + " else: w[k] = w[k].float() # convert to f32 type\n", + " \n", + " self.w = types.SimpleNamespace() # set self.w from w\n", + " self.w.blocks = {}\n", + " for k in w.keys(): # example: \"blocks.0.att.time_first\" => self.w.blocks[0].att.time_first\n", + " parts = k.split('.')\n", + " last = parts.pop()\n", + " here = self.w\n", + " for p in parts:\n", + " if p.isdigit():\n", + " p = int(p)\n", + " if p not in here: here[p] = types.SimpleNamespace()\n", + " here = here[p]\n", + " else:\n", + " if not hasattr(here, p): setattr(here, p, types.SimpleNamespace())\n", + " here = getattr(here, p)\n", + " setattr(here, last, w[k])\n", + "\n", + " def layer_norm(self, x, w):\n", + " return F.layer_norm(x, (self.args.n_embd,), weight=w.weight, bias=w.bias)\n", + "\n", + " @torch.jit.script_method\n", + " def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", + " xk = x * time_mix_k + state[5*i+0] * (1 - time_mix_k)\n", + " xr = x * time_mix_r + state[5*i+0] * (1 - time_mix_r)\n", + " state[5*i+0] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + " return r * (vw @ k)\n", + "\n", + " @torch.jit.script_method\n", + " def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_first, time_decay, kw, vw, rw, ow):\n", + " xk = x * time_mix_k + state[5*i+1] * (1 - time_mix_k)\n", + " xv = x * time_mix_v + state[5*i+1] * (1 - time_mix_v)\n", + " xr = x * time_mix_r + state[5*i+1] * (1 - time_mix_r)\n", + " state[5*i+1] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = kw @ xk\n", + " v = vw @ xv\n", + " \n", + " aa = state[5*i+2]\n", + " bb = state[5*i+3]\n", + " pp = state[5*i+4]\n", + " ww = time_first + k\n", + " qq = torch.maximum(pp, ww)\n", + " e1 = torch.exp(pp - qq)\n", + " e2 = torch.exp(ww - qq)\n", + " a = e1 * aa + e2 * v\n", + " b = e1 * bb + e2\n", + " wkv = a / b\n", + " ww = pp + time_decay\n", + " qq = torch.maximum(ww, k)\n", + " e1 = torch.exp(ww - qq)\n", + " e2 = torch.exp(k - qq)\n", + " state[5*i+2] = e1 * aa + e2 * v\n", + " state[5*i+3] = e1 * bb + e2\n", + " state[5*i+4] = qq\n", + " return ow @ (r * wkv)\n", + "\n", + " def forward(self, token, state):\n", + " with torch.no_grad():\n", + " if state == None:\n", + " state = torch.zeros(self.args.n_layer * 5, self.args.n_embd)\n", + " for i in range(self.args.n_layer): state[5*i+4] = -1e30 # -infinity\n", + " \n", + " x = self.w.emb.weight[token]\n", + " x = self.layer_norm(x, self.w.blocks[0].ln0)\n", + " for i in range(self.args.n_layer):\n", + " att = self.w.blocks[i].att\n", + " x = x + self.time_mixing(self.layer_norm(x, self.w.blocks[i].ln1), state, i, \n", + " att.time_mix_k, att.time_mix_v, att.time_mix_r, att.time_first, att.time_decay, \n", + " att.key.weight, att.value.weight, att.receptance.weight, att.output.weight)\n", + " ffn = self.w.blocks[i].ffn\n", + " x = x + self.channel_mixing(self.layer_norm(x, self.w.blocks[i].ln2), state, i, \n", + " ffn.time_mix_k, ffn.time_mix_r, \n", + " ffn.key.weight, ffn.value.weight, ffn.receptance.weight)\n", + " \n", + " x = self.w.head.weight @ self.layer_norm(x, self.w.ln_out)\n", + " return x.float(), state\n", + "\n", + "##########################################################################################################" + ] + }, + { + "cell_type": "markdown", + "id": "f1b457af-77a3-4b5e-a6f3-034b0fc6708d", + "metadata": {}, + "source": [ + "采样方法和v2、v3版本相比没有发生变化，代码做了一点优化调整而已。" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "fdf027b6-7df9-4c0f-818e-013e7c49e3cd", + "metadata": {}, + "outputs": [], + "source": [ + "def sample_logits(out, temperature=1.0, top_p=0.8):\n", + " probs = F.softmax(out, dim=-1).numpy()\n", + " sorted_probs = np.sort(probs)[::-1]\n", + " cumulative_probs = np.cumsum(sorted_probs)\n", + " cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])\n", + " probs[probs < cutoff] = 0\n", + " if temperature != 1.0:\n", + " probs = probs.pow(1.0 / temperature)\n", + " probs = probs / np.sum(probs)\n", + " out = np.random.choice(a=len(probs), p=probs)\n", + " return out\n", + "\n", + "########################################################################################################" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "298dbbde-6535-406b-bd43-f2d886799f8c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Using CPU. Loading /data1/ckw/RWKV-4-Pile-430M-20220808-8066 ...\n" + ] + } + ], + "source": [ + "print(f'\\nUsing CPU. Loading {args.MODEL_NAME} ...')\n", + "model = RWKV_RNN(args)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "7d366a89-02cb-4b5e-95ef-52f6376d3607", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Preprocessing context (slow version. see v2/rwkv/model.py for fast version)\n" + ] + } + ], + "source": [ + "print(f'\\nPreprocessing context (slow version. see v2/rwkv/model.py for fast version)')\n", + "init_state = None\n", + "for token in tokenizer.encode(context).ids:\n", + " init_out, init_state = model.forward(token, init_state)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "5273e7a8-875e-4998-b98e-f81951a7af32", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\n", + "--[ Trial 0 ]----------------- \n", + "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. The machine learning solutions applied to the class are called Persona, which consist of several categories:\n", + "\n", + "\\begin{tabular}{|c|c|c|}\n", + "\\hline\n", + " Name & Description \\\\ \\hline\n", + "\\hline\n", + " \\end{tabular}\n", + "\n", + "DataWhalechina organizes the data in two ways:\n", + "\n", + "\\begin{tabular}{|c|c|c|}\n", + "\\hline\n", + " \\multicolumn{2}{|c}{\n", + "\n", + "--[ Trial 1 ]----------------- \n", + "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. The main goal is to allow learners to learn how to use artificial intelligence in an integrated fashion, by using both AI and deep learning techniques. Datawhalechina aims to teach AI algorithms from scratch and teach them from scratch to become competent with many algorithms that humans could not have.\n", + "\n", + "Applications\n", + "\n", + "Projects \n", + " DeeplearningAI : Encourage AI algorithms to become competent with many algorithms that humans could not have. Datawhalechina aims to be able to combine knowledge from multiple AI\n", + "\n", + "--[ Trial 2 ]----------------- \n", + "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. The company was founded in 2016. The company has graduated 1,000 engineers, who work from the companies headquarters in Shanghai.\n", + "\n", + "In September 2019, the team was reported to have learned over 400,000 artificial intelligence.\n", + "\n", + "In August 2019, the company was reported to have sold 600,000 artificial intelligence to clients in Singapore.\n", + "\n", + "References\n", + "\n", + "External links\n", + " \n", + "\n", + "Category:Human machine interaction\n", + "Category:Learning management systems\n", + "Category:Learning management systemsTechnologies, industry,\n", + "\n" + ] + } + ], + "source": [ + "for TRIAL in range(NUM_TRIALS):\n", + " print(f'\\n\\n--[ Trial {TRIAL} ]-----------------', context, end=\"\")\n", + " all_tokens = []\n", + " out_last = 0\n", + " out, state = init_out.clone(), init_state.clone()\n", + " for i in range(LENGTH_PER_TRIAL):\n", + " token = sample_logits(out, TEMPERATURE, TOP_P)\n", + " all_tokens += [token]\n", + " tmp = tokenizer.decode(all_tokens[out_last:])\n", + " if '\\ufffd' not in tmp: # only print when we have a valid utf-8 string\n", + " print(tmp, end=\"\", flush=True)\n", + " out_last = i + 1\n", + " out, state = model.forward(token, state) \n", + "print('\\n')" + ] + }, + { + "cell_type": "markdown", + "id": "f1cdc809-c64d-4861-b540-460bc1097e38", + "metadata": {}, + "source": [ + "### 备注：RWKV 的Scaling Law（缩放定律）\n", + "\n", + "RWKV 的缩放定律描述了模型性能随着各种因素变化的数学关系。这些因素包括模型大小（$N$）、数据集大小（$D$）或最优计算预算（$C_{\\min}$）。缩放定律的重要性体现在以下两个方面：\n", + "1. **预测与规划**：它们允许我们在训练大型模型之前，通过插值和外推来预测和规划成本和性能。\n", + "2. **反馈与研究**：它们提供了关于模型失效情况下的重要反馈，指引未来研究方向。\n", + "\n", + "#### 关键内容总结：\n", + "- **与之前的RNN研究对比**：之前的工作指出，LSTM不完全遵循与Transformer相同的对数线性缩放定律。然而，RWKV模型的训练结果表明，RWKV遵循与Transformer相同的一般缩放定律形式。\n", + "- **实验验证**：在[v4的论文](https://arxiv.org/abs/2305.13048)通过训练45个RWKV模型，验证了其损失与计算量之间的线性关系，线性拟合的 $r^2$ 值为0.994，即使外推一个数量级，拟合度仍然很好（$r^2$为0.875）。\n", + "\n", + "这些结果显示了RWKV模型在缩放时的优越性和与Transformer相似的性能缩放行为。" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13f6025d-faea-4647-be05-8fb4cce05991", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/Model_Architecture_Discussions/rwkv-v4/rwkv-v4.ipynb b/Model_Architecture_Discussions/rwkv-v4/rwkv-v4.ipynb deleted file mode 100644 index fffda45..0000000 --- a/Model_Architecture_Discussions/rwkv-v4/rwkv-v4.ipynb +++ /dev/null @@ -1,297 +0,0 @@ -{ - "cells": [ - { - "cell_type": "raw", - "id": "bcd88fb5-6a0f-4c4b-81fd-34be59ea7903", - "metadata": {}, - "source": [ - "模型下载链接：https://hf-mirror.com/BlinkDL/rwkv-4-pile-430m/resolve/main/RWKV-4-Pile-430M-20220808-8066.pth?download=true" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "5b78b7ef-acc6-46cf-88c2-f90a2835e4b3", - "metadata": {}, - "outputs": [], - "source": [ - "########################################################################################################\n", - "# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM\n", - "########################################################################################################\n", - "\n", - "import numpy as np\n", - "np.set_printoptions(precision=4, suppress=True, linewidth=200)\n", - "import types, torch\n", - "from torch.nn import functional as F\n", - "from tokenizers import Tokenizer" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "deacc22b-2896-4b77-b595-3284b0c13544", - "metadata": {}, - "outputs": [], - "source": [ - "tokenizer = Tokenizer.from_file(\"20B_tokenizer.json\")\n", - "\n", - "args = types.SimpleNamespace()\n", - "args.MODEL_NAME = '/data1/ckw/RWKV-4-Pile-430M-20220808-8066'\n", - "args.n_layer = 24\n", - "args.n_embd = 1024\n", - "\n", - "context = \"\\nDataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence.\"\n", - "NUM_TRIALS = 3\n", - "LENGTH_PER_TRIAL = 100\n", - "TEMPERATURE = 1.0\n", - "TOP_P = 0.85\n", - "########################################################################################################" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "0f1b2e2b-9f0d-4db3-b9d9-d43e3e2537ee", - "metadata": {}, - "outputs": [], - "source": [ - "class RWKV_RNN(torch.jit.ScriptModule):\n", - " def __init__(self, args):\n", - " super().__init__()\n", - " self.args = args\n", - " self.eval() # set torch to inference mode\n", - " \n", - " w = torch.load(args.MODEL_NAME + '.pth', map_location='cpu')\n", - " for k in w.keys():\n", - " if '.time_' in k: w[k] = w[k].squeeze()\n", - " if '.time_decay' in k: w[k] = -torch.exp(w[k].float()) # the real time decay is like e^{-e^x}\n", - " else: w[k] = w[k].float() # convert to f32 type\n", - " \n", - " self.w = types.SimpleNamespace() # set self.w from w\n", - " self.w.blocks = {}\n", - " for k in w.keys(): # example: \"blocks.0.att.time_first\" => self.w.blocks[0].att.time_first\n", - " parts = k.split('.')\n", - " last = parts.pop()\n", - " here = self.w\n", - " for p in parts:\n", - " if p.isdigit():\n", - " p = int(p)\n", - " if p not in here: here[p] = types.SimpleNamespace()\n", - " here = here[p]\n", - " else:\n", - " if not hasattr(here, p): setattr(here, p, types.SimpleNamespace())\n", - " here = getattr(here, p)\n", - " setattr(here, last, w[k])\n", - "\n", - " def layer_norm(self, x, w):\n", - " return F.layer_norm(x, (self.args.n_embd,), weight=w.weight, bias=w.bias)\n", - "\n", - " @torch.jit.script_method\n", - " def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", - " xk = x * time_mix_k + state[5*i+0] * (1 - time_mix_k)\n", - " xr = x * time_mix_r + state[5*i+0] * (1 - time_mix_r)\n", - " state[5*i+0] = x\n", - " r = torch.sigmoid(rw @ xr)\n", - " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", - " return r * (vw @ k)\n", - "\n", - " @torch.jit.script_method\n", - " def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_first, time_decay, kw, vw, rw, ow):\n", - " xk = x * time_mix_k + state[5*i+1] * (1 - time_mix_k)\n", - " xv = x * time_mix_v + state[5*i+1] * (1 - time_mix_v)\n", - " xr = x * time_mix_r + state[5*i+1] * (1 - time_mix_r)\n", - " state[5*i+1] = x\n", - " r = torch.sigmoid(rw @ xr)\n", - " k = kw @ xk\n", - " v = vw @ xv\n", - " \n", - " aa = state[5*i+2]\n", - " bb = state[5*i+3]\n", - " pp = state[5*i+4]\n", - " ww = time_first + k\n", - " qq = torch.maximum(pp, ww)\n", - " e1 = torch.exp(pp - qq)\n", - " e2 = torch.exp(ww - qq)\n", - " a = e1 * aa + e2 * v\n", - " b = e1 * bb + e2\n", - " wkv = a / b\n", - " ww = pp + time_decay\n", - " qq = torch.maximum(ww, k)\n", - " e1 = torch.exp(ww - qq)\n", - " e2 = torch.exp(k - qq)\n", - " state[5*i+2] = e1 * aa + e2 * v\n", - " state[5*i+3] = e1 * bb + e2\n", - " state[5*i+4] = qq\n", - " return ow @ (r * wkv)\n", - "\n", - " def forward(self, token, state):\n", - " with torch.no_grad():\n", - " if state == None:\n", - " state = torch.zeros(self.args.n_layer * 5, self.args.n_embd)\n", - " for i in range(self.args.n_layer): state[5*i+4] = -1e30 # -infinity\n", - " \n", - " x = self.w.emb.weight[token]\n", - " x = self.layer_norm(x, self.w.blocks[0].ln0)\n", - " for i in range(self.args.n_layer):\n", - " att = self.w.blocks[i].att\n", - " x = x + self.time_mixing(self.layer_norm(x, self.w.blocks[i].ln1), state, i, \n", - " att.time_mix_k, att.time_mix_v, att.time_mix_r, att.time_first, att.time_decay, \n", - " att.key.weight, att.value.weight, att.receptance.weight, att.output.weight)\n", - " ffn = self.w.blocks[i].ffn\n", - " x = x + self.channel_mixing(self.layer_norm(x, self.w.blocks[i].ln2), state, i, \n", - " ffn.time_mix_k, ffn.time_mix_r, \n", - " ffn.key.weight, ffn.value.weight, ffn.receptance.weight)\n", - " \n", - " x = self.w.head.weight @ self.layer_norm(x, self.w.ln_out)\n", - " return x.float(), state\n", - "\n", - "##########################################################################################################" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "fdf027b6-7df9-4c0f-818e-013e7c49e3cd", - "metadata": {}, - "outputs": [], - "source": [ - "def sample_logits(out, temperature=1.0, top_p=0.8):\n", - " probs = F.softmax(out, dim=-1).numpy()\n", - " sorted_probs = np.sort(probs)[::-1]\n", - " cumulative_probs = np.cumsum(sorted_probs)\n", - " cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])\n", - " probs[probs < cutoff] = 0\n", - " if temperature != 1.0:\n", - " probs = probs.pow(1.0 / temperature)\n", - " probs = probs / np.sum(probs)\n", - " out = np.random.choice(a=len(probs), p=probs)\n", - " return out\n", - "\n", - "########################################################################################################" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "298dbbde-6535-406b-bd43-f2d886799f8c", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Using CPU. Loading /data1/ckw/RWKV-4-Pile-430M-20220808-8066 ...\n" - ] - } - ], - "source": [ - "print(f'\\nUsing CPU. Loading {args.MODEL_NAME} ...')\n", - "model = RWKV_RNN(args)" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "7d366a89-02cb-4b5e-95ef-52f6376d3607", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Preprocessing context (slow version. see v2/rwkv/model.py for fast version)\n" - ] - } - ], - "source": [ - "print(f'\\nPreprocessing context (slow version. see v2/rwkv/model.py for fast version)')\n", - "init_state = None\n", - "for token in tokenizer.encode(context).ids:\n", - " init_out, init_state = model.forward(token, init_state)" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "5273e7a8-875e-4998-b98e-f81951a7af32", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "\n", - "--[ Trial 0 ]----------------- \n", - "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. Founded in 2015 by AI graduate student Yawei Li, DataWhalechina aims to help people learn to think more naturally about data. Learn more.\n", - "\n", - "50% of U.S. high school graduates who take data science courses go on to pursue masters degrees, which cost about $7,000, according to The American Council for an Energy Efficient Economy. Learn more.\n", - "\n", - "More than 600 startups compete for the same creative talent awards in 2016. If there is one award\n", - "\n", - "--[ Trial 1 ]----------------- \n", - "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. Datawhalechina was established in 2013. The company was created to serve the needs of companies that seek to increase the utilization of machine learning technology in their environments. This aims to create a platform that will help increase the adoption of machine learning technology in organizations by creating better decision support tools.\n", - "\n", - "As of 2017, Datawhalechina's team of specialists are spread over the United States, Europe, Asia, Africa and Canada. Their strategy includes providing low-cost software solutions to the\n", - "\n", - "--[ Trial 2 ]----------------- \n", - "DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence.\n", - "The main objective of the organization is to provide diverse students with the information, skills, knowledge and ideas needed to tackle big challenges in their future. The success of the program has prompted the city government to give them more resources to bring more students in the program.\n", - "\n", - "Ethereum\n", - "\n", - "The Ethereum (ETH) blockchain, designed by XRP, is a decentralised ledger technology that enables Bitcoin (BTC) and other cryptocurrencies to be used as payment. It aims to be the largest\n", - "\n" - ] - } - ], - "source": [ - "for TRIAL in range(NUM_TRIALS):\n", - " print(f'\\n\\n--[ Trial {TRIAL} ]-----------------', context, end=\"\")\n", - " all_tokens = []\n", - " out_last = 0\n", - " out, state = init_out.clone(), init_state.clone()\n", - " for i in range(LENGTH_PER_TRIAL):\n", - " token = sample_logits(out, TEMPERATURE, TOP_P)\n", - " all_tokens += [token]\n", - " tmp = tokenizer.decode(all_tokens[out_last:])\n", - " if '\\ufffd' not in tmp: # only print when we have a valid utf-8 string\n", - " print(tmp, end=\"\", flush=True)\n", - " out_last = i + 1\n", - " out, state = model.forward(token, state) \n", - "print('\\n')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d546bfd9-cf80-49bf-8c76-f3918d7d67e4", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.5" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/Model_Architecture_Discussions/rwkv-v5/RWKV-v5-guide.ipynb b/Model_Architecture_Discussions/rwkv-v5/RWKV-v5-guide.ipynb new file mode 100644 index 0000000..c765934 --- /dev/null +++ b/Model_Architecture_Discussions/rwkv-v5/RWKV-v5-guide.ipynb @@ -0,0 +1,869 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 4, + "id": "1fb76974-93ea-4b9c-81b1-55f826e7a361", + "metadata": {}, + "outputs": [], + "source": [ + "########################################################################################################\n", + "# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM\n", + "########################################################################################################\n", + "\n", + "import numpy as np\n", + "np.set_printoptions(precision=4, suppress=True, linewidth=200)\n", + "import types, torch\n", + "import torch.nn as nn\n", + "from torch.nn import functional as F\n", + "\n", + "MyModule = torch.jit.ScriptModule\n", + "MyFunction = torch.jit.script_method" + ] + }, + { + "cell_type": "markdown", + "id": "9c97049c-d3ae-4c72-bff4-d99416f8d650", + "metadata": {}, + "source": [ + "rwkv5又叫eagal" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b1059eca-db4f-4c0b-ae3e-37af49ec7fa1", + "metadata": {}, + "outputs": [], + "source": [ + "import torch" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "1c8d8009-7ee7-4419-aacb-cdc45f287010", + "metadata": {}, + "outputs": [], + "source": [ + "class RWKV_TOKENIZER():\n", + " table: list[list[list[bytes]]]\n", + " good: list[set[int]]\n", + " wlen: list[int]\n", + " def __init__(self, file_name):\n", + " self.idx2token = {}\n", + " sorted = [] # must be already sorted\n", + " lines = open(file_name, \"r\", encoding=\"utf-8\").readlines()\n", + " for l in lines:\n", + " idx = int(l[:l.index(' ')])\n", + " x = eval(l[l.index(' '):l.rindex(' ')])\n", + " x = x.encode(\"utf-8\") if isinstance(x, str) else x\n", + " assert isinstance(x, bytes)\n", + " assert len(x) == int(l[l.rindex(' '):])\n", + " sorted += [x]\n", + " self.idx2token[idx] = x\n", + "\n", + " self.token2idx = {}\n", + " for k, v in self.idx2token.items():\n", + " self.token2idx[v] = int(k)\n", + "\n", + " # precompute some tables for fast matching\n", + " self.table = [[[] for j in range(256)] for i in range(256)]\n", + " self.good = [set() for i in range(256)]\n", + " self.wlen = [0 for i in range(256)]\n", + "\n", + " for i in reversed(range(len(sorted))): # reverse order - match longer tokens first\n", + " s = sorted[i]\n", + " if len(s) >= 2:\n", + " s0 = int(s[0])\n", + " s1 = int(s[1])\n", + " self.table[s0][s1] += [s]\n", + " self.wlen[s0] = max(self.wlen[s0], len(s))\n", + " self.good[s0].add(s1)\n", + "\n", + " def encodeBytes(self, src: bytes) -> list[int]:\n", + " src_len: int = len(src)\n", + " tokens: list[int] = []\n", + " i: int = 0\n", + " while i < src_len:\n", + " s: bytes = src[i : i + 1]\n", + "\n", + " if i < src_len - 1:\n", + " s1: int = int(src[i + 1])\n", + " s0: int = int(src[i])\n", + " if s1 in self.good[s0]:\n", + " sss: bytes = src[i : i + self.wlen[s0]]\n", + " try:\n", + " s = next(filter(sss.startswith, self.table[s0][s1]))\n", + " except:\n", + " pass\n", + " tokens.append(self.token2idx[s])\n", + " i += len(s)\n", + "\n", + " return tokens\n", + "\n", + " def decodeBytes(self, tokens):\n", + " return b''.join(map(lambda i: self.idx2token[i], tokens))\n", + "\n", + " def encode(self, src: str):\n", + " return self.encodeBytes(src.encode(\"utf-8\"))\n", + "\n", + " def decode(self, tokens):\n", + " return self.decodeBytes(tokens).decode('utf-8')\n", + "\n", + " def printTokens(self, tokens):\n", + " for i in tokens:\n", + " s = self.idx2token[i]\n", + " try:\n", + " s = s.decode('utf-8')\n", + " except:\n", + " pass\n", + " print(f'{repr(s)}{i}', end=' ')\n", + " # print(repr(s), i)\n", + " print()\n", + "\n", + "########################################################################################################" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "63a4e8ba-a291-4fdc-aef1-ebfca21840d4", + "metadata": {}, + "outputs": [], + "source": [ + "def sample_logits(out, temperature=1.0, top_p=0.8):\n", + " probs = F.softmax(out, dim=-1).numpy()\n", + " sorted_probs = np.sort(probs)[::-1]\n", + " cumulative_probs = np.cumsum(sorted_probs)\n", + " cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])\n", + " probs[probs < cutoff] = 0\n", + " if temperature != 1.0:\n", + " probs = probs.pow(1.0 / temperature)\n", + " probs = probs / np.sum(probs)\n", + " out = np.random.choice(a=len(probs), p=probs)\n", + " return out\n", + "\n", + "########################################################################################################" + ] + }, + { + "cell_type": "raw", + "id": "cb8c7d5e-08cb-4780-b6d9-ab8bad1417d4", + "metadata": {}, + "source": [ + "可以从这个链接下载模型：\n", + "https://www.modelscope.cn/models/AI-ModelScope/rwkv-5-world/files\n", + "https://www.modelscope.cn/api/v1/models/AI-ModelScope/rwkv-5-world/repo?Revision=master&FilePath=RWKV-5-World-0.1B-v1-20230803-ctx4096.pth" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "94d7d6db-e89e-4209-ae72-6625ba85ef5b", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer = RWKV_TOKENIZER(\"./rwkv_vocab_v20230424.txt\")\n", + "\n", + "# THIS IS NOW UPDATED TO SUPPORT LATEST RWKV-5 WORLD v2 MODELS\n", + "\n", + "args = types.SimpleNamespace()\n", + "args.MODEL_NAME = '/data1/ckw/RWKV-5-World-0.4B-v2-20231113-ctx4096' #这里不用有后缀.pth\n", + "args.n_layer = 24\n", + "args.n_embd = 1024\n", + "args.vocab_size = 65536" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "c8dcf39a-7838-454b-85fc-ec9bd75fa243", + "metadata": {}, + "outputs": [], + "source": [ + "# N_LAYER=\"12\"\n", + "# N_EMBD=\"768\"\n", + "N_LAYER=\"24\"\n", + "N_EMBD=\"1024\"" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "74d7c96a-6fbc-401c-8078-fefb1a6ec5c3", + "metadata": {}, + "outputs": [], + "source": [ + "# context = \"\\nElon Musk has\"\n", + "# context = \"\\n我们发现\"\n", + "context = \"Q:Do you know datawhalechina?\\nA:\"\n", + "NUM_TRIALS = 3\n", + "LENGTH_PER_TRIAL = 100\n", + "LENGTH_PER_TRIAL = 4096\n", + "TEMPERATURE = 1.0\n", + "TOP_P = 0.7" + ] + }, + { + "cell_type": "markdown", + "id": "5ac1c244-71e3-4263-8ad8-4d0cf681ebfd", + "metadata": {}, + "source": [ + "Eagle (RWKV-5) 和 Finch (RWKV-6) 相较于基础的RWKV-4架构在建模上的改进：\n", + "\n", + "1. **改进步骤**：\n", + " - **Eagle的改进**：Eagle模型在RWKV-4的基础上进行了多项改进，包括引入矩阵值的注意力状态（matrix-valued attention states）、在注意力头上应用LayerNorm（层归一化）、使用SiLU（Sigmoid-Weighted Linear Unit）进行注意力门控、并改进了初始化方法。此外，Eagle移除了接受度（receptance）函数中的Sigmoid激活函数。\n", + " - **Finch的改进**：Finch模型进一步引入了对衰减计划（decay schedule）和令牌移位（token-shift）的数据依赖性（data-dependence），使模型在处理时间和令牌数据时更加灵活和精确。\n", + "\n", + "2. **核心架构**：\n", + " - 这些模型的核心架构依然类似于RWKV-4，由一系列堆叠的残差块组成，形状类似于传统的Transformer架构。\n", + " - 每个块包含一个预LayerNorm时间混合子层（Pre-LayerNorm Time-Mixing sub-layer）和一个预LayerNorm通道混合子层（Pre-LayerNorm Channel-Mixing sub-layer），对应于Transformer中的注意力子层和前馈网络子层。\n" + ] + }, + { + "cell_type": "markdown", + "id": "bd3d56d6-59af-41d0-9ac3-2cc5b4fb54ed", + "metadata": {}, + "source": [ + "这个是RWKV 5的Channel Mixing的代码实现，可以对比一下RWKV 4的实现。\n", + "\n", + "\n", + "```python\n", + "@MyFunction\n", + " def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", + " i0 = (2+self.head_size)*i+0\n", + " xk = x * time_mix_k + state[i0] * (1 - time_mix_k)\n", + " xr = x * time_mix_r + state[i0] * (1 - time_mix_r)\n", + " state[i0] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + " return r * (vw @ k)\n", + "```\n", + "\n", + "RWKV 4的Channel Mixing的代码实现为：\n", + "\n", + "\n", + "```python\n", + "@torch.jit.script_method\n", + " def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", + " xk = x * time_mix_k + state[5*i+0] * (1 - time_mix_k)\n", + " xr = x * time_mix_r + state[5*i+0] * (1 - time_mix_r)\n", + " state[5*i+0] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + " return r * (vw @ k)\n", + "```\n", + "\n", + "这里的`i`表示的是RWKV有多少层，在RWKV4的每一层中Channel Mixing记录一个状态，而每一个Time Mixing则记录4个状态，所以一共是5个状态。而RWKV 5中每一层现在记录了`2+self.head_size`个状态，Channel Mixing记录的状态以及计算过程和RWKV 4是完全一样的。" + ] + }, + { + "cell_type": "markdown", + "id": "976f399a-78ba-4fb2-9b52-d19afda8c5d0", + "metadata": {}, + "source": [ + "![](./img/01.png)\n", + "\n", + "图1：RWKV架构概述。左侧：时间混合和通道混合块；右上角：作为RNN单元的RWKV时间混合块；中下部：前馈模块中的令牌移位模块和Eagle时间混合；右下角：Finch时间混合中的令牌移位模块。所有形状注释为简洁起见假设为单头。虚线箭头（左侧，右上角）表示在Finch中有连接，但在Eagle中没有。" + ] + }, + { + "cell_type": "markdown", + "id": "ea37d6cd-1348-450b-9f6e-198d7c1d8368", + "metadata": {}, + "source": [ + "Eagle模型中采用的Token Shift技术：\n", + "\n", + "1. **Token Shift**：\n", + " - Eagle模型从之前的RWKV模型中采用了Token Shift技术，这类似于大小为2的一维因果卷积（1D causal convolution）。\n", + " - 在图1的中心底部可以看到该技术的示意图。\n", + "\n", + "2. **线性插值定义**：\n", + " - 为了更好地介绍Token Shift技术，定义了一些符号。\n", + " - 线性插值（lerp）在时间步$t$和$t-1$之间用于RWKV-4和Eagle Token Shift，定义如下：\n", + " \\begin{align*}\n", + " \\text{lerp}_{\\Box}(a, b) = a + (b - a) \\odot \\mu_{\\Box}\n", + " \\end{align*}\n", + " - 其中，每个$\\mu_{\\Box} \\in \\mathbb{R}^D$是一个可学习的向量。\n", + "\n", + "3. **Token Shift的功能**：\n", + " - Token Shift允许模型学习在每个时间步中分配新信息和旧信息的比例，适用于接受度（receptance）、键（key）、值（value）和门控向量（gate vectors）中的每个通道（$r, k, v, g$），且每个头部（head）独立且唯一地应用这些向量。\n", + " - 这使得即使在单层内，一个单独的头部也可以直接将过去和当前的令牌数据累积到这些向量的不同子空间中，从而形成感应头（induction heads）。\n" + ] + }, + { + "cell_type": "markdown", + "id": "40ca7c1b-8faa-42d5-8c7d-c4f233063856", + "metadata": {}, + "source": [ + "在Eagle和Finch模型中，通道混合模块（Channel Mixing module）的设置及其与RWKV-4架构的异同如下：\n", + "\n", + "1. **模块一致性**：\n", + " - 在Eagle和Finch模型中，通道混合模块与之前的RWKV-4架构基本相同。\n", + " - 唯一的区别在于Eagle模型中，通道混合模块的隐藏维度（hidden dimension）从原来的4D减少到了3.5D。\n", + "\n", + "2. **减少维度的原因**：\n", + " - 这个隐藏维度的减少是为了在Eagle时间混合（Eagle Time Mixing）中引入新的门控权重（gating weights）并确保与之前模型（在相同层数和嵌入维度下）的参数数量相等。\n", + "\n", + "3. **Finch模型中的处理**：\n", + " - 尽管Finch模型中增加了一些新的LoRA权重参数，但并没有进一步减少隐藏维度。\n", + "\n", + "4. **公式一致性**：\n", + " - 通道混合的公式与RWKV-4模型相同，为了符号一致性（notational consistency），再次列出这些公式：\n", + "\n", + "\\begin{align*}\n", + "r'_t &= \\text{lerp}_{r'}(x'_t, x'_{t-1}) W_{r'} \\in \\mathbb{R}^D \\quad \\text{(公式10)} \\\\\n", + "k'_t &= \\text{lerp}_{k'}(x'_t, x'_{t-1}) W_{k'} \\in \\mathbb{R}^{3.5D} \\quad \\text{(公式11)} \\\\\n", + "v'_t &= \\text{ReLU}(k'_t)^2 W_{v'} \\in \\mathbb{R}^D \\quad \\text{(公式12)} \\\\\n", + "o'_t &= \\sigma(r'_t) \\odot v'_t \\in \\mathbb{R}^D \\quad \\text{(公式13)}\n", + "\\end{align*}\n", + "\n", + "这些公式描述了在时间步 \\( t \\) 的通道混合操作：\n", + "- 使用线性插值（lerp）计算 \\( r'_t \\) 和 \\( k'_t \\)。\n", + "- \\( v'_t \\) 通过 \\( k'_t \\) 的ReLU平方值乘以权重矩阵 \\( W_{v'} \\) 得到。\n", + "- \\( o'_t \\) 是 \\( r'_t \\) 的激活函数 \\( \\sigma \\) 的输出与 \\( v'_t \\) 的逐元素乘积。\n", + "\n", + "其中，3.5D 指的是一种表示维度的方式。在深度学习模型中，D 通常代表模型的隐藏层维度（即嵌入维度或特征空间的维度）。例如，如果模型的隐藏维度是256，那么4D表示这个维度被扩展为4倍，也就是1024。\n", + "\n", + "然而，3.5D 是一个不常见的表示方法，通常情况下，我们会看到整数倍的表示（如2D, 4D等）。在这里，3.5D代表的是隐藏维度的3.5倍。\n", + "\n", + "具体来说，如果模型的基础维度是D，那么3.5D就表示：\n", + "\\begin{align*} 3.5D = 3.5 \\times D \\end{align*}\n", + "\n", + "假设D是256，那么3.5D就是：\n", + "\\begin{align*} 3.5 \\times 256 = 896 \\end{align*}\n", + "\n", + "所以，3.5D就是指模型在特定层中使用的特征维度是基础维度的3.5倍。在这个文档中，作者提到从4D减少到3.5D，意味着他们减少了某个层或模块的特征维度，以便引入新的门控权重并保持参数数量的一致性。" + ] + }, + { + "cell_type": "markdown", + "id": "67a5c0d4-57d5-4f9b-a2b7-bddd2250f08a", + "metadata": {}, + "source": [ + "Eagle时间混合（Eagle Time Mixing）的公式及其操作方法如下：\n", + "\n", + "### 公式部分\n", + "\n", + "Eagle时间混合的公式如下：\n", + "\n", + "\\begin{align*}\n", + "\\Box_t &= \\text{lerp}_{\\Box}(x_t, x_{t-1}) W_{\\Box}, \\quad \\Box \\in \\{r, k, v, g\\} \\tag{4} \\\\\n", + "w &= \\exp(-\\exp(\\omega)) \\tag{5} \\\\\n", + "\\text{wk} \\mathbf{v}_t &= \\text{diag}(u) \\cdot k_t^\\top \\cdot v_t + \\sum_{i=1}^{t-1} \\text{diag}(w)^{t-1-i} \\cdot k_i^\\top \\cdot v_i \\in \\mathbb{R}^{(D/h) \\times (D/h)} \\tag{6} \\\\\n", + "o_t &= \\text{concat} \\left( \\text{SiLU}(g_t) \\odot \\text{LayerNorm}(r_t \\cdot \\text{wk} \\mathbf{v}_t) \\right) W_o \\in \\mathbb{R}^D \\tag{7}\n", + "\\end{align*}\n", + "\n", + "### 解释部分\n", + "\n", + "- **LayerNorm的操作**：LayerNorm在每个头部（head）上独立操作，这相当于在h个组上执行GroupNorm（Wu & He，2018）。值得注意的是，$w$ 是由 $\\omega \\in \\mathbb{R}^{D/h}$ 通过公式 $w = \\exp(-\\exp(\\omega))$ 计算得到的，$\\omega$ 是实际的头部可训练参数。这确保了 $w$ 在区间 (0,1) 内，从而保证 $\\text{diag}(w)$ 是一个收缩矩阵。\n", + "\n", + "- **wkv_t 计算**：wkv_t 的注意力计算可以用递归形式写为：\n", + " \\begin{align*}\n", + " \\text{wk} \\mathbf{v}' &= s + \\text{diag}(u) \\cdot k^\\top \\cdot v \\tag{8} \\\\\n", + " s' &= \\text{diag}(w) \\cdot s + k^\\top \\cdot v \\tag{9}\n", + " \\end{align*}\n", + "\n", + "- **解释RWKV的 wkv_t 项**：RWKV的 wk\\mathbf{v}_t 项可以被认为是归一化 $k^\\top v$ 项的基于衰减的等价物。值得注意的是，对于给定的头部 $j$，递归状态 $s$ 是 $k^\\top v$ 的和，其中 $s$ 的每个通道在每个时间步通过相应的 $w$ 通道单独衰减。在应用接受度向量、门控和输出权重之前，当前令牌的 $k^\\top v$ 被乘以一个每通道的学习提升 $u$ 并与状态相加，见图1右上角。这给当前令牌相对于包含在衰减状态历史中的过去令牌的和一个特殊的处理。接受度乘以这个和，类似于线性注意力中的查询项。\n" + ] + }, + { + "cell_type": "markdown", + "id": "78a86c2f-962d-4826-baf4-bc19bc40b6e3", + "metadata": {}, + "source": [ + "这里的最大的改进应该是现在的计算是分成了`H = self.n_head`个头，然后每个头的计算结果都被存到了state里。相比于RWKV-4，这种改进可以类比于Transformer的单头自注意力机制改到多头注意力机制。\n", + "```python\n", + " @MyFunction\n", + " def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_mix_g, time_first, time_decay, kw, vw, rw, gw, ow, ln_w, ln_b):\n", + " H = self.n_head\n", + " S = self.head_size\n", + "\n", + " i1 = (2+S)*i+1\n", + " xk = x * time_mix_k + state[i1] * (1 - time_mix_k)\n", + " xv = x * time_mix_v + state[i1] * (1 - time_mix_v)\n", + " xr = x * time_mix_r + state[i1] * (1 - time_mix_r)\n", + " xg = x * time_mix_g + state[i1] * (1 - time_mix_g)\n", + " state[i1] = x\n", + "\n", + " r = (rw @ xr).view(H, 1, S)\n", + " k = (kw @ xk).view(H, S, 1)\n", + " v = (vw @ xv).view(H, 1, S)\n", + " g = F.silu(gw @ xg)\n", + "\n", + " s = state[(2+S)*i+2:(2+S)*(i+1), :].reshape(H, S, S)\n", + "\n", + " x = torch.zeros(H, S)\n", + " a = k @ v\n", + " x = r @ (time_first * a + s)\n", + " s = a + time_decay * s\n", + " \n", + " state[(2+S)*i+2:(2+S)*(i+1), :] = s.reshape(S, -1)\n", + " x = x.flatten()\n", + "\n", + " x = F.group_norm(x.unsqueeze(0), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).squeeze(0) * g # same as gn(x/8, eps=1e-5)\n", + " return ow @ x\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "bd093a96-fdc5-460d-b39f-fe3735795b42", + "metadata": {}, + "outputs": [], + "source": [ + "class RWKV_RNN(MyModule):\n", + " def __init__(self, args):\n", + " super().__init__()\n", + " self.args = args\n", + " self.eval() # set torch to inference mode\n", + " \n", + " w = torch.load(args.MODEL_NAME + '.pth', map_location='cpu')\n", + " for k in w.keys():\n", + " w[k] = w[k].float() # convert to f32 type\n", + " if '.time_' in k: w[k] = w[k].squeeze()\n", + " if '.time_decay' in k: w[k] = torch.exp(-torch.exp(w[k])).unsqueeze(-1)\n", + " if '.time_faaaa' in k: w[k] = w[k].unsqueeze(-1)\n", + "\n", + " self.n_head = w['blocks.0.att.time_decay'].shape[0]\n", + " self.head_size = w['blocks.0.ln1.weight'].shape[0] // self.n_head\n", + " \n", + " self.w = types.SimpleNamespace() # set self.w from w\n", + " self.w.blocks = {}\n", + " for k in w.keys(): # example: \"blocks.0.att.time_first\" => self.w.blocks[0].att.time_first\n", + " parts = k.split('.')\n", + " last = parts.pop()\n", + " here = self.w\n", + " for p in parts:\n", + " if p.isdigit():\n", + " p = int(p)\n", + " if p not in here: here[p] = types.SimpleNamespace()\n", + " here = here[p]\n", + " else:\n", + " if not hasattr(here, p): setattr(here, p, types.SimpleNamespace())\n", + " here = getattr(here, p)\n", + " setattr(here, last, w[k])\n", + "\n", + " def layer_norm(self, x, w):\n", + " return F.layer_norm(x, (self.args.n_embd,), weight=w.weight, bias=w.bias)\n", + "\n", + " @MyFunction\n", + " def channel_mixing(self, x, state, i:int, time_mix_k, time_mix_r, kw, vw, rw):\n", + " i0 = (2+self.head_size)*i+0\n", + " xk = x * time_mix_k + state[i0] * (1 - time_mix_k)\n", + " xr = x * time_mix_r + state[i0] * (1 - time_mix_r)\n", + " state[i0] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + " return r * (vw @ k)\n", + "\n", + " @MyFunction\n", + " def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_mix_g, time_first, time_decay, kw, vw, rw, gw, ow, ln_w, ln_b):\n", + " H = self.n_head\n", + " S = self.head_size\n", + "\n", + " i1 = (2+S)*i+1\n", + " xk = x * time_mix_k + state[i1] * (1 - time_mix_k)\n", + " xv = x * time_mix_v + state[i1] * (1 - time_mix_v)\n", + " xr = x * time_mix_r + state[i1] * (1 - time_mix_r)\n", + " xg = x * time_mix_g + state[i1] * (1 - time_mix_g)\n", + " state[i1] = x\n", + "\n", + " r = (rw @ xr).view(H, 1, S)\n", + " k = (kw @ xk).view(H, S, 1)\n", + " v = (vw @ xv).view(H, 1, S)\n", + " g = F.silu(gw @ xg)\n", + "\n", + " s = state[(2+S)*i+2:(2+S)*(i+1), :].reshape(H, S, S)\n", + "\n", + " x = torch.zeros(H, S)\n", + " a = k @ v\n", + " x = r @ (time_first * a + s)\n", + " s = a + time_decay * s\n", + " \n", + " state[(2+S)*i+2:(2+S)*(i+1), :] = s.reshape(S, -1)\n", + " x = x.flatten()\n", + "\n", + " x = F.group_norm(x.unsqueeze(0), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).squeeze(0) * g # same as gn(x/8, eps=1e-5)\n", + " return ow @ x\n", + "\n", + " def forward(self, token, state):\n", + " with torch.no_grad():\n", + " if state == None:\n", + " state = torch.zeros(self.args.n_layer * (2+self.head_size), self.args.n_embd)\n", + " \n", + " x = self.w.emb.weight[token]\n", + " x = self.layer_norm(x, self.w.blocks[0].ln0)\n", + " for i in range(self.args.n_layer):\n", + " # print(i)\n", + " att = self.w.blocks[i].att\n", + " x = x + self.time_mixing(self.layer_norm(x, self.w.blocks[i].ln1), state, i, \n", + " att.time_mix_k, att.time_mix_v, att.time_mix_r, att.time_mix_g, att.time_faaaa, att.time_decay, \n", + " att.key.weight, att.value.weight, att.receptance.weight, att.gate.weight, att.output.weight,\n", + " att.ln_x.weight, att.ln_x.bias)\n", + " ffn = self.w.blocks[i].ffn\n", + " x = x + self.channel_mixing(self.layer_norm(x, self.w.blocks[i].ln2), state, i, \n", + " ffn.time_mix_k, ffn.time_mix_r, \n", + " ffn.key.weight, ffn.value.weight, ffn.receptance.weight)\n", + " \n", + " x = self.w.head.weight @ self.layer_norm(x, self.w.ln_out)\n", + " return x.float(), state" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "a330cd34-7ed0-4a6c-92a3-19797d34ee77", + "metadata": {}, + "outputs": [], + "source": [ + "# context = \"Q:Do you know datawhalechina?\\nA:\"\n", + "context = '\\nQ:DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. How do you think of it?'" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "ad824161-413d-460c-9ffe-9dbfb739f86b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'/data1/ckw/RWKV-5-World-0.4B-v2-20231113-ctx4096'" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "args.MODEL_NAME" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "f0e2f841-4cda-48d4-b055-7adf00f2fe73", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(24, 1024)" + ] + }, + "execution_count": 34, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "args.n_layer,args.n_embd" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "id": "aba8a4d4-9a77-4191-a7ef-d5e6100ca3c1", + "metadata": {}, + "outputs": [], + "source": [ + "# args.n_layer = 24\n", + "# args.n_embd = 1024" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "id": "dd44f7bc-e8d6-4242-beb5-89a866990751", + "metadata": {}, + "outputs": [], + "source": [ + "# args.n_layer = 12\n", + "# args.n_embd = 768" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "id": "2a96d9dc-8b5e-40cc-bb36-24c9bdeac29e", + "metadata": {}, + "outputs": [], + "source": [ + "# args.MODEL_NAME='../models/rwkv-5-world-1b5'" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "id": "b7d07606-31b4-4c21-9f89-554d89c2c866", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Using CPU. Loading /data1/ckw/RWKV-5-World-0.4B-v2-20231113-ctx4096 ...\n", + "\n", + "Preprocessing context (slow version. see v2/rwkv/model.py for fast version)\n" + ] + } + ], + "source": [ + "print(f'\\nUsing CPU. Loading {args.MODEL_NAME} ...')\n", + "model = RWKV_RNN(args)\n", + "\n", + "print(f'\\nPreprocessing context (slow version. see v2/rwkv/model.py for fast version)')\n", + "init_state = None" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "id": "ce42cfad-0274-4d5d-950d-fb89ff11ed2c", + "metadata": {}, + "outputs": [], + "source": [ + "init_state = None" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "id": "3e02a81c-1447-4936-a241-4d00ecf8e862", + "metadata": {}, + "outputs": [], + "source": [ + "LENGTH_PER_TRIAL=1024" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "id": "4a00ea05-d6fd-4052-b13a-8107fb268420", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\n", + "--[ Trial 0 ]----------------- \n", + "Q:DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. How do you think of it?\n", + "QI: I think that the group of students is actually the whole AI community.\n", + "Q: In the first episode, how do you think you, a student, can use AI to solve a problem?\n", + "QI: It's a great opportunity to help develop and build knowledge, so that if we see AI problems, we can help solve them.\n", + "Q: How do you think that students can also participate in the teaching of AI?\n", + "QI: It is very important to let the students to think that there is an AI problem, and we can solve it by teaching AI.\n", + "Q: How do you think the research that we did on AI can be used to develop AI technologies?\n", + "QI: The research is interesting and it can be used to develop AI technologies.\n", + "Q: Do you think that students can learn from your research?\n", + "QI: I think so.\n", + "Q: You also talk about the use of AI in real-life applications. What do you think of that?\n", + "QI: I think it's a good thing to see.\n", + "Q: What are the major challenges that you see as being faced by the AI community?\n", + "QI: One is how to find data that can help us solve problems. The other is how to find a good dataset.\n", + "Q: You also talk about how we should deal with the big data problem. How do you think about that?\n", + "QI: We should not think that it is impossible to handle big data. There are a lot of big data, but there is a problem of how to handle them.\n", + "Q: What is the role of AI in industry?\n", + "QI: AI plays an important role in industry. AI has helped us improve the quality of services. We have a lot of new applications that we are using AI to solve.\n", + "Q: How do you think about AI and humans in the future?\n", + "QI: AI is not just for humans. It is also used for us to learn, for example.\n", + "Q: What do you think about the use of AI in the field of tourism?\n", + "QI: It's not that easy to use AI in tourism. There are so many problems.\n", + "Q: Do you think that AI will be a part of tourism in the future?\n", + "QI: I think so. It is very important for us to see.\n", + "Q: What do you think about AI and data sharing?\n", + "QI: It is not that easy to use AI in data sharing.\n", + "Q: What are the ways that you see AI in tourism?\n", + "QI: AI can be used to solve problems.\n", + "Q: How do you think about the relationship between AI and data?\n", + "QI: We need to use AI in the future to help us solve problems.\n", + "Q: How do you think about the relationship between AI and data sharing?\n", + "Q: What do you think about the future of AI?\n", + "Q: What are the main issues that AI is facing?\n", + "Q: What are the biggest challenges that you see in the field of AI?\n", + "Q: What do you think about the future of AI?\n", + "Q: What are the biggest challenges that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main challenges that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "Q: What are the main trends that you see in the field of AI?\n", + "\n", + "--[ Trial 1 ]----------------- \n", + "Q:DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. How do you think of it?\n", + "M: We are always looking for data to make sure that we are doing the right thing. We are currently looking at how to do this through the webinars and learning events. We have had speakers from different areas, such as from Silicon Valley, who have participated in the series. The current speaker, Marco Aurelio, was from Hong Kong. He was doing a presentation on Artificial Intelligence.\n", + "Q:How do you think of the audience that you are aiming to reach?\n", + "M: We are aiming at the general audience. We are also targeting people in the financial industry, who are also interested in artificial intelligence.\n", + "Q:What are your biggest challenges?\n", + "M: One of the biggest challenges is that the audience is very educated. They know about artificial intelligence and data. But the difficulty is that we have to explain the whole technology to them.\n", + "Q:How do you see the future of artificial intelligence?\n", + "M: It is an interesting future. It is really interesting. We are starting to see many different developments. The technology is really getting better and better. There are different ways of data that are being created. We have the development of machines to pick words and sentences and machines to make the machines think.\n", + "Q:How do you see the future of Artificial Intelligence?\n", + "M: We are constantly working on how to make the future of artificial intelligence more human-like.\n", + "Tags: dataWhalechina\n", + "Previous PostFuture is one of the hottest topics in Artificial Intelligence\n", + "Next PostOpinions about the future of Artificial Intelligence are changing\n", + "Cotton Developer News: Hands-On With Artificial Intelligence\n", + "Headlines from the data Whalechina Network: October 6, 2019\n", + "DataWhalechina Network: July 30, 2019\n", + "Cotton Developer News: July 22, 2019\n", + "Headlines from the data Whalechina Network: June 22, 2019\n", + "DataWhalechina Network: May 18, 2019\n", + "Archives Select Month July 2019 June 2019 May 2019 April 2019 March 2019 February 2019 January 2019 December 2018 November 2018 October 2018 September 2018 August 2018 July 2018 June 2018 May 2018 April 2018 March 2018 February 2018 January 2018 December 2017 November 2017 October 2017 September 2017 August 2017 July 2017 June 2017 May 2017 April 2017 March 2017 February 2017 January 2017 December 2016 November 2016 October 2016 September 2016 August 2016 July 2016 June 2016 May 2016 April 2016 March 2016 February 2016 January 2016 December 2015 November 2015 October 2015 September 2015 August 2015 July 2015 June 2015 May 2015 April 2015 March 2015 February 2015 January 2015 December 2014 November 2014 October 2014 September 2014 August 2014 July 2014 June 2014 May 2014 April 2014 March 2014 February 2014 January 2014 December 2013 November 2013 October 2013 September 2013 August 2013 July 2013 June 2013 May 2013 April 2013 March 2013 February 2013 January 2013 December 2012 November 2012 October 2012 September 2012 August 2012 July 2012 June 2012 May 2012 April 2012 March 2012 February 2012 January 2012 December 2011 November 2011 October 2011 September 2011 August 2011 July 2011 June 2011 May 2011 April 2011 March 2011 February 2011 January 2011 December 2010 November 2010 October 2010 September 2010 August 2010 July 2010 June 2010 May 2010 April 2010 March 2010 February 2010 January 2010 December 2009 November 2009 October 2009 September 2009 August 2009 July 2009 June 2009 May 2009 April 2009 March 2009 February 2009 January 2009 December 2008 November 2008 October 2008 September 2008 August 2008 July 2008 June 2008 May 2008 April 2008 March 2008 February 2008 January 2008 December 2007 November 2007 October 2007 September 2007 August 2007 July 2007 June 2007 May 2007 April 2007 March 2007 February 2007 January 2007 December 2006 November 2006 October 2006 September 2006 August 2006 July 2006 June 2006 May 2006 April 2006 March 2006 February 2006 January 2006 December 2005 November 2005 October 2005 September 2005 August 2005 July 2005 June 2005 May 2005 April 2005 March 2005 February 2005 January 2005 December 2004 November 2004 October 2004 September 2004 August 2004 July 2004 June 2004 May 2004 April 2004 March 2004 February 2004 January 2004 December 2003 November 2003 October 2003 September 2003 August 2003 July 2003 June 2003 May 2003 April 2003 March 2003 February 2003 January 2003 December 2002 November 2002 October 2002 September 2002 August 2002 July 2002\n", + "\n", + "--[ Trial 2 ]----------------- \n", + "Q:DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. How do you think of it?\n", + "Q:As AI continues to grow, what are some of the most promising applications of artificial intelligence?\n", + "Q:How do you think artificial intelligence will affect the future of AI?\n", + "Q:How does AI's role in education differ from the way it was used in the past?\n", + "Q:What are some of the challenges AI will face in the future?\n", + "Q:What is your vision for AI?\n", + "Q:What are your key takeaways from this conference?\n", + "Q:What do you hope to accomplish with AI?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are your current trends and plans for AI?\n", + "Q:How can AI be applied in education?\n", + "Q:What do you think will be the biggest impact of AI on education?\n", + "Q:What is your vision for AI in the future?\n", + "Q:How does AI change the way we teach and learn?\n", + "Q:What are your hopes for the future of AI?\n", + "Q:How does AI's role in education differ from the way it was used in the past?\n", + "Q:What are some of the challenges AI will face in the future?\n", + "Q:What is your vision for the future of AI?\n", + "Q:What are your key takeaways from this conference?\n", + "Q:What is your vision for the future of AI?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are your hopes for the future of AI?\n", + "Q:What are your key takeaways from this conference?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q:What is your vision for AI's role in education?\n", + "Q:What are your goals for the future of AI?\n", + "Q:What are some of the biggest challenges AI will face in the future?\n", + "Q\n", + "\n" + ] + } + ], + "source": [ + "for token in tokenizer.encode(context):\n", + " init_out, init_state = model.forward(token, init_state)\n", + "\n", + "for TRIAL in range(NUM_TRIALS):\n", + " print(f'\\n\\n--[ Trial {TRIAL} ]-----------------', context, end=\"\")\n", + " all_tokens = []\n", + " out_last = 0\n", + " out, state = init_out.clone(), init_state.clone()\n", + " for i in range(LENGTH_PER_TRIAL):\n", + " token = sample_logits(out, TEMPERATURE, TOP_P)\n", + " all_tokens += [token]\n", + " try:\n", + " tmp = tokenizer.decode(all_tokens[out_last:])\n", + " if '\\ufffd' not in tmp: # only print when we have a valid utf-8 string\n", + " print(tmp, end=\"\", flush=True)\n", + " out_last = i + 1\n", + " except:\n", + " pass\n", + " out, state = model.forward(token, state) \n", + "print('\\n')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3d3eaf3-252a-43da-9414-e1c6f6c681fc", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kewei-ai", + "language": "python", + "name": "kewei-ai" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/Model_Architecture_Discussions/rwkv-v5/RWKV_v5_demo.ipynb b/Model_Architecture_Discussions/rwkv-v5/RWKV_v5_demo.ipynb index 96b8401..da62f99 100644 --- a/Model_Architecture_Discussions/rwkv-v5/RWKV_v5_demo.ipynb +++ b/Model_Architecture_Discussions/rwkv-v5/RWKV_v5_demo.ipynb @@ -21,6 +21,16 @@ "MyFunction = torch.jit.script_method" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "0e6c9297-472f-4fd8-ad19-d8072b5040f8", + "metadata": {}, + "outputs": [], + "source": [ + "rwkv5又叫eagal" + ] + }, { "cell_type": "code", "execution_count": 2, diff --git a/Model_Architecture_Discussions/rwkv-v5/img/01.png b/Model_Architecture_Discussions/rwkv-v5/img/01.png new file mode 100644 index 0000000..4f03eb4 Binary files /dev/null and b/Model_Architecture_Discussions/rwkv-v5/img/01.png differ diff --git a/Model_Architecture_Discussions/rwkv-v6/RWKV-v6-guide.ipynb b/Model_Architecture_Discussions/rwkv-v6/RWKV-v6-guide.ipynb new file mode 100644 index 0000000..918b76b --- /dev/null +++ b/Model_Architecture_Discussions/rwkv-v6/RWKV-v6-guide.ipynb @@ -0,0 +1,514 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 5, + "id": "f64be1c0-02a8-4ea9-ae05-85b66e803cac", + "metadata": {}, + "outputs": [], + "source": [ + "########################################################################################################\n", + "# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM\n", + "########################################################################################################\n", + "\n", + "import numpy as np\n", + "np.set_printoptions(precision=4, suppress=True, linewidth=200)\n", + "import types, torch\n", + "import torch.nn as nn\n", + "from torch.nn import functional as F\n", + "\n", + "MyModule = torch.jit.ScriptModule\n", + "MyFunction = torch.jit.script_method" + ] + }, + { + "cell_type": "markdown", + "id": "12f7fb18-23e7-44a1-883e-0bd673b8fb0b", + "metadata": {}, + "source": [ + "![](./img/01.png)\n", + "\n", + "图1：RWKV架构概述。左侧：时间混合和通道混合块；右上角：作为RNN单元的RWKV时间混合块；中下部：前馈模块中的令牌移位模块和Eagle时间混合；右下角：Finch时间混合中的令牌移位模块。所有形状注释为简洁起见假设为单头。虚线箭头（左侧，右上角）表示在Finch中有连接，但在Eagle中没有。" + ] + }, + { + "cell_type": "markdown", + "id": "b3d66b1f-5c04-44d7-9bbc-76f844707327", + "metadata": {}, + "source": [ + "首先RWKV 6相比于RWKV 5在Token Shift上进行了改进，具体看下面的中间底部和右下角的图，分别是RWKV 4/5的Token Shift方式和RWKV 6的Token Shift方式。" + ] + }, + { + "cell_type": "markdown", + "id": "20cd8fa2-c1a4-4fdc-ba5d-858b25df1bcf", + "metadata": {}, + "source": [ + "具体内容如下：\n", + "\n", + "### 公式部分\n", + "\n", + "Finch Token Shift中使用的数据依赖线性插值（ddlerp）定义如下：\n", + "\n", + "\\begin{align*}\n", + "\\text{lora}_{\\Box}(x) &= \\lambda_{\\Box} + \\tanh(x A_{\\Box}) B_{\\Box} \\tag{14} \\\\\n", + "\\text{ddlerp}_{\\Box}(a, b) &= a + (b - a) \\odot \\text{lora}_{\\Box}(a + (b - a) \\odot \\mu_{x}) \\tag{15}\n", + "\\end{align*}\n", + "\n", + "### 解释部分\n", + "\n", + "- **可学习向量和矩阵**：\n", + " - $\\mu_{x}$ 和每个 $\\lambda_{\\Box}$ 引入了维度为 $D$ 的可训练向量。\n", + " - $A_{\\Box} \\in \\mathbb{R}^{D \\times 32}$ 和 $B_{\\Box} \\in \\mathbb{R}^{32 \\times D}$ 引入了新的可训练权重矩阵。\n", + " - 对于公式中提到的LoRA$_{\\omega}$的特殊情况，引入了双倍大小的可训练权重矩阵：$A_{\\omega} \\in \\mathbb{R}^{D \\times 64}$ 和 $B_{\\omega} \\in \\mathbb{R}^{64 \\times D}$。\n", + "\n", + "- **未来模型扩展**：\n", + " - 图1中右下角显示了一个示意图。\n", + " - 未来7B及更大规模的Finch模型预计将进一步增加这些权重矩阵的大小（可能翻倍或更多）。\n", + "\n", + "### 功能与作用\n", + "\n", + "这种带有数据依赖性的Token Shift新形式旨在扩展模型超越RWKV-4/Eagle风格的Token Shift的能力，使得每个通道分配的新旧数据量现在依赖于当前和前一个时间步的输入。\n", + "\n", + "### 详细解释\n", + "\n", + "- **数据依赖线性插值（ddlerp）**：\n", + " - ddlerp通过公式14和公式15实现，它结合了当前时间步和前一个时间步的信息来计算插值。\n", + " - $\\text{lora}_{\\Box}(x)$利用了一个$\\lambda_{\\Box}$向量和通过$\\tanh$函数处理的$x A_{\\Box}$与$B_{\\Box}$的乘积来生成。\n", + "\n", + "- **模型能力扩展**：\n", + " - 通过这种数据依赖的Token Shift，Finch模型能够更灵活地处理时间步之间的信息传递，使得模型在处理复杂序列数据时更加精确和高效。\n", + "\n", + "总结来说，Finch在Token Shift上引入了数据依赖的线性插值，利用可训练的向量和矩阵来增强模型的灵活性和能力，使其能够更好地处理时间步之间的信息，从而提高了模型的整体性能。" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "1261d4e1-df4e-410b-a4fa-45452c4b6fb1", + "metadata": {}, + "outputs": [], + "source": [ + "class RWKV_TOKENIZER():\n", + " table: list[list[list[bytes]]]\n", + " good: list[set[int]]\n", + " wlen: list[int]\n", + " def __init__(self, file_name):\n", + " self.idx2token = {}\n", + " sorted = [] # must be already sorted\n", + " lines = open(file_name, \"r\", encoding=\"utf-8\").readlines()\n", + " for l in lines:\n", + " idx = int(l[:l.index(' ')])\n", + " x = eval(l[l.index(' '):l.rindex(' ')])\n", + " x = x.encode(\"utf-8\") if isinstance(x, str) else x\n", + " assert isinstance(x, bytes)\n", + " assert len(x) == int(l[l.rindex(' '):])\n", + " sorted += [x]\n", + " self.idx2token[idx] = x\n", + "\n", + " self.token2idx = {}\n", + " for k, v in self.idx2token.items():\n", + " self.token2idx[v] = int(k)\n", + "\n", + " # precompute some tables for fast matching\n", + " self.table = [[[] for j in range(256)] for i in range(256)]\n", + " self.good = [set() for i in range(256)]\n", + " self.wlen = [0 for i in range(256)]\n", + "\n", + " for i in reversed(range(len(sorted))): # reverse order - match longer tokens first\n", + " s = sorted[i]\n", + " if len(s) >= 2:\n", + " s0 = int(s[0])\n", + " s1 = int(s[1])\n", + " self.table[s0][s1] += [s]\n", + " self.wlen[s0] = max(self.wlen[s0], len(s))\n", + " self.good[s0].add(s1)\n", + "\n", + " def encodeBytes(self, src: bytes) -> list[int]:\n", + " src_len: int = len(src)\n", + " tokens: list[int] = []\n", + " i: int = 0\n", + " while i < src_len:\n", + " s: bytes = src[i : i + 1]\n", + "\n", + " if i < src_len - 1:\n", + " s1: int = int(src[i + 1])\n", + " s0: int = int(src[i])\n", + " if s1 in self.good[s0]:\n", + " sss: bytes = src[i : i + self.wlen[s0]]\n", + " try:\n", + " s = next(filter(sss.startswith, self.table[s0][s1]))\n", + " except:\n", + " pass\n", + " tokens.append(self.token2idx[s])\n", + " i += len(s)\n", + "\n", + " return tokens\n", + "\n", + " def decodeBytes(self, tokens):\n", + " return b''.join(map(lambda i: self.idx2token[i], tokens))\n", + "\n", + " def encode(self, src: str):\n", + " return self.encodeBytes(src.encode(\"utf-8\"))\n", + "\n", + " def decode(self, tokens):\n", + " return self.decodeBytes(tokens).decode('utf-8')\n", + "\n", + " def printTokens(self, tokens):\n", + " for i in tokens:\n", + " s = self.idx2token[i]\n", + " try:\n", + " s = s.decode('utf-8')\n", + " except:\n", + " pass\n", + " print(f'{repr(s)}{i}', end=' ')\n", + " # print(repr(s), i)\n", + " print()\n", + "\n", + "########################################################################################################" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "725bc55e-7f3f-4c1c-9664-ad84bf68e943", + "metadata": {}, + "outputs": [], + "source": [ + "#采样方式没有变化\n", + "def sample_logits(out, temperature=1.0, top_p=0.8):\n", + " probs = F.softmax(out, dim=-1).numpy()\n", + " sorted_probs = np.sort(probs)[::-1]\n", + " cumulative_probs = np.cumsum(sorted_probs)\n", + " cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])\n", + " probs[probs < cutoff] = 0\n", + " if temperature != 1.0:\n", + " probs = probs.pow(1.0 / temperature)\n", + " probs = probs / np.sum(probs)\n", + " out = np.random.choice(a=len(probs), p=probs)\n", + " return out\n", + "\n", + "########################################################################################################" + ] + }, + { + "cell_type": "raw", + "id": "812fac97-a6b8-423c-831d-fe7397883437", + "metadata": {}, + "source": [ + "模型下载地址：https://hf-mirror.com/BlinkDL/rwkv-6-world/resolve/main/RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "434ff88e-b94e-4f8b-86a3-7fefca32cddb", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer = RWKV_TOKENIZER(\"./rwkv_vocab_v20230424.txt\")\n", + "\n", + "args = types.SimpleNamespace()\n", + "args.MODEL_NAME = '/data1/ckw/RWKV-x060-World-1B6-v2.1-20240328-ctx4096'\n", + "args.n_layer = 24\n", + "args.n_embd = 2048\n", + "args.vocab_size = 65536\n", + "\n", + "context = \"\\nDatawhale is \"\n", + "# context = \"\\n我们发现\"\n", + "NUM_TRIALS = 3\n", + "LENGTH_PER_TRIAL = 100\n", + "TEMPERATURE = 1.0\n", + "TOP_P = 0.7" + ] + }, + { + "cell_type": "markdown", + "id": "ee395717-e599-40fd-a4b9-ddfc800babea", + "metadata": {}, + "source": [ + "相比于RWKV 5的Channel Mixing（见下面代码）来说，RWKV6的Channel Mixing没有变化，这里的`time_maa_k`和RWKV 5中的`time_mix_k`是相同形状的可学习参数，都是一个维度为D（模型的隐藏层维度）的张量。" + ] + }, + { + "cell_type": "markdown", + "id": "54571fa3-52c6-4b05-ab97-9f06d76a522e", + "metadata": {}, + "source": [ + "Finch在时间混合（Time Mixing）上做了以下改进，具体内容如下：\n", + "\n", + "### 公式部分\n", + "\n", + "Finch时间混合的公式如下：\n", + "\n", + "\\begin{align*}\n", + "\\Box_t &= \\text{lerp}_{\\Box}(x_t, x_{t-1}) W_{\\Box}, \\quad \\Box \\in \\{r, k, v, g\\} \\tag{16} \\\\\n", + "d_t &= \\text{lora}_d(\\text{ddlerp}_d(x_t, x_{t-1})) \\tag{17} \\\\\n", + "w_t &= \\exp(-\\exp(d_t)) \\tag{18} \\\\\n", + "\\text{wk} \\mathbf{v}_t &= \\text{diag}(u) \\cdot k_t^\\top \\cdot v_t + \\sum_{i=1}^{t-1} \\left( \\prod_{j=1}^{i-1} w_j \\right) \\cdot k_i^\\top \\cdot v_i \\in \\mathbb{R}^{(D/h) \\times (D/h)} \\tag{19} \\\\\n", + "o_t &= \\text{concat} \\left( \\text{SiLU}(g_t) \\odot \\text{LayerNorm}(r_t \\cdot \\text{wk} \\mathbf{v}_t) \\right) W_o \\in \\mathbb{R}^D \\tag{20}\n", + "\\end{align*}\n", + "\n", + "### 解释部分\n", + "\n", + "- **可学习向量和矩阵**：\n", + " - $\\Box_t$ 是通过线性插值（lerp）计算得到的，适用于接受度（receptance）、键（key）、值（value）和门控向量（gate vectors）。\n", + " - $d_t$ 是通过 $\\text{lora}_d$ 函数对 $\\text{ddlerp}_d(x_t, x_{t-1})$ 进行处理得到的。\n", + " - $w_t$ 是由 $d_t$ 计算得到的，用于控制衰减的动态变化。\n", + "\n", + "- **时间混合计算**：\n", + " - $\\text{wk} \\mathbf{v}_t$ 是通过当前键值对 $k_t^\\top \\cdot v_t$ 和所有之前时间步的键值对 $k_i^\\top \\cdot v_i$ 的加权和计算得到的，权重由 $w_t$ 控制。\n", + " - 输出 $o_t$ 是通过连接（concat） $\\text{SiLU}(g_t)$ 和 $\\text{LayerNorm}(r_t \\cdot \\text{wk} \\mathbf{v}_t)$ 的结果得到的。\n", + "\n", + "- **递归形式**：\n", + " \\begin{align*}\n", + " \\text{wk} \\mathbf{v}' &= s + \\text{diag}(u) \\cdot k^\\top \\cdot v \\tag{21} \\\\\n", + " s' &= \\text{diag}(w) \\cdot s + k^\\top \\cdot v \\tag{22}\n", + " \\end{align*}\n", + "\n", + "### 功能与作用\n", + "\n", + "与Eagle不同，Finch中的 $w_t$ 不是在整个序列中固定的。每个通道的 $w_t$ 可以随时间动态变化，具体取决于数据输入，这也是Finch中衰减的核心变化。\n", + "\n", + "### 详细解释\n", + "\n", + "- **动态衰减**：\n", + " - Finch引入的数据依赖衰减使得每个通道的 $w_t$ 可以根据当前和之前的输入动态变化，而不是固定的学习向量。\n", + " - 这种动态衰减机制通过新的LoRA机制应用到学习向量上，增加了模型的灵活性。\n", + "\n", + "- **高级Token-Shift**：\n", + " - 新的时间衰减 $w_t$ 进一步应用了LoRA机制，允许每个通道的 $w_t$ 基于当前和之前的令牌混合来变化。\n", + "\n", + "总结来说，Finch在时间混合上通过引入数据依赖的动态衰减机制和高级Token-Shift，实现了更高的灵活性和精确度，使模型能够更好地处理和融合时间步之间的信息，从而提高了整体性能。" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "3869854a-a4e3-4652-9698-b0d81bbbd645", + "metadata": {}, + "outputs": [], + "source": [ + "class RWKV_RNN(MyModule):\n", + " def __init__(self, args):\n", + " super().__init__()\n", + " self.args = args\n", + " self.eval() # set torch to inference mode\n", + " \n", + " w = torch.load(args.MODEL_NAME + '.pth', map_location='cpu')\n", + "\n", + " for k in w.keys():\n", + " w[k] = w[k].float() # convert to f32 type\n", + " if '.time_' in k: w[k] = w[k].squeeze()\n", + " if '.time_faaaa' in k: w[k] = w[k].unsqueeze(-1)\n", + "\n", + " self.n_head = w['blocks.0.att.time_faaaa'].shape[0]\n", + " self.head_size = w['blocks.0.ln1.weight'].shape[0] // self.n_head\n", + " \n", + " self.w = types.SimpleNamespace() # set self.w from w\n", + " self.w.blocks = {}\n", + " for k in w.keys(): # example: \"blocks.0.att.time_first\" => self.w.blocks[0].att.time_first\n", + " parts = k.split('.')\n", + " last = parts.pop()\n", + " here = self.w\n", + " for p in parts:\n", + " if p.isdigit():\n", + " p = int(p)\n", + " if p not in here: here[p] = types.SimpleNamespace()\n", + " here = here[p]\n", + " else:\n", + " if not hasattr(here, p): setattr(here, p, types.SimpleNamespace())\n", + " here = getattr(here, p)\n", + " setattr(here, last, w[k])\n", + "\n", + " def layer_norm(self, x, w):\n", + " return F.layer_norm(x, (self.args.n_embd,), weight=w.weight, bias=w.bias)\n", + "\n", + " @MyFunction\n", + " def channel_mixing(self, x, state, i:int, time_maa_k, time_maa_r, kw, vw, rw):\n", + " i0 = (2+self.head_size)*i+0\n", + " sx = state[i0] - x\n", + " xk = x + sx * time_maa_k\n", + " xr = x + sx * time_maa_r\n", + " state[i0] = x\n", + " r = torch.sigmoid(rw @ xr)\n", + " k = torch.square(torch.relu(kw @ xk)) # square relu, primer paper\n", + " return r * (vw @ k)\n", + "\n", + " @MyFunction\n", + " def time_mixing(self, x, state, i:int, x_maa, w_maa, k_maa, v_maa, r_maa, g_maa, tm_w1, tm_w2, td_w1, td_w2, time_first, time_decay, kw, vw, rw, gw, ow, ln_w, ln_b):\n", + " H = self.n_head\n", + " S = self.head_size\n", + "\n", + " i1 = (2+S)*i+1\n", + " sx = state[i1] - x\n", + " state[i1] = x\n", + " xxx = x + sx * x_maa\n", + " xxx = torch.tanh(xxx @ tm_w1).view(5, 1, -1)\n", + " xxx = torch.bmm(xxx, tm_w2).view(5, -1)\n", + " mw, mk, mv, mr, mg = xxx.unbind(dim=0)\n", + "\n", + " xw = x + sx * (w_maa + mw)\n", + " xk = x + sx * (k_maa + mk)\n", + " xv = x + sx * (v_maa + mv)\n", + " xr = x + sx * (r_maa + mr)\n", + " xg = x + sx * (g_maa + mg)\n", + "\n", + " w = (time_decay + (torch.tanh(xw @ td_w1) @ td_w2).float()).view(H, S, 1)\n", + " w = torch.exp(-torch.exp(w.float()))\n", + "\n", + " r = (rw @ xr).view(H, 1, S)\n", + " k = (kw @ xk).view(H, S, 1)\n", + " v = (vw @ xv).view(H, 1, S)\n", + " g = F.silu(gw @ xg)\n", + "\n", + " s = state[(2+S)*i+2:(2+S)*(i+1), :].reshape(H, S, S)\n", + "\n", + " x = torch.zeros(H, S)\n", + " a = k @ v\n", + " x = r @ (time_first * a + s)\n", + " s = a + w * s\n", + " \n", + " state[(2+S)*i+2:(2+S)*(i+1), :] = s.reshape(S, -1)\n", + " x = x.flatten()\n", + "\n", + " x = F.group_norm(x.unsqueeze(0), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).squeeze(0) * g # same as gn(x/8, eps=1e-5)\n", + " return ow @ x\n", + "\n", + " def forward(self, token, state):\n", + " with torch.no_grad():\n", + " if state == None:\n", + " state = torch.zeros(self.args.n_layer * (2+self.head_size), self.args.n_embd)\n", + " \n", + " x = self.w.emb.weight[token]\n", + " x = self.layer_norm(x, self.w.blocks[0].ln0)\n", + " for i in range(self.args.n_layer):\n", + " att = self.w.blocks[i].att\n", + " x = x + self.time_mixing(self.layer_norm(x, self.w.blocks[i].ln1), state, i,\n", + " att.time_maa_x, att.time_maa_w, att.time_maa_k, att.time_maa_v, att.time_maa_r, att.time_maa_g, att.time_maa_w1, att.time_maa_w2,\n", + " att.time_decay_w1, att.time_decay_w2, att.time_faaaa, att.time_decay,\n", + " att.key.weight, att.value.weight, att.receptance.weight, att.gate.weight, att.output.weight,\n", + " att.ln_x.weight, att.ln_x.bias)\n", + " ffn = self.w.blocks[i].ffn\n", + " x = x + self.channel_mixing(self.layer_norm(x, self.w.blocks[i].ln2), state, i, \n", + " ffn.time_maa_k, ffn.time_maa_r, \n", + " ffn.key.weight, ffn.value.weight, ffn.receptance.weight)\n", + " \n", + " x = self.w.head.weight @ self.layer_norm(x, self.w.ln_out)\n", + " return x.float(), state" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "5235be83-a574-41f6-8546-bc415e2aeacf", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Using CPU. Loading /data1/ckw/RWKV-x060-World-1B6-v2.1-20240328-ctx4096 ...\n", + "\n", + "Preprocessing context (slow version. see v2/rwkv/model.py for fast version)\n" + ] + } + ], + "source": [ + "print(f'\\nUsing CPU. Loading {args.MODEL_NAME} ...')\n", + "model = RWKV_RNN(args)\n", + "\n", + "print(f'\\nPreprocessing context (slow version. see v2/rwkv/model.py for fast version)')\n", + "init_state = None" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "035ce374-c1c0-43b8-9ba9-37df297baae6", + "metadata": {}, + "outputs": [], + "source": [ + "for token in tokenizer.encode(context):\n", + " init_out, init_state = model.forward(token, init_state)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "1e17d6ef-c02c-4d27-8cf6-9a262f75f77f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\n", + "--[ Trial 0 ]----------------- \n", + "Datawhale is ➡️‼️\n", + "https://twitter.com/datawhale_cn/status/1463997087819689985\n", + "#Data #AI #DataAnalytics #AIOps #DataOps #MachineLearning #DataScience #DataLakeAnalytics #Hadoop #Amazon #Google #AWS #Azure #Dataprep #DevOps #OSS #Linux #Unix #BigData #BigDataOps #DataArchitecture #DataScienceOps #MachineLearningOps\n", + "\n", + "--[ Trial 1 ]----------------- \n", + "Datawhale is 🤓\n", + "\n", + "--[ Trial 2 ]----------------- \n", + "Datawhale is 🤯. They have a solid team, a really good SaaS product and the tools to support their users. That said, I have to take a serious look at the privacy and security of their platform before I buy into their story. I think this is a case of big companies buying into the hype, and they're not taking into account all the realities that go into building a privacy-focused product.\n", + "P.S. You can still apply to Datawhale's Program.\n", + "\n" + ] + } + ], + "source": [ + "for TRIAL in range(NUM_TRIALS):\n", + " print(f'\\n\\n--[ Trial {TRIAL} ]-----------------', context, end=\"\")\n", + " all_tokens = []\n", + " out_last = 0\n", + " out, state = init_out.clone(), init_state.clone()\n", + " for i in range(LENGTH_PER_TRIAL):\n", + " token = sample_logits(out, TEMPERATURE, TOP_P)\n", + " all_tokens += [token]\n", + " try:\n", + " tmp = tokenizer.decode(all_tokens[out_last:])\n", + " if '\\ufffd' not in tmp: # only print when we have a valid utf-8 string\n", + " print(tmp, end=\"\", flush=True)\n", + " out_last = i + 1\n", + " except:\n", + " pass\n", + " out, state = model.forward(token, state) \n", + "print('\\n')" + ] + }, + { + "cell_type": "markdown", + "id": "172c33e0-6d5b-4143-b85d-86777a2f5739", + "metadata": {}, + "source": [ + "v6和v5相比，感觉更喜欢使用emoj了哈哈" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kewei-ai", + "language": "python", + "name": "kewei-ai" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/Model_Architecture_Discussions/rwkv-v6/RWKV_v6_demo.ipynb b/Model_Architecture_Discussions/rwkv-v6/RWKV_v6_demo.ipynb index ec76c2f..cba84e9 100644 --- a/Model_Architecture_Discussions/rwkv-v6/RWKV_v6_demo.ipynb +++ b/Model_Architecture_Discussions/rwkv-v6/RWKV_v6_demo.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 5, "id": "f64be1c0-02a8-4ea9-ae05-85b66e803cac", "metadata": {}, "outputs": [], @@ -23,7 +23,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 6, "id": "1261d4e1-df4e-410b-a4fa-45452c4b6fb1", "metadata": {}, "outputs": [], @@ -109,7 +109,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 7, "id": "725bc55e-7f3f-4c1c-9664-ad84bf68e943", "metadata": {}, "outputs": [], @@ -139,7 +139,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 13, "id": "434ff88e-b94e-4f8b-86a3-7fefca32cddb", "metadata": {}, "outputs": [], @@ -160,18 +160,6 @@ "TOP_P = 0.7" ] }, - { - "cell_type": "markdown", - "id": "b275ed7e-6708-4e5e-b76e-60c3a2a4a6b6", - "metadata": {}, - "source": [ - "首先RWKV 6相比于RWKV 5在Token Shift上进行了改进，具体看下面的中间底部和右下角的图，分别是RWKV 4/5的Token Shift方式和RWKV 6的Token Shift方式。\n", - "\n", - "![](./img/01.png)\n", - "\n", - "相比于RWKV 5的Channel Mixing（见下面）来说，RWKV6的Channel Mixing没有变化，这里的`time_maa_k`和RWKV 5中的`time_mix_k`是相同形状的可学习参数，都是一个维度为D（模型的隐藏层维度）的张量。" - ] - }, { "cell_type": "code", "execution_count": 14, diff --git a/Model_Architecture_Discussions/rwkv-v6/img/01.png b/Model_Architecture_Discussions/rwkv-v6/img/01.png index cc9cd19..4f03eb4 100644 Binary files a/Model_Architecture_Discussions/rwkv-v6/img/01.png and b/Model_Architecture_Discussions/rwkv-v6/img/01.png differ diff --git a/README.md b/README.md index 08f34a4..05ca42f 100644 --- a/README.md +++ b/README.md @@ -82,8 +82,12 @@ | --- | --- | --- | | ChatGLM3 | [chatglm3.ipynb](./Model_Architecture_Discussions/ChatGLM3/加载模型权重.ipynb) | [@Tangent-90C](https://github.com/Tangent-90C) | | Llama3 | [llama3.ipynb](./Model_Architecture_Discussions/llama3/llama3-from-scratch.ipynb) | [@A10-research](https://www.aaaaaaaaaa.org/) | -| RWKV V2 | [rwkv-v2.ipynb](./Model_Architecture_Discussions/rwkv-v2/rwkv-v2.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | -| RWKV V3 | [rwkv-v3.ipynb](./Model_Architecture_Discussions/rwkv-v3/rwkv-v3.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| RWKV V2 | [rwkv-v2.ipynb](./Model_Architecture_Discussions/rwkv-v2/rwkv-v2-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| RWKV V3 | [rwkv-v3.ipynb](./Model_Architecture_Discussions/rwkv-v3/rwkv-v3-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| RWKV V4 | [rwkv-v4.ipynb](./Model_Architecture_Discussions/rwkv-v4/rwkv-v4-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| RWKV V5 | [rwkv-v5.ipynb](./Model_Architecture_Discussions/rwkv-v5/rwkv-v5-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| RWKV V6 | [rwkv-v6.ipynb](./Model_Architecture_Discussions/rwkv-v6/rwkv-v6-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | +| ChatGLM4 | [rwkv-v3.ipynb](./Model_Architecture_Discussions/ChatGLM4/chatglm4-guide.ipynb) | [@Ethan-Chen-plus](https://github.com/Ethan-Chen-plus) | ---