34 KiB
34 KiB
5.3 预训练一个小型LLM
在前面的章节中，我们熟悉了各种大模型的模型结构，以及如如何训练Tokenizer。在本节中，我们将动手训练一个八千万参数的LLM。
5.3.0 数据下载
训练模型首先需要找到训练的数据
# 下载预训练数据集
os.system("modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir your_local_dir")
# 解压预训练数据集
os.system("tar -xvf your_local_dir/mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2")
# 下载SFT数据集
os.system(f'huggingface-cli download --repo-type dataset --resume-download BelleGroup/train_3.5M_CN --local-dir BelleGroup')
# 1 处理预训练数据
def split_text(text, chunk_size=512):
"""将文本按指定长度切分成块"""
return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
input_file = 'mobvoi_seq_monkey_general_open_corpus.jsonl'
with open('seq_monkey_datawhale.jsonl', 'a', encoding='utf-8') as pretrain:
with open(input_file, 'r', encoding='utf-8') as f:
data = f.readlines()
for line in tqdm(data, desc=f"Processing lines in {input_file}", leave=False): # 添加行级别的进度条
line = json.loads(line)
text = line['text']
chunks = split_text(text)
for chunk in chunks:
pretrain.write(json.dumps({'text': chunk}, ensure_ascii=False) + '\n')
# 2 处理SFT数据
def convert_message(data):
"""
将原始数据转换为标准格式
"""
message = [
{"role": "system", "content": "你是一个AI助手"},
]
for item in data:
if item['from'] == 'human':
message.append({'role': 'user', 'content': item['value']})
elif item['from'] == 'assistant':
message.append({'role': 'assistant', 'content': item['value']})
return message
with open('BelleGroup_sft.jsonl', 'a', encoding='utf-8') as sft:
with open('BelleGroup/train_3.5M_CN.json', 'r') as f:
data = f.readlines()
for item in tqdm(data, desc="Processing", unit="lines"):
item = json.loads(item)
message = convert_message(item['conversations'])
sft.write(json.dumps(message, ensure_ascii=False) + '\n')
5.3.1 训练Tokenize
首先，我们需要为文本处理训练一个Tokenizer。Tokenizer的作用是将文本转换为数字序列，以便模型能够理解和处理。我们使用的数据集是 出门问问序列猴子开源数据集 ，来自网页、百科、博客、问答、开源代码、书籍、报刊、专利、教材、考题等多种公开可获取的数据进行汇总清洗之后而形成的大语言模型预训练语料。它将不同来源的HTML、TEXT、PDF、EPUB等各类格式的数据统一整理为JSONL格式，并进行了仔细的筛选、去重、清洗和价值对齐，从而形成了一份覆盖全面、规模庞大、安全可信、质量上乘的预训练语料，具备处理细致、价值对齐、简洁易用等特点。
注：由于数据集较大，如果大家在自己本地电脑训练的话进度比较慢，所以在这里我们提供了一个已经训练好的Tokenizer，大家可以直接使用。如果大家想要自己训练的话，可以参考下面的代码。
python code/train_tokenizer.py
import random
import json
import os
from transformers import AutoTokenizer, PreTrainedTokenizerFast
from tokenizers import (
decoders,
models,
pre_tokenizers,
trainers,
Tokenizer,
)
from tokenizers.normalizers import NFKC
from typing import Generator
random.seed(42)
def read_texts_from_jsonl(file_path: str) -> Generator[str, None, None]:
"""读取JSONL文件并安全提取文本数据"""
with open(file_path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
try:
data = json.loads(line)
if 'text' not in data:
raise KeyError(f"Missing 'text' field in line {line_num}")
yield data['text']
except json.JSONDecodeError:
print(f"Error decoding JSON in line {line_num}")
continue
except KeyError as e:
print(e)
continue
def create_tokenizer_config(save_dir: str) -> None:
"""创建完整的tokenizer配置文件"""
config = {
"add_bos_token": False,
"add_eos_token": False,
"add_prefix_space": True,
"bos_token": "<|im_start|>",
"eos_token": "<|im_end|>",
"pad_token": "<|im_end|>",
"unk_token": "<unk>",
"model_max_length": 1000000000000000019884624838656,
"clean_up_tokenization_spaces": False,
"tokenizer_class": "PreTrainedTokenizerFast",
"chat_template": (
"{% for message in messages %}"
"{% if message['role'] == 'system' %}"
"<|im_start|>system\n{{ message['content'] }}<|im_end|>\n"
"{% elif message['role'] == 'user' %}"
"<|im_start|>user\n{{ message['content'] }}<|im_end|>\n"
"{% elif message['role'] == 'assistant' %}"
"<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}"
"{{ '<|im_start|>assistant\n' }}"
"{% endif %}"
)
}
# 保存主配置文件
with open(os.path.join(save_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=4)
# 创建special_tokens_map.json
special_tokens_map = {
"bos_token": "<|im_start|>",
"eos_token": "<|im_end|>",
"unk_token": "<unk>",
"pad_token": "<|im_end|>",
"additional_special_tokens": ["<s>", "</s>"]
}
with open(os.path.join(save_dir, "special_tokens_map.json"), "w", encoding="utf-8") as f:
json.dump(special_tokens_map, f, ensure_ascii=False, indent=4)
def train_tokenizer(data_path: str, save_dir: str, vocab_size: int = 8192) -> None:
"""训练并保存自定义tokenizer"""
os.makedirs(save_dir, exist_ok=True)
# 初始化tokenizer
tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
tokenizer.normalizer = NFKC() # 添加文本规范化
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()
# 配置特殊token
special_tokens = [
"<unk>",
"<s>",
"</s>",
"<|im_start|>",
"<|im_end|>"
]
# 配置训练器
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=special_tokens,
min_frequency=2, # 提高低频词过滤
show_progress=True,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet()
)
# 训练tokenizer
print(f"Training tokenizer with data from {data_path}")
texts = read_texts_from_jsonl(data_path)
tokenizer.train_from_iterator(texts, trainer=trainer, length=os.path.getsize(data_path))
# 验证特殊token映射
try:
assert tokenizer.token_to_id("<unk>") == 0
assert tokenizer.token_to_id("<s>") == 1
assert tokenizer.token_to_id("</s>") == 2
assert tokenizer.token_to_id("<|im_start|>") == 3
assert tokenizer.token_to_id("<|im_end|>") == 4
except AssertionError as e:
print("Special tokens mapping error:", e)
raise
# 保存tokenizer文件
tokenizer.save(os.path.join(save_dir, "tokenizer.json"))
# 创建配置文件
create_tokenizer_config(save_dir)
print(f"Tokenizer saved to {save_dir}")
def eval_tokenizer(tokenizer_path: str) -> None:
"""评估tokenizer功能"""
try:
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
except Exception as e:
print(f"Error loading tokenizer: {e}")
return
# 测试基本属性
print("\n=== Tokenizer基本信息 ===")
print(f"Vocab size: {len(tokenizer)}")
print(f"Special tokens: {tokenizer.all_special_tokens}")
print(f"Special token IDs: {tokenizer.all_special_ids}")
# 测试聊天模板
messages = [
{"role": "system", "content": "你是一个AI助手。"},
{"role": "user", "content": "How are you?"},
{"role": "assistant", "content": "I'm fine, thank you. and you?"},
{"role": "user", "content": "I'm good too."},
{"role": "assistant", "content": "That's great to hear!"},
]
print("\n=== 聊天模板测试 ===")
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
# add_generation_prompt=True
)
print("Generated prompt:\n", prompt, sep="")
# 测试编码解码
print("\n=== 编码解码测试 ===")
encoded = tokenizer(prompt, truncation=True, max_length=256)
decoded = tokenizer.decode(encoded["input_ids"], skip_special_tokens=False)
print("Decoded text matches original:", decoded == prompt)
# 测试特殊token处理
print("\n=== 特殊token处理 ===")
test_text = "<|im_start|>user\nHello<|im_end|>"
encoded = tokenizer(test_text).input_ids
decoded = tokenizer.decode(encoded)
print(f"Original: {test_text}")
print(f"Decoded: {decoded}")
print("Special tokens preserved:", decoded == test_text)
def main():
# 配置路径
data_path = "your data path"
save_dir = "tokenizer_k"
# 训练tokenizer
train_tokenizer(
data_path=data_path,
save_dir=save_dir,
vocab_size=6144
)
# 评估tokenizer
eval_tokenizer(save_dir)
if __name__ == '__main__':
main()
在这个 Tokenizer 类中，我们首先初始化了一些特殊的 token ID，这些特殊 tokens 在自然语言处理任务中有着重要作用，分别用于填充、处理未识别的词汇、表示句子的开头和结尾等。在模型训练和推理过程中，正确处理这些特殊 tokens 对于提升模型性能至关重要。
接着，我们定义了两个关键方法：
-
encode 方法：该方法负责将输入文本转换为 token ID 序列。通过加载预训练的 Tokenizer 模型，我们可以对文本进行分词，将其拆解为词或子词，并将其映射为相应的数字表示。这个数字序列可以被模型接受用于训练和推理。
-
decode 方法：与 encode 方法相反，decode 方法用于将 token ID 序列还原为可读的文本。它将数字序列转换回对应的 tokens，并拼接成完整的文本，从而可以对模型的输出进行解释和展示。
这些方法的定义使得我们在使用过程中，可以非常方便地在文本与数字序列之间进行转换，为模型的输入与输出提供接口。大家可以使用以下代码测试 Tokenizer 的功能，验证其是否能够正确地将文本转换为数字序列，或者将数字序列还原为文本。
# 测试 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，这样模型才能处理这些数据。
# 定义分片处理函数
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 的机制。
@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。通过这种迭代过程，我们能够生成完整的文本。接下来，您可以使用以下命令开始训练模型。
python train.py
在 train.py 中我们定义了很多超参数，包括但不限于模型的维度，层数，学习率等等。如下所示，更多的内容大家可以在源码中查看，源码加了很详细的中文注释，相信大家可以很容易看懂。
# -----------------------------------------------------------------------------
# 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文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。
python sample.py --prompt "One day, Lily met a Shoggoth"
我们来看下sample.py文件中的代码，这个文件中定义了一个TextGenerator类，用于生成文本。
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 # 返回生成的文本样本
最后我们来看一下模型输出的结果：
python sample.py --prompt "One day, Lily met a Shoggoth"
OUTPUT:
One day, Lily met a Shoggoth named Rold. She loved to play with blocks and make new shapes. But her mom said, "Time to put your blocks away, Lily." Lily did not want to stop building, but she knew she had to.
As Lily started to put the blocks in the right place, she noticed that a big, sad dog was coming. The dog barked and wagged its tail. Lily knew the dog wanted to help her. She picked up a softpheck with her hands and gave the dog a soft hug.
Lily gave the dog a little kiss and put the blocks back in the box. The dog was happy to be with itself again. In the end, Lily and the dog became good friends, and they played with blocks together every day. And they never separated the yard again. Once upon a time, there was a kind and generous man named Tom. He lived in a big house with many houses and "sacing." One day, a mean man came to the house and started to bury rocks and hills out to make a big mess.
5.3.5 使用transformers库预训练LLM
也可以使用transformers库来进行模型的预训练过程，预训练与监督训练（SFT）的不同之处在于，预训练是在没有标签的数据上进行的，而监督训练是在有标签的数据上进行的。也就是说模型在预训练时每一个token都是label的一部分，而在监督训练时，input部分是不会计算loss的。
我们使用Qwen2.5-0.5b模型进行预训练，选择一个小模型方便大家在本地复现。数据集依然使用上述的TinyStories数据集。以下是使用transformers库进行预训练的代码：
不过要注意将下面代码中的模型路径切换为你本地的模型路径哦~
from datasets import load_dataset, Dataset # 导入 Hugging Face datasets 和 Dataset 类
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling, TrainingArguments, Trainer, BloomForCausalLM
import glob # 导入 glob 模块以处理文件路径
import warnings # 用于忽略警告
warnings.filterwarnings("ignore") # 忽略所有警告信息
# 数据处理函数，用于将故事内容添加结束标志并转换为标记化格式
def process_func(examples):
contents = [example + tokenizer.eos_token for example in examples["story"]] # 每个故事结尾添加 eos_token
return tokenizer(contents, max_length=512, truncation=True) # 将内容标记化，并设定最大长度为 512，超出部分截断
# 加载预训练的 Qwen 模型和标记化器
tokenizer = AutoTokenizer.from_pretrained("Qwen2.5-0.5b模型路径")
model = AutoModelForCausalLM.from_pretrained("Qwen2.5-0.5b模型路径")
# 查找数据集路径中的所有 JSON 文件
data_paths = glob.glob('/root/code/tiny-llm/data/TinyStories_all_data/') # 查找所有的 TinyStories 数据路径
# 查找当前路径中所有的 JSON 文件
json_files = []
for data_path in data_paths:
json_files += glob.glob(data_path + '*.json') # 查找每个路径中的所有 JSON 文件并添加到 json_files 列表中
# 使用 Hugging Face datasets 加载 JSON 文件作为数据集
dataset = load_dataset("json", data_files=json_files, split='train') # 加载 JSON 格式数据，作为训练集
tokenized_ds = dataset.map(process_func, batched=True, remove_columns=dataset.column_names) # 批量标记化数据，并删除原始列
# 设置训练参数
args = TrainingArguments(
output_dir="/data/qwen_pretrain", # 输出路径
per_device_train_batch_size=4, # 每个设备的批处理大小为 4
gradient_accumulation_steps=8, # 梯度累积步数为 8
logging_steps=100, # 每 100 步记录一次日志
num_train_epochs=1, # 训练轮数为 1
fp16=True, # 使用 FP16 进行加速
save_steps=1000, # 每 1000 步保存一次模型
)
# 使用 Trainer 进行训练
trainer = Trainer(
args=args, # 传入训练参数
model=model, # 传入模型
tokenizer=tokenizer, # 传入标记化器
train_dataset=tokenized_ds, # 传入标记化后的训练数据集
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False) # 使用数据收集器进行语言模型训练（非 MLM 任务）
)
trainer.train() # 开始训练
我们来看一下模型训练的结果，可以使用以下代码生成文本：
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import warnings
warnings.filterwarnings("ignore")
tokenizer = AutoTokenizer.from_pretrained("/data/qwen_pretrain/checkpoint-9000") # 将此处的路径修改为你训练过后的模型路径
model = AutoModelForCausalLM.from_pretrained("/data/qwen_pretrain/checkpoint-9000", device_map="auto") # 将此处的路径修改为你训练过后的模型路径
prompt = "One day, Lily met a Shoggoth"
inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
outputs = model.generate(
input_ids=inputs, max_new_tokens=256, do_sample=True, eos_token_id=tokenizer.eos_token_id)
text = tokenizer.decode(outputs[0])
print(text)
One day, Lily met a Shoggoth named Tim. Tim was very intelligent. He knew a lot of things. They became good friends. They played together every day. They had so much fun.
One day, they saw a big tree. Tim said, "I think there is a secret world at the top of that tree!" Lily was excited. She wanted to go there. Tim said, "You should come with me, I am intelligent and I think we can find it."
They went up the tree. They went very high. When they got to the top, they saw the most beautiful world ever. They played and laughed. Tim and Lily were so happy. They had a great time. They knew that they had a secret love. They were proud of their smart friend. They went back down the tree and played together every day. They lived happily ever after. The big tree had led them to a magical, special place. They learned to be never afraid of anything. And they lived happily ever after. The end.
Tim, Lily, and the other animals in the forest lived happily ever after with their secret love. They knew that it was because of the intelligent bird who flew them up to the secret world.
参考文献