掘金 人工智能 02月17日
轻松上手训练自己的大模型
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

本文详细介绍了如何使用PyTorch创建一个参数量仅为230万的LLM。文章从环境配置、数据预处理(使用TinyShakespeare数据集)、模型构建(包括基础神经网络模型和后续改进)以及训练过程等方面进行了详尽的阐述。重点介绍了LLaMA架构的关键概念,如RMSNorm、SwiGLU激活函数和RoPE位置嵌入,并逐步展示了如何将这些技术应用到小型LLM的构建中。最后,通过训练和推理过程,验证了模型的性能,并为读者提供了一个实践性的LLM构建指南。

🛠️**环境配置与数据准备**: 使用PyTorch等库,并利用TinyShakespeare数据集进行模型训练,该数据集包含约100万个字符,远小于原始LLaMA的训练数据集。

🧠**LLaMA架构核心概念**: 详细介绍了RMSNorm、SwiGLU激活函数和RoPE位置嵌入等关键技术,这些技术显著提升了模型的性能和效率。

📊**基础模型构建与改进**: 构建了一个简单的神经网络模型作为基础,然后逐步应用LLaMA的技术进行改进,通过损失函数评估模型性能,并进行训练和优化。

🚀**模型训练与推理**: 通过指定训练参数(如epochs、batch_size)并使用Adam优化器,对模型进行训练。训练完成后,使用generate函数进行文本生成,验证模型的推理能力。

制作自己的大型语言模型 (LLM) 是一件很酷的事情,许多大公司(如 Google、Twitter 和 Facebook)都在做这件事。他们发布了这些模型的不同版本,如 70 亿、130 亿或 700 亿。甚至较小的社区也在这样做。您可能已经阅读过有关创建自己的 LLM 的博客或观看过视频,但它们通常谈论的都是理论,而不是实际步骤和代码。在这篇博文中,我将尝试制作一个只有 230 万个参数的 LLM,有趣的是,我们不需要花哨的 GPU。我们会保持简单并使用基本数据集,这样您就可以看到创建自己的百万参数 LLM 是多么容易。目录先决条件确保你对面向对象编程 ( OOP ) 和神经网络 ( NN ) 有基本的了解。熟悉PyTorch也会对编码有所帮助。了解Transformer 架构在深入研究使用 LLaMA 方法创建我们自己的 LLM 之前,必须了解 LLaMA 的架构。下面是 vanilla Transformer 和 LLaMA 之间的比较图。如果您不熟悉 vanilla Transformer 架构,可以阅读此博客获取基本指南。让我们更详细地了解一下 LLaMA 的基本概念:使用 RMSNorm 进行预规范化:在 LLaMA 方法中,采用了一种称为 RMSNorm 的技术来规范化每个 Transformer 子层的输入。该方法受到 GPT-3 的启发,旨在优化与层规范化相关的计算成本。RMSNorm 提供与 LayerNorm 类似的性能,但显著减少了运行时间(减少了 7%∼64%)。它通过强调重新缩放不变性和根据均方根 (RMS) 统计量调节总和输入来实现这一点。主要动机是通过删除均值统计量来简化 LayerNorm。感兴趣的读者可以在此处探索 RMSNorm 的详细实现。SwiGLU 激活函数:LLaMA 引入了 SwiGLU 激活函数,灵感来自 PaLM。要理解 SwiGLU,首先必须掌握 Swish 激活函数。SwiGLU 扩展了 Swish,并包含一个自定义层,该层具有密集网络,用于拆分和乘以输入激活。目的是通过引入更复杂的激活函数来增强模型的表达能力。有关 SwiGLU 的更多详细信息,请参阅相关论文。RoPE:旋转嵌入 (RoPE) 是 LLaMA 中使用的一种位置嵌入。它使用旋转矩阵对绝对位置信息进行编码,并在自注意力公式中自然包含显式相对位置依赖性。RoPE 具有诸多优势,例如可扩展到各种序列长度,并且随着相对距离的增加,标记间依赖性逐渐​​减弱。这是通过与旋转矩阵相乘来编码相对位置来实现的,从而导致相对距离衰减——这是自然语言编码的理想特征。对数学细节感兴趣的人可以参考RoPE 论文。除了这些概念之外,LLaMA 论文还介绍了其他重要方法,包括使用具有特定参数的AdamW 优化器、xformers 库中可用的因果多头注意力运算符等高效实现,以及手动实现的 Transformer 层后向函数,以优化后向传递期间的计算。特别感谢Anush Kumar对 LLaMA 每个重要方面的深入讲解。设置环境我们将在整个项目中使用一系列 Python 库,因此让我们导入它们:import torchfrom torch import nnfrom torch.nn import functional as Fimport numpy as npfrom matplotlib import pyplot as pltimport timeimport pandas as pdimport urllib.request此外,我正在创建一个存储模型参数的配置对象。MASTER_CONFIG = { }这种方法保持了灵活性,允许在未来根据需要添加更多参数。数据预处理在原始 LLaMA 论文中,采用了各种开源数据集来训练和评估模型。不幸的是,对于较小的项目来说,使用大量数据集可能不切实际。因此,对于我们的实现,我们将采取一种更为温和的方法,创建一个大幅缩小版的 LLaMA。鉴于无法访问大量数据的限制,我们将专注于使用 TinyShakespeare 数据集训练简化版的 LLaMA。此开源数据集可在此处获取,包含来自各种莎士比亚作品的约 40,000 行文本。这一选择受到Karpathy 的 Makemore 系列的影响,该系列为训练语言模型提供了宝贵的见解。虽然 LLaMA 是在包含1.4 万亿个标记的庞大数据集上进行训练的,但我们的数据集 TinyShakespeare 包含大约100 万个字符。首先,让我们通过下载获取数据集:url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"file_name = "tinyshakespeare.txt"urllib.request.urlretrieve(url, file_name)该 Python 脚本从指定的 URL 获取 tinyshakespeare 数据集,并将其保存在本地,文件名为 “tinyshakespeare.txt”接下来,让我们确定词汇量,它代表数据集中字符的唯一数量。以下是代码片段:lines = open("tinyshakespeare.txt", 'r').read()vocab = sorted(list(set(lines)))print('Printing the first 10 characters of the vocab list:', vocab[:10])print('Total number of characters in our dataset (Vocabulary Size):', len(vocab))现在,我们创建整数到字符 ( itos ) 和字符到整数 ( stoi ) 之间的映射。代码如下:itos = {i: ch for i, ch in enumerate(vocab)}stoi = {ch: i for i, ch in enumerate(vocab)}在原始的 LLaMA 论文中,使用了 Google 的SentencePiece 字节对编码标记器。但是,为了简单起见,我们将选择基本的字符级标记器。让我们创建稍后将应用于数据集的编码和解码函数:def encode(s): return [stoi[ch] for ch in s]def decode(l): return ''.join([itos[i] for i in l])decode(encode("morning"))最后一行将输出morning确认编码和解码功能的正确功能。我们现在将数据集转换为 torch 张量,并指定其数据类型以便使用PyTorch进行进一步操作:dataset = torch.tensor(encode(lines), dtype=torch.int8)print(dataset.shape)输出 istorch.Size([1115394]) 表示我们的数据集包含大约一百万个 token。值得注意的是,这比包含1.4 万亿个 token 的LLaMA 数据集要小得多。我们将创建一个函数,负责将数据集拆分为训练集、验证集或测试集。在机器学习或深度学习项目中,这种拆分对于开发和评估模型至关重要,同样的原则也适用于复制大型语言模型 (LLM) 方法:def get_batches(data, split, batch_size, context_window, config=MASTER_CONFIG): train = data[:int(.8 len(data))] val = data[int(.8 len(data)): int(.9 len(data))] test = data[int(.9 len(data)):] batch_data = train if split == 'val': batch_data = val if split == 'test': batch_data = test ix = torch.randint(0, batch_data.size(0) - context_window - 1, (batch_size,)) x = torch.stack([batch_data[i:i+context_window] for i in ix]).long() y = torch.stack([batch_data[i+1:i+context_window+1] for i in ix]).long() return x, y现在我们的分裂函数已经定义,让我们建立对此过程至关重要的两个参数:MASTER_CONFIG.update({ 'batch_size': 8, 'context_window': 16 })batch_size 决定每次随机分割处理多少个批次,而 context_window 指定每个批次的每个输入(x)和目标(y)序列中的字符数。让我们从数据集中批次 8 和上下文窗口 16 的训练分割中打印一个随机样本:xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])decoded_samples = [(decode(xs[i].tolist()), decode(ys[i].tolist())) for i in range(len(xs))]print(decoded_samples)评估策略现在,我们将创建一个专门用于评估我们自己创建的 LLaMA 架构的函数。在定义实际模型方法之前这样做的原因是为了能够在训练过程中进行持续评估。@torch.no_grad() def evaluate_loss(model, config=MASTERCONFIG): out = {} model.eval() for split in ["train", "val"]: losses = [] for in range(10): xb, yb = get_batches(dataset, split, config['batch_size'], config['contextwindow']) , loss = model(xb, yb) losses.append(loss.item()) out[split] = np.mean(losses) model.train() return out我们使用损失作为衡量模型在训练迭代过程中表现的指标。我们的函数迭代训练和验证分割,计算每个分割的 10 个批次的平均损失,最后返回结果。然后使用 model.train() 将模型重新设置为训练模式。建立基础神经网络模型我们正在构建一个基本的神经网络,稍后我们将使用 LLaMA 技术对其进行改进。class SimpleBrokenModel(nn.Module): def init(self, config=MASTER_CONFIG): super().init() self.config = config self.embedding = nn.Embedding(config['vocab_size'], config['d_model']) self.linear = nn.Sequential( nn.Linear(config['d_model'], config['d_model']), nn.ReLU(), nn.Linear(config['d_model'], config['vocab_size']), ) print("Model parameters:", sum([m.numel() for m in self.parameters()]))在当前架构中,嵌入层的词汇量为 65,代表我们数据集中的字符。由于这是我们的基础模型,因此我们使用 ReLU 作为线性层中的激活函数;但是,稍后将使用 LLaMA 中使用的 SwiGLU 替换它。为了为我们的基础模型创建前向传递,我们必须在 NN 模型中定义一个前向函数。class SimpleBrokenModel(nn.Module): def init(self, config=MASTER_CONFIG): ... def forward(self, idx, targets=None): x = self.embedding(idx) a = self.linear(x) logits = F.softmax(a, dim=-1) if targets is not None: loss = F.cross_entropy(logits.view(-1, self.config['vocab_size']), targets.view(-1)) return logits, loss else: return logits print("Model parameters:", sum([m.numel() for m in self.parameters()]))此正向传递函数将字符索引 (idx) 作为输入,应用嵌入层,将结果传递到线性层,应用 softmax 激活以获得概率分布 (logits)。如果提供了目标,它会计算交叉熵损失并返回 logits 和损失。如果没有提供目标,它只返回 logits。要实例化此模型,我们可以直接调用该类并打印简单神经网络模型中的参数总数。我们将线性层的维度设置为 128,并在配置对象中指定此值:MASTER_CONFIG.update({ 'd_model': 128,})model = SimpleBrokenModel(MASTER_CONFIG)print("Total number of parameters in the Simple Neural Network Model:", sum([m.numel() for m in model.parameters()]))我们的简单神经网络模型包含大约 33,000 个参数。类似地,为了计算对数和损失,我们只需要将分割的数据集输入到我们的模型中:xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = model(xs, ys)为了训练我们的基础模型并记录其性能,我们需要指定一些参数。我们总共训练了 1000 个时期。将批次大小从 8 增加到 32,并将 log_interval 设置为 10,表示代码将每 10 个批次打印或记录有关训练进度的信息。为了进行优化,我们将使用 Adam 优化器。MASTER_CONFIG.update({ 'epochs': 1000, 'log_interval': 10, 'batch_size': 32, })model = SimpleBrokenModel(MASTER_CONFIG)optimizer = torch.optim.Adam( model.parameters(), )让我们执行训练过程并捕获基础模型的损失,包括参数总数。此外,为了清晰起见,每行都进行了注释:def train(model, optimizer, scheduler=None, config=MASTER_CONFIG, print_logs=False): losses = [] start_time = time.time() for epoch in range(config['epochs']): optimizer.zero_grad() xs, ys = get_batches(dataset, 'train', config['batch_size'], config['context_window']) logits, loss = model(xs, targets=ys) loss.backward() optimizer.step() if scheduler: scheduler.step() if epoch % config['log_interval'] == 0: batch_time = time.time() - start_time x = evaluate_loss(model) losses += [x] if print_logs: print(f"Epoch {epoch} | val loss {x['val']:.3f} | Time {batch_time:.3f} | ETA in seconds {batch_time (config['epochs'] - epoch)/config['log_interval'] :.3f}") start_time = time.time() if scheduler: print("lr: ", scheduler.get_lr()) print("Validation loss: ", losses[-1]['val']) return pd.DataFrame(losses).plot()train(model, optimizer)训练前的初始交叉熵损失为 4.17,经过 1000 个 epoch 后,该损失降至 3.93。在这种情况下,交叉熵反映了选择错误单词的可能性。我们的模型在 logits 上加入了一个 softmax 层,它将数字向量转换为概率分布。我们使用内置的 F.cross_entropy 函数,需要直接传入未归一化的 logits。因此,我们将相应地修改我们的模型。class SimpleModel(nn.Module): def init(self, config): ... def forward(self, idx, targets=None): x = self.embedding(idx) logits = self.linear(x) if targets is not None: ...让我们重新创建更新后的 SimpleModel 并对其进行 1000 个 epoch 的训练以观察任何变化:model = SimpleModel(MASTER_CONFIG)xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = model(xs, ys)optimizer = torch.optim.Adam(model.parameters())train(model, optimizer)将损失降低到 2.51 后,让我们探索一下我们的语言模型(约有33,000 个参数)如何在推理过程中生成文本。我们将创建一个“generate”函数,稍后在复制 LLaMA 时会用到它:def generate(model, config=MASTER_CONFIG, max_newtokens=30): idx = torch.zeros(5, 1).long() for in range(max_new_tokens): logits = model(idx[:, -config['context_window']:]) last_time_step_logits = logits[ :, -1, : ] p = F.softmax(last_time_step_logits, dim=-1) idx_next = torch.multinomial( p, num_samples=1 ) idx = torch.cat([idx, idx_next], dim=-1) return [decode(x) for x in idx.tolist()]generate(model)使用我们约 33K 个参数的基本模型,生成的文本看起来并不好。不过,既然我们已经用这个简单的模型奠定了基础,我们将在下一节中继续构建 LLaMA 架构。复制LLaMA架构在博客的前面部分,我们介绍了基本概念,现在,我们将这些概念整合到我们的基础模型中。LLaMA 对原始 Transformer 进行了三项架构修改:RMSNorm 用于预归一化旋转嵌入SwiGLU 激活函数我们会将这些修改逐一纳入我们的基础模型,并在其基础上进行迭代和构建。预归一化的 RMSNorm:我们正在定义具有以下功能的 RMSNorm 函数:class RMSNorm(nn.Module): def init(self, layer_shape, eps=1e-8, bias=False): super(RMSNorm, self).init() self.register_parameter("scale", nn.Parameter(torch.ones(layer_shape))) def forward(self, x): """ Assumes shape is (batch, seq_len, d_model) """ ff_rms = torch.linalg.norm(x, dim=(1,2)) x[0].numel() -.5 raw = x / ff_rms.unsqueeze(-1).unsqueeze(-1) return self.scale[:x.shape[1], :].unsqueeze(0) * raw我们定义 RMSNorm 类。在初始化期间,它会注册一个比例参数。在前向传递中,它会计算输入张量的Frobenius 范数,然后对张量进行归一化。最后,通过注册的比例参数对张量进行缩放。此函数旨在用于 LLaMA 中以替换 LayerNorm 操作。现在是时候将 LLaMA 的第一个实现概念 RMNSNorm 合并到我们的简单 NN 模型中了。以下是更新后的代码:class SimpleModel_RMS(nn.Module): def init(self, config): super().init() self.config = config self.embedding = nn.Embedding(config['vocab_size'], config['d_model']) self.rms = RMSNorm((config['context_window'], config['d_model'])) self.linear = nn.Sequential( ... ) print("Model parameters:", sum([m.numel() for m in self.parameters()])) def forward(self, idx, targets=None): x = self.embedding(idx) x = self.rms(x) logits = self.linear(x) if targets is not None: ...让我们用 RMNSNorm 执行修改后的 NN 模型,并观察模型中更新的参数数量以及损失:model = SimpleModel_RMS(MASTER_CONFIG)xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = model(xs, ys)optimizer = torch.optim.Adam(model.parameters())train(model, optimizer)验证损失略有减少,我们更新后的 LLM 的参数现在总计约为 55,000 个。旋转嵌入:接下来,我们将实现旋转位置嵌入。在 RoPE 中,作者建议通过旋转嵌入来嵌入序列中标记的位置,在每个位置应用不同的旋转。让我们创建一个模拟 RoPE 实际论文实现的函数:def get_rotary_matrix(context_window, embedding_dim): R = torch.zeros((context_window, embedding_dim, embedding_dim), requires_grad=False) for position in range(context_window): for i in range(embedding_dim // 2): theta = 10000. (-2. (i - 1) / embedding_dim) m_theta = position theta R[position, 2 i, 2 i] = np.cos(m_theta) R[position, 2 i, 2 i + 1] = -np.sin(m_theta) R[position, 2 i + 1, 2 i] = np.sin(m_theta) R[position, 2 i + 1, 2 i + 1] = np.cos(m_theta) return R我们根据指定的上下文窗口和嵌入维度生成一个旋转矩阵,遵循提出的 RoPE 实现。您可能熟悉涉及注意力头的 transformers 架构,因此在复制 LLaMA 时,我们同样需要创建注意力头。首先,让我们使用之前为旋转嵌入开发的 get_rotary_matrix 函数创建一个带掩码的注意力头。*此外,为了清晰起见,每行都进行了注释:class RoPEAttentionHead(nn.Module): def init(self, config): super().init() self.config = config self.w_q = nn.Linear(config['d_model'], config['d_model'], bias=False) self.w_k = nn.Linear(config['d_model'], config['d_model'], bias=False) self.w_v = nn.Linear(config['d_model'], config['d_model'], bias=False) self.R = get_rotary_matrix(config['context_window'], config['d_model']) def get_rotary_matrix(context_window, embedding_dim): R = torch.zeros((context_window, embedding_dim, embedding_dim), requires_grad=False) for position in range(context_window): for i in range(embedding_dim//2): ... return R def forward(self, x, return_attn_weights=False): b, m, d = x.shape q = self.w_q(x) k = self.w_k(x) v = self.w_v(x) q_rotated = (torch.bmm(q.transpose(0, 1), self.R[:m])).transpose(0, 1) k_rotated = (torch.bmm(k.transpose(0, 1), self.R[:m])).transpose(0, 1) activations = F.scaled_dot_product_attention( q_rotated, k_rotated, v, dropout_p=0.1, is_causal=True ) if return_attn_weights: attn_mask = torch.tril(torch.ones((m, m)), diagonal=0) attn_weights = torch.bmm(q_rotated, k_rotated.transpose(1, 2)) / np.sqrt(d) + attn_mask attn_weights = F.softmax(attn_weights, dim=-1) return activations, attnweights return activations现在我们有一个返回注意力权重的单个掩蔽注意力头,下一步是创建一个多头注意力机制。class RoPEMaskedMultiheadAttention(nn.Module): def init(self, config): super().init() self.config = config self.heads = nn.ModuleList([ RoPEMaskedAttentionHead(config) for in range(config['n_heads']) ]) self.linear = nn.Linear(config['n_heads'] config['d_model'], config['d_model']) self.dropout = nn.Dropout(.1) def forward(self, x): heads = [h(x) for h in self.heads] x = torch.cat(heads, dim=-1) x = self.linear(x) x = self.dropout(x) return x原始论文在其较小的 7b LLM 变体中使用了 32 个主管,但由于限制,我们将在我们的方法中使用 8 个主管。MASTER_CONFIG.update({ 'n_heads': 8,})现在我们已经实现了旋转嵌入和多头注意力,让我们用更新的代码重写我们的 RMNSorm 神经网络模型。我们将测试其性能,计算损失,并检查参数数量。我们将这个更新的模型称为“RopeModel”*class RopeModel(nn.Module): def init(self, config): super().init() self.config = config self.embedding = nn.Embedding(config['vocab_size'], config['d_model']) self.rms = RMSNorm((config['context_window'], config['d_model'])) self.rope_attention = RoPEMaskedMultiheadAttention(config) self.linear = nn.Sequential( nn.Linear(config['d_model'], config['d_model']), nn.ReLU(), ) self.last_linear = nn.Linear(config['d_model'], config['vocab_size']) print("model params:", sum([m.numel() for m in self.parameters()])) def forward(self, idx, targets=None): x = self.embedding(idx) x = self.rms(x) x = x + self.rope_attention(x) x = self.rms(x) x = x + self.linear(x) logits = self.last_linear(x) if targets is not None: loss = F.cross_entropy(logits.view(-1, self.config['vocab_size']), targets.view(-1)) return logits, loss else: return logits让我们使用 RMNSNorm、旋转嵌入和 Masked Multi Head Attentions 执行修改后的 NN 模型,以观察模型中更新的参数数量以及损失:model = RopeModel(MASTER_CONFIG)xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = model(xs, ys)optimizer = torch.optim.Adam(model.parameters())train(model, optimizer)验证损失再次略有下降,我们更新后的 LLM 的参数现在总计约为 55,000 个。让我们对模型进行更多次训练,看看我们重新创建的 LLaMA LLM 的损失是否继续减少。MASTER_CONFIG.update({ "epochs": 5000, "log_interval": 10,})train(model, optimizer)验证损失持续减少,这表明更多次的训练可以进一步减少损失,尽管效果并不显著。SwiGLU 激活函数:如前所述,LLaMA 的创建者使用 SwiGLU 而不是 ReLU,因此我们将在代码中实现 SwiGLU 方程。class SwiGLU(nn.Module): """ Paper Link -> https://arxiv.org/pdf/2002.05202v1.pdf """ def init(self, size): super().init() self.config = config self.linear_gate = nn.Linear(size, size) self.linear = nn.Linear(size, size) self.beta = torch.randn(1, requires_grad=True) self.beta = nn.Parameter(torch.ones(1)) self.register_parameter("beta", self.beta) def forward(self, x): swish_gate = self.linear_gate(x) torch.sigmoid(self.beta self.linear_gate(x)) out = swish_gate self.linear(x) return out在 python 中实现 SwiGLU 方程后,我们需要将其集成到我们修改后的 LLaMA 语言模型 ( RopeModel ) 中。class RopeModel(nn.Module): def init(self, config): super().init() self.config = config self.embedding = nn.Embedding(config['vocab_size'], config['d_model']) self.rms = RMSNorm((config['context_window'], config['d_model'])) self.rope_attention = RoPEMaskedMultiheadAttention(config) self.linear = nn.Sequential( nn.Linear(config['d_model'], config['d_model']), SwiGLU(config['d_model']), ) self.last_linear = nn.Linear(config['d_model'], config['vocab_size']) print("model params:", sum([m.numel() for m in self.parameters()])) def forward(self, idx, targets=None): x = self.embedding(idx) x = self.rms(x) x = x + self.rope_attention(x) x = self.rms(x) x = x + self.linear(x) logits = self.last_linear(x) if targets is not None: loss = F.cross_entropy(logits.view(-1, self.config['vocab_size']), targets.view(-1)) return logits, loss else: return logits让我们使用 RMNSNorm、旋转嵌入、Masked Multi Head Attentions 和 SwiGLU 执行修改后的 NN 模型,以观察模型中更新的参数数量以及损失:model = RopeModel(MASTER_CONFIG)xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = model(xs, ys)optimizer = torch.optim.Adam(model.parameters())train(model, optimizer)验证损失再次略有下降,我们更新后的 LLM 的参数现在总计约为 60,000 个。到目前为止,我们已经成功实现了论文中的关键组件,即 RMNSorm、RoPE 和 SwiGLU。我们观察到这些实现导致损失略有减少。现在我们将为 LLaMA 添加层来检查其对损失的影响。原始论文对 7b 版本使用了 32 层,但我们只使用 4 层。让我们相应地调整模型设置。MASTER_CONFIG.update({ 'n_layers': 4, })让我们首先创建一个单层来了解其影响。class LlamaBlock(nn.Module): def init(self, config): super().init() self.config = config self.rms = RMSNorm((config['context_window'], config['d_model'])) self.attention = RoPEMaskedMultiheadAttention(config) self.feedforward = nn.Sequential( nn.Linear(config['d_model'], config['d_model']), SwiGLU(config['d_model']), ) def forward(self, x): x = self.rms(x) x = x + self.attention(x) x = self.rms(x) x = x + self.feedforward(x) return x创建 LlamaBlock 类的实例并将其应用于随机张量。block = LlamaBlock(MASTER_CONFIG)random_input = torch.randn(MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'], MASTER_CONFIG['d_model'])output = block(random_input)成功创建单层后,我们现在可以使用它来构建多层。此外,我们将模型类从“ropemodel”重命名为“Llama”,因为我们已经复制了 LLaMA 语言模型的每个组件。class Llama(nn.Module): def init(self, config): super().init() self.config = config self.embeddings = nn.Embedding(config['vocab_size'], config['d_model']) self.llamablocks = nn.Sequential( OrderedDict([(f"llama{i}", LlamaBlock(config)) for i in range(config['n_layers'])]) ) self.ffn = nn.Sequential( nn.Linear(config['d_model'], config['d_model']), SwiGLU(config['d_model']), nn.Linear(config['d_model'], config['vocab_size']), ) print("model params:", sum([m.numel() for m in self.parameters()])) def forward(self, idx, targets=None): x = self.embeddings(idx) x = self.llama_blocks(x) logits = self.ffn(x) if targets is None: return logits else: loss = F.cross_entropy(logits.view(-1, self.config['vocab_size']), targets.view(-1)) return logits, loss让我们使用 RMNSNorm、旋转嵌入、Masked Multi Head Attentions、SwiGLU 和 N_layers 执行修改后的 LLaMA 模型,以观察模型中更新的参数数量以及损失:llama = Llama(MASTER_CONFIG)xs, ys = get_batches(dataset, 'train', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = llama(xs, ys)optimizer = torch.optim.Adam(llama.parameters())train(llama, optimizer)虽然存在过度拟合的可能性,但探索延长训练周期数是否会导致损失进一步减少至关重要。另外,请注意,我们当前的 LLM 有超过 200 万个参数。让我们对它进行更多次的训练。MASTER_CONFIG.update({ 'epochs': 10000,})train(llama, optimizer, scheduler=None, config=MASTER_CONFIG)这里的损失是 1.08,我们可以实现更低的损失,而不会遇到明显的过拟合。这表明该模型表现良好。让我们再次训练模型,这次加入一个调度程序train(llama, optimizer, config=MASTER_CONFIG)到目前为止,我们已经在自定义数据集上成功实现了 LLaMA 架构的精简版本。现在,让我们检查一下 200 万参数语言模型生成的输出。# Generate text using the trained LLM (llama) with a maximum of 500 tokensgenerated_text = generate(llama, MASTER_CONFIG, 500)[0]print(generated_text)尽管一些生成的单词可能不是完美的英语,但我们仅有 200 万个参数的 LLM 已经显示出对英语的基本理解。现在,让我们看看我们的模型在测试集上的表现如何。xs, ys = get_batches(dataset, 'test', MASTER_CONFIG['batch_size'], MASTER_CONFIG['context_window'])logits, loss = llama(xs, ys)print(loss)测试集上的计算损失约为 1.236。检查生成的输出的变化的一个简单方法是进行大量时期的训练并观察结果。尝试超参数超参数调整是训练神经网络的关键步骤。在原始的 Llama 论文中,作者采用了余弦退火学习计划。然而,在我们的实验中,它表现不佳。以下是使用不同学习计划试验超参数的示例:MASTER_CONFIG.update({ "epochs": 1000})llama_with_cosine = Llama(MASTER_CONFIG)llama_optimizer = torch.optim.Adam( llama.parameters(), betas=(.9, .95), weight_decay=.1, eps=1e-9, lr=1e-3)scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(llama_optimizer, 300, eta_min=1e-5)train(llama_with_cosine, llama_optimizer, scheduler=scheduler)保存你的语言模型(法学硕士)您可以使用以下命令保存整个 LLM 或仅保存参数:torch.save(llama, 'llama_model.pth')torch.save(llama.state_dict(), 'llama_model_params.pth')要将 PyTorch 模型保存到 Hugging Face 的 Transformers 库,可以使用 save_pretrained 方法。以下是示例:from transformers import GPT2LMHeadModel, GPT2Configllama_config = GPT2Config.from_dict(MASTER_CONFIG)llama_transformers = GPT2LMHeadModel(config=llama_config)llama_transformers.load_state_dict(llama.state_dict())output_dir = "llama_model_transformers"llama_transformers.save_pretrained(output_dir)GPT2Config 用于创建与 GPT-2 兼容的配置对象。然后,创建 GPT2LMHeadModel 并加载 Llama 模型中的权重。最后,调用 save_pretrained 将模型和配置保存在指定的目录中。然后您可以使用 Transformers 库加载模型:from transformers import GPT2LMHeadModel, GPT2Configoutput_dir = "llama_model_transformers"llama_transformers = GPT2LMHeadModel.from_pretrained(output_dir)结论在这篇博文中,我们逐步介绍了如何实施 LLaMA 方法来构建您自己的小型语言模型 (LLM)。建议将您的模型扩展到大约 1500 万个参数,因为 1000 万到 2000 万之间的小型模型往往能更好地理解英语。一旦您的 LLM 精通语言,您就可以针对特定用例对其进行微调。

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

LLM PyTorch TinyShakespeare RMSNorm SwiGLU
相关文章