手把手从零打造 Llama3：解锁下一代预训练模型

引言

        Llama3 相较于 Llama2，不仅在模型架构上做了显著优化，尤其是全局查询注意力机制（GQA）的引入，使得模型在大规模数据处理上表现更加出色。同时，Llama3 采用了与 GPT 一致的 tiktoken 分词器，大幅提升了分词效率。

        本篇文章将带你从头构建 Llama3 预训练流程，深入了解其关键细节和实现方式，让你掌握这一下一代模型的核心技术。

1. 启动训练脚本

        在这一步中，我们将实现 Llama3 的预训练框架。核心代码包含模型初始化、数据加载与批处理，以及训练循环的定义。为了便于操作，所有模型的细节和实现细节都在 llama_model.py 文件中定义。

import math
import os
import time
from contextlib import nullcontext
from datetime import datetime
from functools import partial

import torch
from LLama_content.llama_model import Transformer, ModelArgs
from LLama_content.llama_model import Task

# 定义输出目录和训练配置
out_dir = "output"
eval_interval = 2000
log_interval = 1
eval_iters = 100
eval_only = False
always_save_checkpoint = False
init_from = "scratch"

# 数据配置
batch_size = 128
max_seq_len = 256
vocab_size = 4096

# 模型配置
dim = 288
n_layers = 8
n_heads = 8
n_group = 4
multiple_of = 32
dropout = 0.0

# AdamW 优化器配置
gradient_accumulation_steps = 4
learning_rate = 5e-4
max_iters = 100000
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

# 学习率衰减
decay_lr = True
warmup_iters = 1000

# 系统设置
device = "cuda:0"
dtype = "bfloat16"

# 保存配置参数
config = {k: globals()[k] for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str))}

# 设置随机种子
torch.manual_seed(1337)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device_type = "cuda" if "cuda" in device else "cpu"
ptdtype = torch.float16

# 混合精度训练设置
ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

# 定义数据迭代器
iter_batches = partial(
    Task.iter_batches,
    batch_size=batch_size,
    max_seq_len=max_seq_len,
    vocab_size=vocab_size,
    vocab_source='custom',
    device=device,
)

# 模型初始化
model_args = dict(dim=dim, n_layers=n_layers, n_heads=n_heads, n_group=n_group, vocab_size=vocab_size, multiple_of=multiple_of, max_seq_len=max_seq_len, dropout=dropout)
gptconf = ModelArgs(**model_args)
model = Transformer(gptconf)
model.to(device)

# AdamW 优化器
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)

        通过上面的代码，我们已经完成了模型的初始化和优化器的配置。Llama3 的核心是其 Transformer 架构，模型参数通过 ModelArgs 配置，其中包含层数、注意力头数和隐藏层维度等关键参数。

2. 数据迭代与训练循环

        在这里，我们将定义数据批处理和训练循环，确保模型能够高效地迭代并更新参数。

# 定义评估函数
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ["train", "val"]:
        batch_iter = iter_batches(split=split)
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = next(batch_iter)
            with ctx:
                logits = model(X, Y)
                loss = raw_model.last_loss
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# 定义学习率调度
def get_lr(it):
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    if it > max_iters:
        return 0.0
    return learning_rate

# 训练循环
train_batch_iter = iter_batches(split="train")
X, Y = next(train_batch_iter)
t0 = time.time()
local_iter_num = 0
running_mfu = -1.0
while iter_num < max_iters:
    lr = get_lr(iter_num)
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    if iter_num % eval_interval == 0:
        losses = estimate_loss()
        print(f"Step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
        if losses["val"] < best_val_loss:
            best_val_loss = losses["val"]
            torch.save(model.state_dict(), os.path.join(out_dir, "checkpoint.pt"))

    for _ in range(gradient_accumulation_steps):
        with ctx:
            logits = model(X, Y)
            loss = raw_model.last_loss / gradient_accumulation_steps

        X, Y = next(train_batch_iter)
        loss.backward()

    optimizer.step()
    optimizer.zero_grad()

    iter_num += 1
    local_iter_num += 1

    if iter_num % log_interval == 0:
        print(f"Step {iter_num}: loss {loss.item():.4f}, lr {lr:e}")

        在这里，训练循环每隔一定步数会进行一次评估，评估过程中模型会通过验证集计算损失值。在每次迭代时，采用了梯度累积来模拟更大的批次训练，确保模型能够高效利用 GPU 资源。

3. 模型推理与结果生成

        完成预训练后，我们可以通过模型生成推理结果。以下是简化后的推理代码，展示了如何使用 Llama3 生成文本。

@torch.no_grad()
def generate_text(model, prompt, max_length):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    output = model.generate(input_ids, max_new_tokens=max_length)
    return tokenizer.decode(output[0], skip_special_tokens=True)

prompt = "从零开始预训练 Llama3"
generated_text = generate_text(model, prompt, max_length=100)
print(generated_text)

        通过调用 generate_text 函数，我们可以在给定的提示下生成文本。这个函数接收输入提示，并使用模型生成基于上下文的自然语言输出。

结语

        通过本文的介绍，相信你对 Llama3 模型的预训练过程有了更深入的理解。从模型初始化、数据迭代到训练循环，再到推理生成，整个流程覆盖了 Llama3 的核心训练架构。

        Llama3 在模型层面的改进，以及使用更高效的分词机制，令其在性能和速度上都得到了显著提升。对于需要大规模文本预训练的场景，Llama3 提供了强大的工具和灵活的接口。如果你有更大的数据集和算力支持，推荐进一步优化预训练过程，以获得更高质量的语言模型。

        未来我们将继续探索如何利用 Llama3 进行更复杂的任务，如微调和下游任务的应用。

如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！

欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持！