当前位置：首页 > news >正文

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

news 2024/12/23 6:07:41

引言

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 partialimport 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_losslosses[k] = loss.item()out[split] = losses.mean()model.train()return out# 定义学习率调度
def get_lr(it):if it < warmup_iters:return learning_rate * it / warmup_itersif it > max_iters:return 0.0return 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"] = lrif 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_stepsX, Y = next(train_batch_iter)loss.backward()optimizer.step()optimizer.zero_grad()iter_num += 1local_iter_num += 1if 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 函数，我们可以在给定的提示下生成文本。这个函数接收输入提示，并使用模型生成基于上下文的自然语言输出。