目 录CONTENT

文章目录

在本地 GPU 上预训练 Llama 模型

Administrator
2025-12-22 / 0 评论 / 0 点赞 / 0 阅读 / 0 字

📢 转载信息

原文链接:https://machinelearningmastery.com/pretraining-a-llama-model-on-your-local-gpu/

原文作者:Adrian Tam


像 Llama 这样的仅解码器语言模型通常使用自监督学习目标在大规模文本上进行训练。这被称为预训练,以区别于后续针对特定任务的微调步骤。在本文中,您将学习如何在本地 GPU 上预训练一个 Llama 模型。具体来说,您将学习如何:

  • 准备训练数据
  • 运行预训练

让我们开始吧。

在本地 GPU 上预训练 Llama 模型
照片来自 Hongbin。保留部分权利。

概述

本文分为三个部分;它们是:

  • 使用特殊标记训练分词器
  • 准备训练数据
  • 运行预训练

使用特殊标记训练分词器

您将使用的模型架构与上一篇文章中创建的模型相同。这是一个具有 50,000 词汇大小的 12 层 Llama 模型。用于预训练的数据是 HuggingFaceFW/fineweb 数据集。

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

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

# Load FineWeb 10B sample (using only a slice for demo to save memory)
dataset = datasets.load_dataset("HuggingFaceFW/fineweb", "sample-10BT", split="train", streaming=True)

def get_texts(dataset: datasets.Dataset, limit: int = 100_000) -> Iterator[str]:
    """Get texts from the dataset until the limit is reached or the dataset is exhausted."""
    count = 0
    for sample in dataset:
        yield sample["text"]
        count += 1
        if limit and count >= limit:
            break

# Initialize a BPE model
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
trainer = trainers.BpeTrainer(
    vocab_size=50_000,
    min_frequency=2,
    special_tokens=["[PAD]", "[BOT]", "[EOT]", "[UNK]"],
    show_progress=True,
)

# Train and save the tokenizer to disk
texts = get_texts(dataset, limit=100_000)
tokenizer.train_from_iterator(texts, trainer=trainer)
tokenizer.save("bpe_50k.json")

此分词器使用字节级的 BPE(字节对编码)算法。通常情况下它不会产生任何未知标记,但您仍然为未知标记设置了一个特殊标记。此外,您为文本开始 ([BOT])、文本结束 ([EOT]) 和填充 ([PAD]) 设置了特殊标记。这些对于下一标记预测非常有用。

此代码将自动使用您的 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 Dataset 对象,以便您的模型可以消费,如下所示:

class PretrainingBatchDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, tokenizer, seq_length, batch_size, device):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.device = device
        self.seq_length = seq_length
        self.batch_size = batch_size
        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):
        """Get a sequence of token ids from the dataset. [BOT] and [EOT] tokens
        are added. Clipped and padded to the sequence length.
        """
        seq = self.dataset[index]["text"]
        tokens: list[int] = [self.bot] + self.tokenizer.encode(seq).ids + [self.eot]
        # pad to target sequence length
        toklen = len(tokens)
        if toklen < self.seq_length+1:
            pad_length = self.seq_length+1 - toklen
            tokens += [self.pad] * pad_length

        # return the sequence
        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 数据加载器一起使用它们。由于 PyTorch 的 CrossEntropyLoss 在计算训练损失时需要此类型来识别填充标记,因此您还必须将数据类型设置为 int64

您可以通过创建一个 DataLoader 对象并从中抽取一个批次来测试数据集:

batch_size = 8
seq_length = 512
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
dataloader = torch.utils.data.DataLoader(
    PretrainingBatchDataset(dataset, tokenizer, seq_length, batch_size, device),
    batch_size=batch_size
)

for x, y in dataloader:
    print(x)
    print(y)
    break

运行预训练

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

使用上一篇文章中的模型代码,我们首先创建一个模型对象:

# Create pretraining model with default config
test_config = LlamaConfig()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = LlamaForPretraining(test_config).to(device)

这是一个用于演示目的的小型模型。它只有 1.71 亿个参数,远小于您可以在互联网上找到的任何大型语言模型。

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

  • 此数据集只有一个训练分割。为简单起见,数据未被洗牌,未创建保留集,并且训练循环不包含任何评估步骤。
  • 下一标记预测是针对整个词汇表的分类问题。自然地,损失指标是交叉熵损失。您应该确保在计算损失时不使用填充标记,因为它们不是有效输入。
  • 将序列长度设置为 512。训练模型所需的资源随序列长度呈 $O(N^2)$ 扩展。因此,您更倾向于保持其较短,但过短的序列长度会阻止模型理解更长的上下文。
  • 遵循训练大型语言模型的最佳实践,使用带有预热期的余弦学习率调度器。预热期可以设置为固定的步数或总训练步数的百分比(例如 0.1% 到 2%)。我们在这里将其设置为 1,000 步。
  • 确定序列长度后,调整批次大小以适应您的 GPU 内存。您可以从 8 开始,经验表明它可以适应 12GB 的 VRAM。
  • 对于 HuggingFaceFW/fineweb 10B 数据集中的 1400 万个样本和 100 亿个标记,您可能不需要训练很多个 epoch。事实上,许多大型语言模型只在非常大的数据集上训练 1-3 个 epoch。

让我们将这些参数组合起来定义训练配置:

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

# DataLoader, optimizer, scheduler, and loss function
model.train()
dataloader = torch.utils.data.DataLoader(
    PretrainingBatchDataset(dataset, tokenizer, seq_length, batch_size, device),
    batch_size=batch_size
# ... 内容被截断



🚀 想要体验更好更全面的AI调用?

欢迎使用青云聚合API,约为官网价格的十分之一,支持300+全球最新模型,以及全球各种生图生视频模型,无需翻墙高速稳定,文档丰富,小白也可以简单操作。

0

评论区