原文towardsdatascience.com/training-llm-from-scratch-in-rust-03381bbd7204在这篇配套文章中我将展示我在 Rust 中从头开始训练一个类似 GPT 模型的实现。没有 GPU只有 CPU性能比原生 C 代码高出 30 倍。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/41aa162bccc173da174237588256a530.png图片由 GoogleDeepMind 在 Unsplash 提供在我上一篇文章medium.com/towards-data-science/writing-llms-in-rust-looking-for-an-efficient-matrix-multiplication-e9539b0cb9d3中我介绍了矩阵乘法的问题注意力算法如何使用矩阵乘法执行平均过程以及如何高效地实现——至少对我来说——一个矩阵乘法函数在 Rust 中使用 Blas。在这篇新文章中我想展示我在 Rust 中实现 llm.c 的第一个构建块即使用 Rust 从头开始训练一个类似 GPT 的模型。这是我学习更多关于 Rust 生态系统并了解它与 C 相比如何的方法。特别是我希望我的代码能够从 GPT 权重开始仅使用 CPU 训练类似 GPT 的模型—— 因此没有 GPU 或 TPU。我的目标是了解我们可以在简单的笔记本电脑上推动这些模型多远以及 Rust 生态系统可以用于此的程度。最终此代码也可能有助于使用给定的输入语料库微调 GPT 模型。所有相关的代码片段都可以在 这里 找到。我希望每个人都能记住以下要点如何在 Rust 中处理内存参数和 GPT 参数注意力前向是如何实现的如何利用 Rayon 的并行性以及线程锁1. 处理所有参数和内存布局类似 GPT 的模型具有大量的参数和张量嵌入矩阵、层归一化参数、查询、键、值矩阵、注意力输出、前馈层输出等等。如果我们处理的是 PyTorch所有这些都会自动包含在代码范式之中因此无需担心独立的张量对象或者这些张量如何适应内存。相反在我们的 Rust 实现中我们确实需要担心并管理所有这些参数。所有参数都存储为向量Vecf32单个向量在内存中连续因此我们可以有一种简单的内存方法——我们可以提高缓存局部性加快矩阵乘法操作、加载和参数保存过程一个更简单的接口可以与 Blas 的sgemm一起使用以实现更好的矩阵乘法从磁盘进行简单的读写操作另一方面我们必须为此付出代价并与数组切片交互我们必须知道每个数组的大小我们需要谨慎处理切片操作即正确选择张量的正确部分以避免索引边界错误对于每个参数我们必须有一个offset变量。例如wte占用vocab_size*channels个浮点数所以params_memory[0..vocab_size*channels]。然后下一个张量wpe将占用params_memory[vocab_size*channels..(vocab_size*channels max_seq_len*channels)]let wteamp;self.params_memory[0..vocab_size*channels];let wpeamp;self.params_memory[vocab_size*channels..(vocab_size*channelsmax_seq_len*channels)];总体而言唯一的风险是正确切片params_memory数组。如果我们知道大小我们就不会进行无效的内存访问我们也在params_memory变量中有一个单一的真实来源。核心功能用于参数和内存分配是gpt_build_from_checkpoint。从这一点我们正在读取输入文件file.read_f32::LittleEndian()。LittleEndian用于读取和写入数字无论是小端还是大端字节直接到/从字节数组。然后创建参数如下model.params.wtemodel.params_memory[offset..offsetmodel.param_size[0]].to_vec();offsetmodel.param_sizes[0];2a. 编码器前向编码器后向在构建类似 GPT 的模型时需要考虑的第二步是创建单词和位置嵌入即高维向量嵌入。这是通过encoder_forward实现的它返回一个大小为[B, T, C]的激活张量。需要记住的一个重要事项是 B、T 和 C 维度的含义。输入数据被细分为块或大小为 B 的批次。每个批次有一个块大小T。假设取一个文本语料库每个句子可能是一个批次每个句子可能被分为大小为 T 的块。通道大小C是最“奇怪”的一个。我们的思维立刻转向图像处理领域在那里我们有 R、G 和 B 通道。在我们的句子语料库中通道参数是我们创建嵌入的维度。在我们的情况下C 768。此参数直接在检查点文件中读取let(max_t,v,l,nh,c)(model_header[2]asusize,model_header[3]asusize,model_header[4]asusize,model_header[5]asusize,model_header[6]asusize,);特别是对于encoder_forward我们正在处理一个包含我们标记 ID 集合的输入向量inp形状为[B x T]。单词嵌入矩阵wte的大小为[V x C]其中V是词汇量大小即模型可以表示的唯一标记数量在我们的情况下是 50,000。另一方面位置嵌入wpe的大小为[max_t x C]其中max_t是最大序列长度即 1024。为了适应这些值我们正在使用 Rust 切片方法例如在encoder_forward中let out_start_idxb_idx*t*ct_idx*c;//slicing let out_btamp;mut out[out_start_idx..out_start_idxc];//slicing let ixinp[b_idx*tt_idx]asusize;//take theinputvalues let wte_start_idxix*c;//slicing let wte_ixamp;wte[wte_start_idx..wte_start_idxc];//take wte values let wpe_start_idxt_idx*c;//slicing let wpe_tamp;wpe[wpe_start_idx..wpe_start_idxc];//take wpe valuesforiin0..c{out_bt[i]wte_ix[i]wpe_t[i];}要欣赏 Rust 切片考虑B2, T3, C4。这意味着输出长度为B x T x C 24因此out:[out[b0,t0,:],out[b0,t1,:],out[b0,t2,:],out[b1,t0,:],out[b1,t1,:],out[b1,t2,:]]其中out[b, t, :]是 4 个元素。因此对于b1, t2即第二个批次和第三个标记切片从out_start_idx b_idx x t x c t_ids x c 1 x 3 x 4 2 x 4 20开始。2b. 层归一化和注意力GPT 类型的模型在其架构中有一个归一化步骤这样我们就可以稳定训练并提高性能。LayerNorm 确保每个单独的层都有一个归一化的高斯分布。我们在通道维度C上对每个向量进行归一化确保均值为零方差为 1。在实现layernorm_forward时我们使用变量eps 1e-5f32;以防止在计算1 / sqrt( var eps )时除以零。归一化完成后我们可以开始讨论注意力层。在我们的代码中注意力是一个多头自注意力输入维度C被分割到每个头中使用C/nh我们为每个标记计算查询Q、键K和值V我们计算注意力分数最后我们根据注意力分数对值进行加权求和值得记住注意力做了什么。考虑到我们有B, T, and C输入元素我们想要做的是从输入字符串中取出最多T个标记并让算法理解它们是如何“相互关联”的。例如第五个标记应该只考虑它之前的标记即第一个、第二个、第三个和第四个标记。这样流始终从当前标记到上一个时间戳标记。要了解所有标记之间的相互关联程度我们只需要计算 t-th 标记有多少次可能与前一个 (t – 1)-th 标记连接。为了进行高效的平均我们使用矩阵乘法技巧。特别是我们使用三个向量来帮助我们进行平均。前两个向量Q和K是查询和键向量。查询向量回答的问题是“我在寻找什么”而键向量回答的问题是“我包含什么”。现在在K和Q之间进行点积将返回这两个向量对齐的程度即标记内容我包含什么和标记关联我在寻找什么之间的对齐。为了使代码更高效我们正在使用自注意力头。这意味着每个K和Q向量都将具有大小B, T, head_size。计算返回一个大小为B, T, T的权重向量。这意味着权重将有B行正如我们的批次一样多。对于每个批次我们将有一个T x T的方阵这是我们的标记的大小。因此对于每个 t-th 行和 t-th 列的组合我们将有一个“统计”权重表示这两个标记一起出现的可能性。最后一步是对权重向量W使用值向量V进行查询。值向量只是一个简单的线性神经网络层它应用于输入标记。这直接跟在 softmax 处理之后。输出将具有大小B, T, head_size。我们在这里的挑战将是将每个头大小维度的所有通道维度连接起来。让我们转到实际方面。我们注意力前向函数的输入项是out输出缓冲区这将是一个大小为[B, T, C]的张量preatt存储预 softmax 注意力分数的张量att存储最终 softmax 后概率的张量inp输入特征我们将从中派生出查询、键和值向量b, t, c, nh分别是批大小、序列长度、总通道数词汇大小和注意力过程中的头数首先我们准备所有常数。选择c3 c x 3是为了最终连接Q, K, V向量。主循环处理所有头对于所有头我们遍历所有标记然后遍历所有批次。偏移量再次计算如下let query_startb_idx*t*c3t_idx*c3h*hs;let preatt_startb_idx*nh*t*th*t*tt_idx*t;let att_startb_idx*nh*t*th*t*tt_idx*t;这样就可以提取查询向量let query_vecamp;inp[query_start..query_starths];特别是我们有当前头和标记的hs维度查询向量。记住查询表示“这个标记在之前的标记中寻找什么”。然后我们构建键矩阵let mut keys_matVec::with_capacity((t_idx1)*hs);fort2in0..t_idx{let key_startb_idx*t*c3t2*c3h*hsc;//c to skip Qandaccess K keys_mat.extend_from_slice(amp;inp[key_start..key_starths]);}这里with_capacity构造一个新的、空的VecT其容量至少为指定的容量。该向量将能够容纳至少capacity个元素而不进行重新分配。如果capacity为 0则向量不会分配新元素。我们将所有键收集到当前时间步t_idx 1– 记住键是hs维度的正如我们可以从key_start看到的那样。接下来我们使用 Blas 进行预注意力分数的计算。let mut preatt_rowvec![0.0f32;t_idx1];unsafe{sgemm(Layout::RowMajor,Transpose::None,Transpose::None,(t_idx1)asi32,1,hsasi32,1.0,amp;keys_mat,hsasi32,query_vec,1,0.0,amp;mut preatt_row,1,);}这里的矩阵乘法是(t_idx 1) x hs * hs x 1 (t_idx 1) x 1这给出了当前标记的查询与每个先前标记的键的匹配度得分更详细的解释见这里。这些是 logits并且它们通过 softmax 进行归一化并存储在数组att中。最后我们有值向量与注意力分数之间的矩阵乘法。这给我们一个加权求和对于每个标记我们知道它与所有之前看到的标记的分数。3. 尝试利用 Rayon 并行化在真正使用代码之前我想花点时间谈谈 Rayon以及我们如何利用它进行数据并行化。Rayon 是一个数据并行库它允许我们轻松地在多个线程上运行。正如我在关于 matmul 的帖子中之前看到的我们可以使用并行迭代器par_iter()、par_chunks()和par_chunks_mut()。这些迭代器可以直接在所有需要的线程上分区数据负载而不需要你做原始和脏活。这给了我们在使用和安全性方面的一些简单性。你可能在代码中看到如下行out.par_chunks_mut(oc).for_each(|row|{for(o,val)inrow.iter_mut().enumerate(){*valbias[o];}});forrowinout.chunks_mut(oc){//...}Rayon 将out数组分割成大小为oc的块并在并行线程中处理它们。每个线程都得到一个单独的块来处理这样就不会有重叠或对相同数据的竞争。这可以添加到layernorm函数以及encoder函数中因为我们能够处理更大的数据集并确保更好的并行化。然而并非所有闪亮的东西都是金子。一些操作如将梯度累积到单个数组中或在多个线程中汇总统计数据需要共享状态。共享状态意味着我们不能分割数据但我们需要所有数据同时存在因此需要同步。实现共享状态很复杂因为我们需要防止线程在不协调的情况下同时写入相同的内存地址。因此我们需要MutexT。Mutex提供了互斥排他这样一次只有一个线程可以锁定互斥锁确保它是唯一修改包含数据的线程。use std::sync::Mutex;let shared_dataMutex::new(vec![0.0f32;size]);(batches_in_parallel).for_each(|batch|{let mut guardshared_data.lock().unwrap();for(g,val)inguard.iter_mut().zip(batch){*gval;}});如果你看到我的attention_backward函数你会看到它被分成多个子块。这主要是为了避免错误Cannot borrow as mutable more than once at a time。此外在这里我强烈使用 Rayon 和 Mutex以允许在过程中进行一些并行处理。实际上在反向传播过程中我们需要计算相对于输入dinp的梯度相对于预 softmax 注意力分数dpreatt以及相对于注意力概率datt的梯度。这是在整个批次上进行的因此不出所料我们确实需要并行化以避免瓶颈你将在下面看到这是最耗时的步骤。我们想要的并行过程是针对批次和每个注意力头这样我们就可以独立处理这些。然而因为我总是避免错误Cannot borrow as mutable more than once at a time我们需要每个线程计算其局部结果并将所有这些合并成最终的全球梯度。为此我需要使用Mutexlet global_dinpMutex::new(vec![0.0f32;dinp.len()]);let global_dattMutex::new(vec![0.0f32;datt.len()]);let global_dpreattMutex::new(vec![0.0f32;dpreatt.len()]);这样我就可以在循环中使用 Rayon 进行并行工作(0..b).into_par_iter().for_each(|b_idx|{let mut local_dinpvec![0.0f32;dinp.len()];let mut local_dattvec![0.0f32;datt.len()];let mut local_dpreattvec![0.0f32;dpreatt.len()];});让每个线程在不同的b_idx上工作这样我们就可以局部计算梯度。所有这些都是在隔离状态下完成的因此每个线程都在其局部数组上工作。在进行局部计算后我们需要将所有数组组合成全局梯度数组{let mut g_dinpglobal_dinp.lock().unwrap();g_dinp.iter_mut().zip(local_dinp.iter()).for_each(|(g,l)|*gl);}这种lock机制防止了线程以不一致的方式交错写入。当一个线程持有互斥锁时它具有独占访问权。在最后一步我们将切片复制到那里以将局部结果传播到最终数组中dinp.copy_from_slice(amp;global_dinp.lock().unwrap());datt.copy_from_slice(amp;global_datt.lock().unwrap());dpreatt.copy_from_slice(amp;global_dpreatt.lock().unwrap());是时候享受乐趣了代码、性能和推理现在是时候与代码玩耍了。所有这些计算都是在 MacBook ProM216 GB 内存上运行的。首先确保使用python prepro_tinyshakespeare.py下载所需的数据。这将下载输入语料库到data文件夹中。文本被转换为输入训练和验证标记分别为tiny_shakespeare_train.bin和tiny_shakespeare_val.bin。文本使用 GPT-2 标记器进行标记化。然后你可以使用以下命令构建 rust 代码cd llm bash build.sh经过 2000 步之后你可能得到一个类似的推理输出3792,Is340,it922,good11,,611,if345,you423,have26246,pity11,,284,to423,have281,an45618,ornament11,,257,a1486,design11,,198,2514,To9280,dance11,,7365,bat258,he11,,18044,breathe290,and4545,teach30,?440,O11,,611,if340,it307,be2081,true11,,198,1026,It318,is2081,true356,we743,may307,be991,still2877,living11,,611,if340,it307,be2081,true25,:198,46,O11,,2652,stay11,,393,or314,I2236,shall307,be2636,dead13,.628,其中我正在打印标记 ID 及其文本值。代码在 16 个线程上运行。要选择线程数你可以修改此代码中的行和此bash 构建脚本中的行。图 1 显示了每一步的前向和反向传递时间。耗时以ms为单位。总体来看我们可以看到前向传递有相当好的优化平均时间为272.01 /- 57.71 ms。必须做一些工作来使反向传递更高效因为它的平均时间为472.63 /- 51.75 ms。这些耗时比 Karpathy 的原始提交——作为我 Rust 代码的主要灵感来源——平均每步需要30 秒要好 30 倍。…/Images/c182273420335c98b3e3cd67a2834d81.png图 1LLM 训练中前向和反向传递的耗时。耗时以毫秒ms为单位代码在针对小型莎士比亚数据集运行了 2000 步。[图片由作者提供]同时我们可以测量和跟踪训练损失如图 2 所示。总体来看有一个趋势从最初的平均 4.5 下降到最后的 3.2 左右。…/Images/1614f4509d6fcb84c1165b05984f5e58.png图 22000 步后的训练损失。[图片由作者提供]进一步的推理示例在 2000 步之后使用 GPT-2 标记化生成Tis come;youll bear it,this fierce protestor..JULIET:We will.O the loaning makers.First watch man:Youandyour lads,your actions must be controlld by Sir John这可能不是 LLM 的最佳结果但这是在 2000 步之后仅仅在 30 分钟训练后对小型莎士比亚数据集进行微调的结果。结论如果你已经到达这里非常感谢阅读我的文章。我希望你已经全面地查看了我的代码并准备好微调 GPT 模型。文章展示了我深入 Rust 语言并学习如何优化训练 GPT-like 模型代码的方法。特别是我们学习了如何在 Rust 中实现一个类似 GPT 的模型主要障碍有哪些例如矩阵乘法、设置参数以及如何处理内存管理、张量切片和并行化。如何利用 Rayon使用数据并行循环使用Mutex同步数据。该项目的目标是运行基于简单 CPU 的笔记本电脑上的 GPT 训练甚至在这样硬件上Rust 代码看起来也比 C 语言版本要好得多。许多更多项目可以从这个初始种子开始从微调到在 Rust 中实现更好的优化。如果你想与我联系你可以发送电子邮件到 [email protected]。