本文介绍了如何在本地单张GPU上预训练一个Llama模型。预训练是指使用自监督学习目标在大量文本上训练仅解码器语言模型（如Llama），以区别于后续针对特定任务的微调步骤。你将学习如何准备训练数据并运行预训练。

概览

本文分为三个部分：

1. 训练带有特殊标记的分词器
2. 准备训练数据
3. 运行预训练

训练带有特殊标记的分词器

我们将使用的模型架构与之前文章创建的相同。这是一个12层的Llama模型，词汇表大小为50,000。用于预训练的数据集是 HuggingFaceFW/fineweb。

要准备训练数据，首先需要设置分词器。以下代码在 HuggingFaceFW/fineweb 数据集上训练一个BPE分词器，并将其保存到文件：

from typing import Iterator

import datasets
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, normalizers

# 加载 FineWeb 10B 样本（为演示节省内存，仅使用一小部分）
dataset = datasets.load_dataset("HuggingFaceFW/fineweb", "sample-10BT", split="train", streaming=True)

def get_texts(dataset: datasets.Dataset, limit: int = 100_000) -> Iterator[str]:
    """从数据集中获取文本，直到达到限制数量或数据集耗尽。"""
    count = 0
    for sample in dataset:
        yield sample["text"]
        count += 1
        if limit and count >= limit:
            break

# 初始化 BPE 模型
tokenizer = Tokenizer(models.BPE(byte_fallback=True, unk_token="[UNK]"))
tokenizer.normalizer = normalizers.NFKC()
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True, use_regex=False)
tokenizer.decoder = decoders.ByteLevel()

# 训练器
trainer = trainers.BpeTrainer(
    vocab_size=50_000,
    min_frequency=2,
    special_tokens=["[PAD]", "[BOT]", "[EOT]", "[UNK]"],
    show_progress=True,
)

# 训练并将分词器保存到磁盘
texts = get_texts(dataset, limit=100_000)
tokenizer.train_from_iterator(texts, trainer=trainer)
tokenizer.save("bpe_50k.json")

这个分词器在字节级别使用BPE算法。通常它不会产生未知标记，但我们仍为其设置了一个特殊标记。此外，我们还为文本开头、文本结尾和填充设置了特殊标记，这些对于下一个标记预测非常有用。

此代码会自动使用所有CPU核心。在高端计算机上运行此代码需要几分钟。训练好的分词器将保存到名为 bpe_50k.json 的文件中。训练完成后，可以使用以下代码重新加载它：

from tokenizers import Tokenizer

tokenizer = Tokenizer.from_file("bpe_50k.json")

请注意，我们用50,000的词汇量训练了分词器。这对于单语言模型来说非常有用。但是，如果要训练多语言模型，则需要更大的词汇量。

准备训练数据

预训练语言模型意味着预测序列中的下一个标记。对于训练数据，需要对文本进行分词，以创建一个整数标记ID的张量及其移位一个位置的版本作为预测目标。

如上一节所示，可以通过迭代数据集对象来加载数据集并将文本作为字符串打印出来：

dataset = datasets.load_dataset("HuggingFaceFW/fineweb", "sample-10BT", split="train")
for sample in dataset:
    print(sample["text"])
    break

与通常用于语言模型训练的数据集相比，这个数据集很小。然而，它仍然足够大，包含多样的人类语言样本。

对于预训练，需要创建一个PyTorch数据集对象，以便模型可以使用它，如下所示：

class PretrainingDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, tokenizer, seq_length, device):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.device = device
        self.seq_length = seq_length
        self.bot = tokenizer.token_to_id("[BOT]")
        self.eot = tokenizer.token_to_id("[EOT]")
        self.pad = tokenizer.token_to_id("[PAD]")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        """从数据集中获取一个标记ID序列。添加 [BOT] 和 [EOT] 标记。
        裁剪并填充到序列长度。
        """
        seq = self.dataset[index]["text"]
        tokens: list[int] = [self.bot] + self.tokenizer.encode(seq).ids + [self.eot]
        # 填充到目标序列长度
        toklen = len(tokens)
        if toklen < self.seq_length+1:
            pad_length = self.seq_length+1 - toklen
            tokens += [self.pad] * pad_length
        # 返回序列
        x = torch.tensor(tokens[:self.seq_length], dtype=torch.int64, device=self.device)
        y = torch.tensor(tokens[1:self.seq_length+1], dtype=torch.int64, device=self.device)
        return x, y

这是为预训练准备文本数据的最简单方法。我们封装了Hugging Face数据集对象，在 __len__ 方法中匹配样本数量。在 __getitem__ 方法中，将特定文本样本分词为整数标记ID的张量。添加文本开头和文本结尾标记有助于预训练：当只提供文本开头标记时，模型可以预测句子的第一个标记。当提供整个序列时，模型应预测结尾。

Transformer模型不限制传递的长度，除了位置编码能处理的最大序列长度。但是，当将多个序列作为一个批次传递时，需要确保所有序列具有相同的长度，以便可以将它们堆叠成单个张量。我们为较短的序列添加填充标记，并将较长的序列裁剪到目标序列长度。

预训练是自监督学习。预期输出的标签已在输入序列中。因此，将 x 设置为输入序列，并将其移位一个位置的版本作为目标序列 y。希望它们是PyTorch张量而不是Python列表，以便可以与PyTorch数据加载器一起使用。还必须将数据类型设置为int64，因为PyTorch的CrossEntropyLoss有此限制，需要在计算训练损失时识别填充标记。

可以通过创建数据加载器对象并从中抽取一个批次来测试数据集：

batch_size = 8
seq_length = 512
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
dataloader = torch.utils.data.DataLoader(
    PretrainingDataset(dataset, tokenizer, seq_length, device),
    batch_size=batch_size
)
for x, y in dataloader:
    print(x)
    print(y)
    break

运行预训练

一旦从数据集中准备好输入和目标数据，在语言模型上运行预训练与训练其他深度学习模型没有什么不同。

使用上一篇文章中的模型代码，首先创建一个模型对象：

# 使用默认配置创建预训练模型
model_config = LlamaConfig()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = LlamaForPretraining(model_config).to(device)

这是一个用于演示目的的小型模型。它只有1.71亿个参数，比你在互联网上找到的任何大型语言模型都要小得多。

接下来，应该定义训练参数。根据硬件情况，可能需要调整批量大小，但保持序列长度适中有助于模型学习上下文。以下是使用的策略：

• 该数据集只有一个训练拆分。为简单起见，数据不打乱，不创建验证集，训练循环也不包含任何评估步骤。
• 下一个标记预测是整个词汇表上的分类问题。自然地，损失函数是交叉熵。应确保填充标记不用于计算损失，因为它们不是有效的输入。
• 将序列长度设置为512。训练模型所需的资源与序列长度呈 (O(N^2)) 关系。因此，倾向于保持较短的长度，但太短的序列长度会阻止模型理解更长的上下文。
• 遵循训练大语言模型的最佳实践，使用带有预热期的余弦学习率调度器。预热期可以设置为固定的步数或占总训练步数的百分比（例如0.1%-2%）。这里将其设置为1000步。
• 一旦确定了序列长度，调整批量大小以适应GPU内存。可以从8开始，经验表明这适合12GB的显存。
• HuggingFaceFW/fineweb 10B 数据集有1400万个样本和100亿个标记，可能不需要训练多个周期。事实上，许多大语言模型在非常大的数据集上只训练1-3个周期。

将这些参数组合在一起来定义训练配置：

# 训练参数
epochs = 3
learning_rate = 1e-3
batch_size = 8
seq_length = 512
num_warmup_steps = 1000
PAD_TOKEN_ID = tokenizer.token_to_id("[PAD]")

# 数据加载器、优化器、调度器和损失函数
model.train()
dataloader = torch.utils.data.DataLoader(
    PretrainingDataset(dataset, tokenizer, seq_length, device),
    batch_size=batch_size
)
optimizer = torch.optim.AdamW(
    model.parameters(), lr=learning_rate, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.01
)
num_training_steps = len(dataloader) * epochs
warmup_scheduler = lr_scheduler.LinearLR(
    optimizer,
    start_factor=0.1, end_factor=1.0, total_iters=num_warmup_steps
)
cosine_scheduler = lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=num_training_steps - num_warmup_steps,
    eta_min=0
)
scheduler = lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_scheduler],
    milestones=[num_warmup_steps]
)
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN_ID)

AdamW优化器配置的峰值学习率为1e-3。其他参数设置为默认值。PyTorch的余弦调度器与线性调度器结合以实现预热期。它们使用 SequentialLR 调度器组合，并配置为在第1000步从线性调度切换到余弦调度。

请注意，在加载数据集进行训练时，我们没有设置 streaming=True，也没有打乱数据集。这使得数据加载器对象具有确定性。这样，可以轻松确定总训练步数，这有助于设置学习率调度器。

损失函数使用 nn.CrossEntropyLoss，并将填充标记ID设置为忽略索引。这意味着只要参考目标是填充标记，就不计算损失。这对于匹配上一节中创建数据集对象时定义的行为很重要。

按照大语言模型的标准，这是一个小型模型和小型数据集。然而，训练仍然非常慢。在单张GPU上从头开始训练将需要数百小时。重要的是可以对模型进行检查点保存并恢复训练。让我们在训练循环中实现这一点：

# 查找最后一个检查点
if os.path.exists("llama_pretraining_checkpoint.pth"):
    checkpoint = torch.load("llama_pretraining_checkpoint.pth")
    begin_epoch = checkpoint["epoch"]
    begin_batch = checkpoint["batch"]
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    del checkpoint
    print(f"从周期 {begin_epoch} 和批次 {begin_batch} 恢复训练")
else:
    begin_epoch = 0
    begin_batch = 0

# 开始训练
for epoch in range(begin_epoch, epochs):
    dataloader = torch.utils.data.DataLoader(
        PretrainingDataset(
            dataset.skip(begin_batch * batch_size),
            tokenizer,
            seq_length,
            device,
        ),
        batch_size=batch_size
    )
    pbar = tqdm.tqdm(dataloader, desc=f"周期 {epoch+1}/{epochs}")
    for batch_id, batch in enumerate(pbar):
        if (begin_batch + batch_id) % 1000 == 0:
            # 检查点保存模型和优化器状态
            torch.save({
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "epoch": epoch,
                "batch": batch_id + begin_batch,
            }, f"llama_pretraining_checkpoint.pth")
        # 获取批处理数据
        input_ids, target_ids = batch
        # 创建注意力掩码：因果掩码 + 填充掩码
        attn_mask = create_causal_mask(input_ids.shape[1], device) + \
                    create_padding_mask(input_ids, PAD_TOKEN_ID, device)
        # 从模型获取输出
        logits = model(input_ids, attn_mask)
        # 计算损失：logits和目标之间的交叉熵，忽略填充标记
        loss = loss_fn(logits.view(-1, logits.size(-1)), target_ids.view(-1))
        # 反向传播损失并应用梯度裁剪
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        pbar.set_postfix(loss=loss.item())
        pbar.update(1)
    begin_batch = 0   # 为下一个周期重置
    pbar.close()

当对训练进行检查点时，需要保存模型状态、优化器状态和调度器状态。还需要记住周期和批次索引，以便可以从数据集中的同一批次恢复。

使用 tqdm 库的进度条来可视化训练进度。在训练期间，从数据加载器对象中拉取一对输入和目标张量。datasets 库允许跳过任意数量的样本。使用它来创建一个数据加载器对象，以便从之前的检查点恢复。

然后，创建一个注意力掩码来屏蔽填充标记并启用因果掩码以控制自注意力机制。模型输出是一个三维张量，其批量大小和序列长度与输入相同。需要对其进行重塑以适应损失函数，然后使用计算出的损失更新模型。所有这些都是训练深度学习模型的标准流程。

最后，可以保存模型以便将其用于推理：

torch.save(model.state_dict(), "llama_pretraining_model.pth")
torch.save(model.base_model.state_dict(), "llama_model.pth")

根据用例，可能需要保存基础模型、预训练模型或两者。基础模型适用于其他任务，而预训练模型可用作生成模型。

为完整起见，以下是训练的完整代码：

import dataclasses
import os

import datasets
import tqdm
import tokenizers
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim.lr_scheduler as lr_scheduler
from torch import Tensor

# 加载分词器
tokenizer = tokenizers.Tokenizer.from_file("bpe_50K.json")

# 加载数据集
dataset = datasets.load_dataset("HuggingFaceFW/fineweb", "sample-10BT", split="train")

# 构建模型
@dataclasses.dataclass
class LlamaConfig:
    """定义 Llama 模型超参数。"""
    vocab_size: int = 50000  # 分词器词汇表大小
    max_position_embeddings: int = 2048  # 最大序列长度
    hidden_size: int = 768  # 隐藏层维度
    intermediate_size: int = 4*768  # MLP 隐藏层维度
    num_hidden_layers: int = 12  # Transformer 层数
    num_attention_heads: int = 12  # 注意力头数
    num_key_value_heads: int = 3  # GQA 的键值头数


def rotate_half(x: Tensor) -> Tensor:
    """旋转输入张量后半部分的隐藏维度。

    这是旋转位置编码的辅助函数。
    对于形状为 (..., d) 的张量，它返回一个张量，其中最后
    d/2 个维度通过交换和取反进行旋转。

    Args:
        x: 形状为 (..., d) 的输入张量

    Returns:
        形状相同、最后维度旋转后的张量
    """
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)  # 连接旋转后的结果


class RotaryPositionEncoding(nn.Module):
    """旋转位置编码。"""

    def __init__(self, dim: int, max_position_embeddings: int) -> None:
        """初始化 RotaryPositionEncoding 模块

        Args:
            dim: 应用 RoPE 的输入张量的隐藏维度
            max_position_embeddings: 输入张量的最大序列长度
        """
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        # 计算 n\theta_i 矩阵
        N = 10_000.0
        inv_freq = 1.0 / (N ** (torch.arange(0, dim, 2) / dim))
        inv_freq = torch.cat((inv_freq, inv_freq), dim=-1)
        position = torch.arange(max_position_embeddings)
        sinusoid_inp = torch.outer(position, inv_freq)
        # 将余弦和正弦矩阵保存为缓冲区，而不是参数
        self.register_buffer("cos", sinusoid_inp.cos())
        self.register_buffer("sin", sinusoid_inp.sin())

    def forward(self, x: Tensor) -> Tensor:
        """对张量 x 应用 RoPE

        Args:
            x: 形状为 (batch_size, seq_length, num_heads, head_dim) 的输入张量

        Returns:
            形状为 (batch_size, seq_length, num_heads, head_dim) 的输出张量
        """
        batch_size, seq_len, num_heads, head_dim = x.shape
        dtype = x.dtype
        # 将余弦和正弦矩阵转换为 4D 张量，并与 x 具有相同的数据类型
        cos = self.cos.to(dtype)[:seq_len].view(1, seq_len, 1, -1)
        sin = self.sin.to(dtype)[:seq_len].view(1, seq_len, 1, -1)
        # 对 x 应用 RoPE
        output = (x * cos) + (rotate_half(x) * sin)
        return output


class LlamaAttention(nn.Module):
    """带有旋转嵌入的分组查询注意力。"""

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_kv_heads = config.num_key_value_heads  # GQA: H_kv < H_q

        # hidden_size 必须能被 num_heads 整除
        assert (self.head_dim * self.num_heads) == self.hidden_size

        # 用于 Q、K、V 投影的线性层
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

    def forward(self, hidden_states: Tensor, rope: RotaryPositionEncoding, attn_mask: Tensor) -> Tensor:
        bs, seq_len, dim = hidden_states.size()

        # 将输入投影到 Q、K、V
        query_states = self.q_proj(hidden_states).view(bs, seq_len, self.num_heads, self.head_dim)
        key_states = self.k_proj(hidden_states).view(bs, seq_len, self.num_kv_heads, self.head_dim)
        value_states = self.v_proj(hidden_states).view(bs, seq_len, self.num_kv_heads, self.head_dim)

        # 应用旋转位置嵌入
        query_states = rope(query_states)
        key_states = rope(key_states)

        # 将张量从 BSHD 维度转置为 BHSD 维度，用于 scaled_dot_product_attention
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)

        # 使用 PyTorch 优化的注意力实现
        # 设置 is_causal=True 与设置显式注意力掩码不兼容
        attn_output = F.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=attn_mask,
            dropout_p=0.0,
            enable_gqa=True,
        )

        # 将输出张量从 BHSD 维度转置为 BSHD 维度，重塑为 3D，然后投影输出
        attn_output = attn_output.transpose(1, 2).reshape(bs, seq_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)
        return attn_output


class LlamaMLP(nn.Module):
    """带有 SwiGLU 激活的前馈网络。"""

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        # SwiGLU 的两个并行投影
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.act_fn = F.silu  # SwiGLU 激活函数
        # 投影回隐藏层大小
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        # SwiGLU 激活：将 gate 和 up-projected 输入相乘
        gate = self.act_fn(self.gate_proj(x))
        up = self.up_proj(x)
        return self.down_proj(gate * up)


class LlamaDecoderLayer(nn.Module):
    """Llama 模型的单个 Transformer 层。"""

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=1e-5)
        self.self_attn = LlamaAttention(config)
        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=1e-5)
        self.mlp = LlamaMLP(config)

    def forward(self, hidden_states: Tensor, rope: RotaryPositionEncoding, attn_mask: Tensor) -> Tensor:
        # 第一个残差块：自注意力
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_outputs = self.self_attn(hidden_states, rope=rope, attn_mask=attn_mask)
        hidden_states = attn_outputs + residual

        # 第二个残差块：MLP
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states) + residual
        return hidden_states


class LlamaModel(nn.Module):
    """没有预训练头的完整 Llama 模型。"""

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.rotary_emb = RotaryPositionEncoding(
            config.hidden_size // config.num_attention_heads,
            config.max_position_embeddings,
        )

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = nn.RMSNorm(config.hidden_size, eps=1e-5)

    def forward(self, input_ids: Tensor, attn_mask: Tensor) -> Tensor:
        # 将输入的标记 ID 转换为嵌入
        hidden_states = self.embed_tokens(input_ids)
        # 通过所有 Transformer 层，然后通过最后的归一化层
        for layer in self.layers:
            hidden_states = layer(hidden_states, rope=self.rotary_emb, attn_mask=attn_mask)
        hidden_states = self.norm(hidden_states)
        # 返回最终的隐藏状态
        return hidden_states


class LlamaForPretraining(nn.Module):
    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.base_model = LlamaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def forward(self, input_ids: Tensor, attn_mask: Tensor) -> Tensor:
        hidden_states = self.base_model(input_ids, attn_mask)
        return self.lm_head(hidden_states)


def create_causal_mask(seq_len: int, device: torch.device, dtype: torch.dtype = torch.float32) -> Tensor:
    """为自注意力创建一个因果掩码。

    Args:
        seq_len: 序列的长度
        device: 创建掩码的设备
        dtype: 掩码的数据类型

    Returns:
        形状为 (seq_len, seq_len) 的因果掩码
    """
    mask = torch.full((seq_len, seq_len), float('-inf'), device=device, dtype=dtype) \
                .triu(diagonal=1)
    return mask


def create_padding_mask(batch, padding_token_id, device: torch.device, dtype: torch.dtype = torch.float32) -> Tensor:
    """为一批序列创建自注意力的填充掩码。

    Args:
        batch: 一批序列，形状为 (batch_size, seq_len)
        padding_token_id: 填充标记的 ID

    Returns:
        形状为 (batch_size, 1, seq_len, seq_len) 的填充掩码
    """
    padded = torch.zeros_like(batch, device=device, dtype=dtype) \
                  .masked_fill(batch == padding_token_id, float('-inf'))
    mask = padded[:,:,None] + padded[:,None,:]
    return mask[:, None, :, :]


# 生成固定长度填充序列的数据集类
class PretrainingDataset(torch.utils.data.Dataset):
    def __init__(self, dataset: datasets.Dataset, tokenizer: tokenizers.Tokenizer,
                 seq_length: int, device: torch.device = None):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.device = device
        self.seq_length = seq_length
        self.bot = tokenizer.token_to_id("[BOT]")
        self.eot = tokenizer.token_to_id("[EOT]")
        self.pad = tokenizer.token_to_id("[PAD]")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        """从数据集中获取一个标记ID序列。添加 [BOT] 和 [EOT] 标记。
        裁剪并填充到序列长度。
        """
        seq = self.dataset[index]["text"]
        tokens: list[int] = [self.bot] + self.tokenizer.encode(seq).ids + [self.eot]
        # 填充到目标序列长度
        toklen = len(tokens)
        if toklen < self.seq_length+1:
            pad_length = self.seq_length+1 - toklen
            tokens += [self.pad] * pad_length
        # 返回序列
        x = torch.tensor(tokens[:self.seq_length], dtype=torch.int64, device=self.device)
        y = torch.tensor(tokens[1:self.seq_length+1], dtype=torch.int64, device=self.device)
        return x, y

# 使用默认配置创建预训练模型
model_config = LlamaConfig()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = LlamaForPretraining(model_config).to(device)

# 打印模型大小
print(f"模型参数量: {sum(p.numel() for p in model.parameters()) / 1024**2:.2f} M")
print(f"模型缓冲区大小: {sum(p.numel() for p in model.buffers()) / 1024**2:.2f} M")

# 训练参数
epochs = 3
learning_rate = 1e-3
batch_size = 8
seq_length = 512
num_warmup_steps = 1000
PAD_TOKEN_ID = tokenizer.token_to_id("[PAD]")

# 数据加载器、优化器、调度器和损失函数
model.train()
dataloader = torch.utils.data.DataLoader(
    PretrainingDataset(dataset, tokenizer, seq_length, device),
    batch_size=batch_size
)
optimizer = torch.optim.AdamW(
    model.parameters(), lr=learning_rate, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.01
)
num_training_steps = len(dataloader) * epochs
warmup_scheduler = lr_scheduler.LinearLR(
    optimizer,
    start_factor=0.1, end_factor=1.0, total_iters=num_warmup_steps
)
cosine_scheduler = lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=num_training_steps - num_warmup_steps,
    eta_min=0
)
scheduler = lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_scheduler],
    milestones=[num_warmup_steps]
)
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN_ID)

# 查找最后一个检查点
if os.path.exists("llama_pretraining_checkpoint.pth"):
    checkpoint = torch.load("llama_pretraining_checkpoint.pth")
    begin_epoch = checkpoint["epoch"]
    begin_batch = checkpoint["batch"]
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    del checkpoint
    print(f"从周期 {begin_epoch} 和批次 {begin_batch} 恢复训练")
else:
    begin_epoch = 0
    begin_batch = 0

# 开始训练
for epoch in range(begin_epoch, epochs):
    dataloader = torch.utils.data.DataLoader(
        PretrainingDataset(
            dataset.skip(begin_batch * batch_size),
            tokenizer,
            seq_length,
            device,
        ),
        batch_size=batch_size
    )
    pbar = tqdm.tqdm(dataloader, desc=f"周期 {epoch+1}/{epochs}")
    for batch_id, batch in enumerate(pbar):
        if (begin_batch + batch_id) % 1000 == 0:
            # 检查点保存模型和优化器状态
            torch.save({
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "epoch": epoch,
                "batch": batch_id + begin_batch,
            }, f"llama_pretraining_checkpoint.pth")
        # 获取批处理数据
        input_ids, target_ids = batch
        # 创建注意力掩码：因果掩码 + 填充掩码
        attn_mask = create_causal_mask(input_ids.shape[1], device) + \
                    create_padding_mask(input_ids, PAD_TOKEN_ID, device)
        # 从模型获取输出
        logits = model(input_ids, attn_mask)
        # 计算损失：logits和目标之间的交叉熵，忽略填充标记
        loss = loss_fn(logits.view(-1, logits.size(-1)), target_ids.view(-1))
        # 反向传播损失并应用梯度裁剪
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        pbar.set_postfix(loss=loss.item())
        pbar.update(1)
    begin_batch = 0   # 为下一个周期重置
    pbar.close()

# 保存模型
torch.save(model.state_dict(), "llama_pretraining_model.pth")
torch.save(model.base_model.state_dict(), "llama_model.pth")

请注意，这是一个简化的训练方案。专业的模型训练过程会在更大的模型上使用更大的数据集。例如，具有7B-70B参数的Llama 2模型在2万亿个标记上进行训练。用于训练的超参数（如学习率）在实际训练之前会进行调整。

此外，先用较短的序列长度训练模型，然后扩展到更长的序列长度会更有效。众所周知，最初在较低质量的数据上训练模型，然后在结束时使用较高质量的数据会使模型更具表现力。上面的代码中都没有实现这些技术。可以参考之前的文章了解改进训练的技巧。

进一步阅读

• Liu et al (2024) Understanding LLMs: A Comprehensive Overview from Training to Inference
• Grattafiori et al (2024) The Llama 3 Herd of Models
• Groeneveld et al (2024) OLMo: Accelerating the Science of Language Models
• Sebastian Raschka, Build a Large Language Model (From Scratch). Manning Publications 2024

总结

在这篇文章中，你学习了如何在单张GPU上预训练Llama模型。具体来说，你学习了如何：

• 训练一个带有特殊标记（用于下一个标记预测）的分词器
• 准备用于预训练的训练数据
• 在单张GPU上运行带有检查点保存功能的预训练FINISHED
  更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
  对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）