From c579aff59d165aaa41eeeb3df004298f65a8740f Mon Sep 17 00:00:00 2001 From: KMnO4-zx <1021385881@qq.com> Date: Sun, 22 Sep 2024 16:00:36 +0800 Subject: [PATCH] docs: ch05 Done --- README.md | 21 +- docs/chapter5/5.1 模型结构-LLaMA.md | 3 + ...练 Tokenizer.md => 5.2 训练 Tokenizer.md} | 8 +- docs/chapter5/5.3 训练一个小型LLM.md | 485 ++++++++++++++++++ 4 files changed, 503 insertions(+), 14 deletions(-) rename docs/chapter5/{5.3 训练 Tokenizer.md => 5.2 训练 Tokenizer.md} (98%) create mode 100644 docs/chapter5/5.3 训练一个小型LLM.md diff --git a/README.md b/README.md index 39cff2d..c1f71c8 100644 --- a/README.md +++ b/README.md @@ -82,25 +82,26 @@ - 4.2.2 SFT - 4.2.3 RLHF -### 第五章 预训练一个 LLM 志学 +### 第五章 预训练一个 LLM 志学 Done - 5.1 模型架构-LLaMA Done - 5.1.1 LLaMA Attention - 5.1.2 LLaMA Decoder Layer - 5.1.3 LLaMA MLP - 5.1.4 LLaMA RMSNorm - 5.1.5 A Whole LLaMA - - 5.2 预训练数据 - - 5.2.1 预训练数据集 - - 5.2.2 预训练数据处理 - - 5.3 训练 Tokenizer - - 5.3.1 Word-based Tokenizer - - 5.3.2 Character-based Tokenzier - - 5.3.3 Subword Tokenizer + - 5.2 训练 Tokenizer + - 5.2.1 Word-based Tokenizer + - 5.2.2 Character-based Tokenzier + - 5.2.3 Subword Tokenizer - （1）BPE - （2）Word Piece - （3）Unigram - - 5.3.4 训练一个 Tokenizer - - 5.4 预训练 + - 5.2.4 训练一个 Tokenizer + - 5.3 训练一个小型LLM + - 5.3.1 训练Tokenizer + - 5.3.2 数据预处理 + - 5.3.3 训练模型 + - 5.3.4 使用模型生成文本 ### 第六章 微调 LLM - 6.1 微调数据 diff --git a/docs/chapter5/5.1 模型结构-LLaMA.md b/docs/chapter5/5.1 模型结构-LLaMA.md index f5a538a..9da0722 100644 --- a/docs/chapter5/5.1 模型结构-LLaMA.md +++ b/docs/chapter5/5.1 模型结构-LLaMA.md @@ -564,3 +564,6 @@ Number of parameters: 15191712 torch.Size([1, 1, 32000]) ``` +**参考文献** + +- [llama2.c](https://github.com/karpathy/llama2.c) \ No newline at end of file diff --git a/docs/chapter5/5.3 训练 Tokenizer.md b/docs/chapter5/5.2 训练 Tokenizer.md similarity index 98% rename from docs/chapter5/5.3 训练 Tokenizer.md rename to docs/chapter5/5.2 训练 Tokenizer.md index 411f520..474c65a 100644 --- a/docs/chapter5/5.3 训练 Tokenizer.md +++ b/docs/chapter5/5.2 训练 Tokenizer.md @@ -1,4 +1,4 @@ -# 5.3 训练 Tokenizer +# 5.2 训练 Tokenizer 在自然语言处理 (NLP) 中，Tokenizer 是一种将文本分解为较小单位（称为 token）的工具。这些 token 可以是词、子词、字符，甚至是特定的符号。Tokenization 是 NLP 中的第一步，直接影响后续处理和分析的效果。不同类型的 tokenizer 适用于不同的应用场景，以下是几种常见的 tokenizer 及其特点。 @@ -14,7 +14,7 @@ Output: ["Hello", ",", "world", "!", "There", "is", "Datawhale", "."] 在这个例子中，输入的句子被分割成一系列单词和标点符号，每个单词或标点符号都作为一个独立的 token。 -## 5.3.2 Character-based Tokenizer +## 5.2.2 Character-based Tokenizer **Character-based Tokenizer** 将文本中的每个字符视为一个独立的 token。这种方法能非常精细地处理文本，适用于处理拼写错误、未登录词或新词。由于每个字符都是一个独立的 token，因此这种方法可以捕捉到非常细微的语言特征。这对于一些特定的应用场景，如生成式任务或需要处理大量未登录词的任务，特别有用。但是，这种方法也会导致 token 序列变得非常长，增加了模型的计算复杂度和训练时间。此外，字符级的分割可能会丢失一些词级别的语义信息，使得模型难以理解上下文。 @@ -26,7 +26,7 @@ Output: ["H", "e", "l", "l", "o"] 在这个例子中，单词“Hello”被分割成单个字符，每个字符作为一个独立的 token。这种方法能够处理任何语言和字符集，具有极大的灵活性。 -## 5.3.3 Subword Tokenizer +## 5.2.3 Subword Tokenizer **Subword Tokenizer** 介于词和字符之间，能够更好地平衡分词的细粒度和处理未登录词的能力。Subword Tokenizer 的关键思想是将文本分割成比单词更小的单位，但又比字符更大，这样既能处理未知词，又能保持一定的语义信息。常见的子词分词方法包括 BPE、WordPiece 和 Unigram。 @@ -74,7 +74,7 @@ Output: ["new", "est"] 每种 Tokenizer 方法都有其特定的应用场景和优缺点，选择适合的 Tokenizer 对于自然语言处理任务的成功至关重要。 -## 5.3.4 训练一个 Tokenizer +## 5.2.4 训练一个 Tokenizer 这里我们选择使用 BPE 算法来训练一个 Subword Tokenizer。BPE 是一种简单而有效的分词方法，能够处理未登录词和罕见词，同时保持较小的词典大小。我们将使用 Hugging Face 的 `tokenizers` 库来训练一个 BPE Tokenizer。 diff --git a/docs/chapter5/5.3 训练一个小型LLM.md b/docs/chapter5/5.3 训练一个小型LLM.md new file mode 100644 index 0000000..886e7f8 --- /dev/null +++ b/docs/chapter5/5.3 训练一个小型LLM.md @@ -0,0 +1,485 @@ +# 5.3 + +在前面的章节中，我们熟悉了各种大模型的模型结构，以及如如何训练Tokenizer。在本节中，我们将动手训练一个小型的LLM。 + +## 5.3.1 训练Tokenizer + +首先，我们需要为文本处理训练一个Tokenizer。Tokenizer的作用是将文本转换为数字序列，以便模型能够理解和处理。我们使用的数据集是 [TinyStory](https://www.modelscope.cn/datasets/AI-ModelScope/TinyStories) ，它是一个由GPT-3.5和GPT-4生成的小型故事数据集，包含简短的故事，且词汇量有限。在这个任务中，我们采用字符级Tokenizer，将文本中的每个字符映射为对应的数字。通过以下命令可以下载数据集并训练Tokenizer。 + +```bash +python train_vocab.py --download True --vocab_size 4096 +``` + +LLaMA2 的词表大小为 32,000，但由于 TinyStory 数据集较小，词汇量有限，我们将词表大小设置为 4,096。训练完成后，我们得到的 Tokenizer 能够将文本转换为数字序列，也可以将数字序列还原为文本。 + +```python +def download_file(url: str, fname: str, chunk_size=1024): + """发送HTTP GET请求以流式方式获取文件""" + ··· + +def download(): + """执行 download_file 下载数据集""" + ··· + +def train_vocab(vocab_size: int=32000, num_shards: int=20): + """ + vocab_size: int, 词汇表的大小，决定分词器的词汇量。 + num_shards: int, 用于加快词汇表训练的效率，指定要处理的分片数量。 + """ + # 确保词汇表大小为正数 + assert vocab_size > 0, "Vocab size must be positive" + + # SentencePiece 模型的前缀路径，将用于保存分词器 + prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}") + + # 1) 将多个分片中的文本导出为单个文本文件 tiny.txt + tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt") + data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") + shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json"))) + + # 创建 tiny.txt 文件并写入指定数量的分片中的文本 + print(f"Writing temporary file {tiny_file} with {num_shards} shards...") + with open(tiny_file, "w", encoding="utf-8") as of: + # 遍历前 num_shards 个分片 + for shard in tqdm(shard_filenames[:num_shards]): + with open(shard, "r") as f: + data = json.load(f) # 读取分片中的JSON数据 + # 遍历每个例子，将其中的故事文本写入 tiny.txt 文件 + for example in data: + text = example["story"] + text = text.strip() # 去除文本首尾的空白字符 + of.write(text + "\n") # 每个文本写入一行 + + # 输出生成的 tiny.txt 文件的大小 + print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB") + + # 2) 使用 SentencePiece 训练分词器 + print("Will now train the vocab...") + spm.SentencePieceTrainer.train( + input=tiny_file, # 输入文件为之前生成的 tiny.txt + model_prefix=prefix, # 模型前缀路径 + model_type="bpe", # 使用 Byte-Pair Encoding (BPE) 训练分词器 + vocab_size=vocab_size, # 词汇表大小 + self_test_sample_size=0, # 自测样本大小设置为 0 + input_format="text", # 输入文件格式为纯文本 + character_coverage=1.0, # 覆盖所有字符（包括非常见字符） + num_threads=os.cpu_count(), # 使用 CPU 的线程数 + split_digits=True, # 拆分数字 + allow_whitespace_only_pieces=True, # 允许仅由空格组成的词元 + byte_fallback=True, # 启用字节级回退 + unk_surface=r" \342\201\207 ", # UNK token 表示未知字符的方式 + normalization_rule_name="identity" # 使用“identity”归一化规则 + ) + + # 3) 可选的清理操作，询问用户是否删除临时文件 tiny.txt + dec = input(f"Delete the temporary file {tiny_file}? [y/N] ") + if dec.lower() == "y": + os.remove(tiny_file) # 删除临时文件 + print(f"Deleted {tiny_file}") + + # 输出模型保存的路径 + print(f"Trained tokenizer is in {prefix}.model") + print("Done.") +``` + +在本部分中，我们使用了 `SentencePiece` 库来训练自定义的 `Tokenizer`。首先，我们需要从 `TinyStory` 数据集中提取文本内容，作为训练的输入数据。`SentencePiece` 是一种基于子词单元的分词算法，能够有效处理不同语言中的词汇碎片化问题。 + +训练 `Tokenizer` 时，`SentencePiece` 会自动生成两个文件：`tok4096.model` 和 `tok4096.vocab`，其中 `tok4096.model` 是我们训练好的模型文件，位于 `data` 目录下。这个文件可以用于将文本数据转换为 `Token` 序列，也可以将 `Token` 序列还原为文本。 + +为了更便捷地使用这个 `Tokenizer`，我们还在 `tokenizer.py` 文件中定义了一个 `Tokenizer` 类。这个类封装了 `Tokenizer` 的常用操作，例如文本编码和解码功能，并支持加载我们训练好的模型文件。通过这个类，我们可以轻松地将文本转换为模型可接受的数字序列，或将预测结果转化为可读的文本。 + +具体的代码实现和细节可以在 `tokenizer.py` 文件中找到，接下来我们将进一步展示如何使用该类来处理 `TinyStory` 数据集中的故事文本。 + +```python +class Tokenizer: + def __init__(self, tokenizer_model=None): + """ + 初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。 + + 参数: + tokenizer_model: str, 可选，分词器模型的路径，如果不指定则使用默认路径 TOKENIZER_MODEL。 + """ + # 如果提供了分词器模型路径，使用该路径；否则使用默认模型路径 + model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL + # 确保模型文件存在 + assert os.path.isfile(model_path), model_path + + # 加载 SentencePiece 模型 + self.sp_model = SentencePieceProcessor(model_file=model_path) + self.model_path = model_path + + # 获取分词器的特殊token和词汇表大小 + self.n_words: int = self.sp_model.vocab_size() # 词汇表大小 + self.bos_id: int = self.sp_model.bos_id() # 句子开头 (BOS) 的ID + self.eos_id: int = self.sp_model.eos_id() # 句子结尾 (EOS) 的ID + self.pad_id: int = self.sp_model.pad_id() # 填充 (PAD) 的ID + + # 验证分词器词汇表大小是否正确 + assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() + + def encode(self, s: str, bos: bool, eos: bool) -> List[int]: + """ + 将字符串编码为词元ID列表。可以选择是否添加句子开头 (BOS) 和句子结尾 (EOS) 标记。 + + 参数: + s: str, 要编码的字符串。 + bos: bool, 是否在编码的词元列表前添加 BOS 标记。 + eos: bool, 是否在编码的词元列表末尾添加 EOS 标记。 + + 返回: + List[int]: 编码后的词元ID列表。 + """ + # 确保输入是字符串类型 + assert type(s) is str + # 使用SentencePiece将字符串编码为词元ID + t = self.sp_model.encode(s) + # 如果需要BOS标记，将其添加到词元列表开头 + if bos: + t = [self.bos_id] + t + # 如果需要EOS标记，将其添加到词元列表末尾 + if eos: + t = t + [self.eos_id] + return t + + def decode(self, t: List[int]) -> str: + """ + 将词元ID列表解码为字符串。 + + 参数: + t: List[int], 词元ID列表。 + + 返回: + str: 解码后的字符串。s + """ + return self.sp_model.decode(t) +``` + +在这个 `Tokenizer` 类中，我们首先初始化了一些特殊的 token ID，这些特殊 tokens 在自然语言处理任务中有着重要作用，分别用于填充、处理未识别的词汇、表示句子的开头和结尾等。在模型训练和推理过程中，正确处理这些特殊 tokens 对于提升模型性能至关重要。 + +接着，我们定义了两个关键方法： + +1. encode 方法：该方法负责将输入文本转换为 token ID 序列。通过加载预训练的 Tokenizer 模型，我们可以对文本进行分词，将其拆解为词或子词，并将其映射为相应的数字表示。这个数字序列可以被模型接受用于训练和推理。 + +2. decode 方法：与 encode 方法相反，decode 方法用于将 token ID 序列还原为可读的文本。它将数字序列转换回对应的 tokens，并拼接成完整的文本，从而可以对模型的输出进行解释和展示。 + +这些方法的定义使得我们在使用过程中，可以非常方便地在文本与数字序列之间进行转换，为模型的输入与输出提供接口。大家可以使用以下代码测试 `Tokenizer` 的功能，验证其是否能够正确地将文本转换为数字序列，或者将数字序列还原为文本。 + +```python +# 测试 Tokenizer +enc = Tokenizer('./data/tok4096.model') # 加载分词器 +tetx = 'Hello, world!' # 测试文本 +print(enc.encode(text, bos=True, eos=True)) # 编码文本 +print(enc.decode(enc.encode(text, bos=True, eos=True))) # 解码文本 + +OUTPUT: +[1, 346, 2233, 4010, 1475, 4021, 2] +Hello, world! +``` + +## 5.3.2 数据预处理 + +在训练模型之前，首先需要对数据进行预处理。这一步的核心任务是将文本数据转换为模型能够理解的数字序列。具体来说，文本中的每个字符、单词或子词都需要被映射为一个唯一的数字 ID，这样模型才能处理这些数据。 + +```python +# 定义分片处理函数 +def process_shard(args, vocab_size, tokenizer_model_path): + """ 处理数据分片，将其中的文本进行分词并保存为二进制文件 """ + ··· + + +# 定义预处理函数，用于对多个数据分片进行批量处理 +def pretokenize(vocab_size): + """ 预处理所有的数据分片，并将分词后的数据保存为二进制文件 """ + ··· + + +class PretokDataset(torch.utils.data.IterableDataset): + """从磁盘加载已预处理的分词数据，并将其以 PyTorch 张量的形式返回。""" + + def __init__(self, split, max_seq_len, vocab_size, vocab_source): + """ + 初始化数据集。 + + 参数: + split: str, 数据集的分割方式（'train' 或 'test'）。 + max_seq_len: int, 最大序列长度，用于生成输入输出序列。 + vocab_size: int, 词汇表的大小。 + vocab_source: str, 词汇表的来源（'llama2' 或 'custom'）。 + """ + super().__init__() + self.split = split # 数据集划分（训练集或测试集） + self.max_seq_len = max_seq_len # 最大序列长度 + self.vocab_size = vocab_size # 词汇表大小 + self.vocab_source = vocab_source # 词汇表来源 + + def __iter__(self): + """ + 返回迭代器，按批次加载数据并生成模型输入/输出。 + """ + # 获取DataLoader的worker信息（用于并行数据加载） + worker_info = torch.utils.data.get_worker_info() + worker_id = worker_info.id if worker_info else 0 # worker ID + # 获取分布式训练的rank信息（用于多GPU训练） + rank = dist.get_rank() if dist.is_initialized() else 0 + # 基于worker_id和rank生成唯一的随机数种子，确保数据在每个worker和rank之间是唯一的 + seed = 42 + worker_id + 1337 * rank + rng = random.Random(seed) + print(f"Created a PretokDataset with rng seed {seed}") + + # 根据词汇表来源决定数据路径 + if self.vocab_source == "llama2": + # 如果使用 Llama 2 词汇表，.bin 文件和 .json 文件在同一目录下 + bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") + shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin"))) + elif self.vocab_source == "custom": + # 如果使用自定义词汇表，.bin 文件在 tok{N} 目录下 + bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}") + shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin"))) + + # 根据数据集划分使用不同的分片文件 + # 训练集使用所有分片文件，测试集只使用第一个分片 + shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1] + assert len(shard_filenames) > 0, f"在 {bin_dir} 中未找到任何 .bin 文件" + + while True: + # 随机打乱分片文件 + rng.shuffle(shard_filenames) + for shard in shard_filenames: + # 使用 memmap 读取文件，使得数据留在磁盘上，减少内存占用 + m = np.memmap(shard, dtype=np.uint16, mode="r") + # 计算该分片中的批次数量 + num_batches = len(m) // self.max_seq_len + num_batches -= 1 # 去掉最后一个不完整的批次 + assert num_batches > 0, "这个分片文件太小了？请检查。" + # 随机打乱批次索引 + ixs = list(range(num_batches)) + rng.shuffle(ixs) + # 对每个批次生成输入 x 和目标输出 y + for ix in ixs: + start = ix * self.max_seq_len # 批次起始索引 + end = start + self.max_seq_len + 1 # 批次结束索引 + # 将数据转换为 NumPy 数组并拷贝到 RAM 中 + chunk = torch.from_numpy((m[start:end]).astype(np.int64)) + # 模型输入 x 是当前批次的前 max_seq_len 个词元 + x = chunk[:-1] + # 模型输出 y 是下一个词元 + y = chunk[1:] + # 生成 x, y 对 + yield x, y + + +class Task: + @staticmethod + def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs): + ds = PretokDataset(**dataset_kwargs) + dl = torch.utils.data.DataLoader( + ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers + ) + for x, y in dl: + x = x.to(device, non_blocking=True) + y = y.to(device, non_blocking=True) + yield x, y +``` + +在这部分中，首先定义了 `process_shard` 函数，用于处理数据分片。该函数的主要功能是将文本数据分词后，转换为更高效的二进制文件格式，以便后续更快速地加载和处理数据。 + +接下来，我们定义了 `pretokenize` 函数，用于批量处理多个数据分片。通过这一函数，所有数据可以并行处理，进一步加快预处理的速度。 + +然后，我们设计了一个 `PretokDataset` 类，用于加载已预处理好的数据集。我们继承了 `torch.utils.data.IterableDataset` 来定义该数据集，这使得我们可以更灵活、高效地处理数据。在这个类中，核心是 `__iter__` 方法，它负责生成用于训练的数据批次。 + +最后，我们还定义了一个 `Task` 类，专门用于迭代数据集，并生成模型所需的输入和目标输出。这一部分的设计确保了数据流的顺畅对接，为模型训练提供了标准化的数据输入。可以通过以下代码来测试预处理后的数据集。 + +## 5.3.3 训练模型 + +在数据预处理完成后，我们就可以开始训练模型了。我们使用的模型是一个和LLama2结构一样的 Decoder only Transformer模型，使用Pytorch实现。相关代码在`model.py`文件中。此处不再赘述，源码中有详细的中文注释，且我们在之前的文章中也有详细的介绍。 + +在模型这一部分可以重点看一下生成式模型是如何实现生成token的，可以查看`model.py`文件中的`Transforerm`类中的`generate`方法。 + +在完成数据预处理后，我们就可以开始训练模型了。我们使用的模型是一个与 LLaMA2 结构相同的 Decoder-only Transformer 模型，采用 PyTorch 实现。具体的实现细节已经包含在 `model.py` 文件中，在此不再赘述。该源码中包含详细的中文注释，此外我们在之前的文章中也对模型架构进行了深入介绍。 + +在模型部分，建议重点关注生成式模型如何生成 token 的过程。可以参考 `model.py` 文件中的 `Transformer` 类，尤其是 `generate` 方法的实现，它展示了模型如何基于已有的上下文生成后续 token 的机制。 + +```python +@torch.inference_mode() + def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): + """ + 给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。 + 在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。 + """ + for _ in range(max_new_tokens): + # 如果序列上下文过长，截断它到最大长度 + idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:] + + # 前向传播获取序列中最后一个位置的 logits + logits = self(idx_cond) + logits = logits[:, -1, :] # 只保留最后一个时间步的输出 + + if temperature == 0.0: + # 选择最有可能的索引 + _, idx_next = torch.topk(logits, k=1, dim=-1) + else: + # 缩放 logits 并应用 softmax + logits = logits / temperature + if top_k is not None: + v, _ = torch.topk(logits, min(top_k, logits.size(-1))) + logits[logits < v[:, [-1]]] = -float('Inf') + probs = F.softmax(logits, dim=-1) + idx_next = torch.multinomial(probs, num_samples=1) + + # 将采样的索引添加到序列中并继续 + idx = torch.cat((idx, idx_next), dim=1) + + return idx +``` + +在 generate 方法中，我们首先获取序列中最后一个位置的 logits，然后基于这些 logits 生成新的 token。接着，生成的新 token 会被添加到序列中，模型随后会继续生成下一个 token。通过这种迭代过程，我们能够生成完整的文本。接下来，您可以使用以下命令开始训练模型。 + +```bash +python train.py +``` + +在 `train.py` 中我们定义了很多超参数，包括但不限于模型的维度，层数，学习率等等。如下所示，更多的内容大家可以在源码中查看，源码加了很详细的中文注释，相信大家可以很容易看懂。 + +```python +# ----------------------------------------------------------------------------- +# I/O 配置，用于定义输出目录和训练时的日志记录与评估设置 +out_dir = "output" # 模型输出保存路径 +eval_interval = 2000 # 评估间隔步数 +log_interval = 1 # 日志记录间隔步数 +eval_iters = 100 # 每次评估时迭代的步数 +eval_only = False # 如果为True，脚本在第一次评估后立即退出 +always_save_checkpoint = False # 如果为True，在每次评估后总是保存检查点 +init_from = "scratch" # 可以选择从头开始训练（'scratch'）或从已有的检查点恢复（'resume'） + +# 数据配置 +batch_size = 8 # 每个微批次的样本数量，如果使用梯度累积，实际批次大小将更大 +max_seq_len = 256 # 最大序列长度 +vocab_size = 4096 # 自定义词汇表大小 + +# 模型配置 +dim = 288 # 模型的隐藏层维度 +n_layers = 8 # Transformer的层数 +n_heads = 8 # 注意力头的数量 +n_kv_heads = 4 # 模型分组 +multiple_of = 32 # 在某些层的维度必须是该数的倍数 +dropout = 0.0 # Dropout概率 + +# AdamW优化器配置 +gradient_accumulation_steps = 4 # 梯度累积步数，用于模拟更大的批次 +learning_rate = 5e-4 # 最大学习率 +max_iters = 100000 # 总的训练迭代次数 +weight_decay = 1e-1 # 权重衰减系数 +beta1 = 0.9 # AdamW优化器的β1参数 +beta2 = 0.95 # AdamW优化器的β2参数 +grad_clip = 1.0 # 梯度裁剪阈值，0表示不裁剪 + +# 学习率衰减配置 +decay_lr = True # 是否启用学习率衰减 +warmup_iters = 1000 # 学习率预热的步数 + +# 系统设置 +device = "cuda:0" # 设备选择：'cpu'，'cuda'，'cuda:0'等 +dtype = "bfloat16" # 数据类型：'float32'，'bfloat16'，'float16' +``` + +## 5.3.4 使用模型生成文本 + +在模型训练完成后，会在`output`目录下生成一个`ckpt.pt`文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。 + +```bash +python sample.py --prompt "One day, Lily met a Shoggoth" +``` + +我们来看下`sample.py`文件中的代码，这个文件中定义了一个`TextGenerator`类，用于生成文本。 + +```python +class TextGenerator: + def __init__(self, + checkpoint='output/ckpt.pt', # 模型检查点路径 + tokenizer_model_path='tok4096.model', # 分词器模型路径 + seed=1337, # 随机种子，确保可重复性 + device=None, # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU + dtype="float32"): # 数据类型，默认为 float32，可以选择 float16 或 bfloat16 + """ + 初始化 TextGenerator 类，加载模型、设置设备和分词器等。 + """ + # 模型加载配置 + self.checkpoint = checkpoint # 保存的模型检查点路径 + self.tokenizer_model_path = tokenizer_model_path # 分词器模型文件路径 + self.seed = seed # 随机数种子，用于生成的可重复性 + self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') # 根据硬件条件选择设备 + self.dtype = dtype # 模型的浮点数类型 + self.device_type = 'cuda' if 'cuda' in self.device else 'cpu' # 判断当前设备是否为 CUDA + + # 设置随机种子，确保生成的可重复性 + torch.manual_seed(seed) # 设置 CPU 随机种子 + torch.cuda.manual_seed(seed) # 设置 CUDA 随机种子 + torch.backends.cuda.matmul.allow_tf32 = True # 允许 CUDA 使用 TF32 精度进行矩阵乘法运算 + torch.backends.cudnn.allow_tf32 = True # 允许 cuDNN 使用 TF32 精度加速 + + # 根据 dtype 选择适当的自动混合精度上下文 + ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype] + self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype) + + # 加载模型检查点文件 + checkpoint_dict = torch.load(self.checkpoint, map_location=self.device) # 加载模型参数 + gptconf = ModelArgs(**checkpoint_dict['model_args']) # 初始化模型参数 + self.model = Transformer(gptconf) # 实例化 Transformer 模型 + state_dict = checkpoint_dict['model'] # 获取模型状态字典 + + # 去除状态字典中的不必要前缀 + unwanted_prefix = '_orig_mod.' # 这个前缀在保存时可能被添加，现在要去除它 + for k, v in list(state_dict.items()): + if k.startswith(unwanted_prefix): + state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) # 去除不必要的前缀 + + # 加载模型参数到模型中 + self.model.load_state_dict(state_dict, strict=False) + # 计算模型参数量 + num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + print(f"Model has {num_params} parameters.") + # 设置模型为评估模式（evaluation mode），防止训练模式下的 dropout 等操作影响结果 + self.model.eval() + # 将模型放置到正确的设备上（GPU 或 CPU） + self.model.to(self.device) + # 初始化分词器 + self.tokenizer = Tokenizer(tokenizer_model=self.tokenizer_model_path) # 根据指定的路径加载分词器 + + def sample(self, + start="Hello!", # 生成文本的起始提示词，可以是任意字符串 + num_samples=3, # 生成样本的数量，默认生成 3 个样本 + max_new_tokens=256, # 每个样本生成的最大 token 数，默认最多生成 256 个 token + temperature=1.0, # 控制生成的随机性，1.0 为标准，值越大越随机 + top_k=300): # 保留概率最高的 top_k 个 token，限制生成时的选择范围 + """ + 根据给定的起始文本生成样本。 + + :param start: 生成文本的起始提示词 + :param num_samples: 要生成的文本样本数 + :param max_new_tokens: 每个样本生成的最大 token 数 + :param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机 + :param top_k: 限制生成时选择的 token 范围 + :return: 生成的文本样本列表 + """ + # 如果 start 是以 'FILE:' 开头，表示从文件中读取起始文本 + if start.startswith('FILE:'): + with open(start[5:], 'r', encoding='utf-8') as f: + start = f.read() # 读取文件内容作为起始文本 + + # 将起始文本编码为 token id 序列 + start_ids = self.tokenizer.encode(start, bos=True, eos=False) # bos=True 表示加上句首标记，eos=False 表示不加句尾标记 + x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...]) # 将编码后的 token id 转为 PyTorch 张量 + + generated_texts = [] # 用于保存生成的文本样本 + with torch.no_grad(): # 禁用梯度计算，提升效率 + with self.ctx: # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时） + for k in range(num_samples): # 循环生成指定数量的样本 + y = self.model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k) # 生成文本 + generated_texts.append(self.tokenizer.decode(y[0].tolist())) # 解码生成的 token 序列为可读文本 + + return generated_texts # 返回生成的文本样本 +``` + +**参考文献** + +- [llama2.c](https://github.com/karpathy/llama2.c) \ No newline at end of file