스크립트
00:00:00.
00:00:13大家好,感谢大家的到来,
00:00:18我现在开始演讲。那么
00:00:22这次演讲是关于一篇论文,
00:00:27我写的这篇论文非常简单。
00:00:33它的核心观点非常明确。
00:00:36基本上就是两行
00:00:39代数公式,就能让 Transformer 中的 RMS Norm
00:00:42层成本更低、
00:00:45速度更快,并且对它作为
00:00:48Transformer 架构中的一层进行了优化。
00:00:51这类似于
00:00:54Layer Norm 曾经是标准配置,
00:00:57后来被 RMS Norm 所取代。
00:01:00这项研究也延续了这种思考方式。
00:01:04我有幸认识了
00:01:09开源领域的一些朋友,
00:01:11并共同撰写了这篇论文,
00:01:14合作者是 Nils Graf,他算是
00:01:17这个想法的提出者。
00:01:21工作就是从那里展开的。
00:01:23论文已经预印在 arXiv 上。
00:01:26大家可以去看看、阅读并测试一下。
00:01:29还有一个开源代码库。
00:01:31至于这个概念,或者说这个思路
00:01:36和思考方式,最简单
00:01:38理解它的方式可能就是类比 FlashAttention。
00:01:40就像 FlashAttention
00:01:45通过延迟乘法计算
00:01:48来尽量减少内存之间的通信,
00:01:52从而使整个计算过程
00:01:54更加快速一样。
00:01:55我们的思路也与此非常相似。
00:01:58它进行了一些针对性的优化,
00:02:01大大加快了 RMS Norm 的处理速度,
00:02:08进而提升了整个 Transformer 的性能。
00:02:15可能有人会问:为什么要优化 RMS Norm?
00:02:19毕竟这一层几乎不承担什么数学计算量。
00:02:24确实如此。
00:02:26如果你看纯数学计算量的占比,
00:02:29它确实非常小。
00:02:31然而,它所占用的时钟时间(或称挂钟时间),
00:02:34其实相当长。
00:02:36例如,在单步解码过程中,
00:02:40也就是进行推理的时候,
00:02:43RMS Norm 可能会被启动多达 33 次。
00:02:47当然,这也取决于具体的模型等等。
00:02:50论文中列出了具体的测试模型
00:02:52以及具体的测试方法。
00:02:54问题在于如何对其进行优化,
00:02:59以及如何避免矩阵乘法中
00:03:03这种权重的额外开销。
00:03:06之所以缓慢,
00:03:09并不是因为 GPU 做数学计算慢或差,
00:03:15而是因为它们在实际计算之外的其他操作上效率低下。
00:03:20比如启动任务,即启动真正的计算。
00:03:24例如在某些实验中,
00:03:30频繁启动进程多达 33 次,
00:03:32这会耗费大量的开销。
00:03:35比如,将每次归一化操作
00:03:40融合(Fuse)到矩阵乘法中,就能有助于避免这种开销。
00:03:44此外,进行权重折叠(Weight Folding)
00:03:46有助于减少内存之间的数据传输,
00:03:50而数据传输也是 GPU 相对较慢的过程。
00:03:54还有等待时间。
00:03:56例如,延迟 RMS Norm 层中
00:04:00执行的标量除法,
00:04:02也是避免这种等待步骤的一种方法。
00:04:06所以,简而言之,这篇论文所做的
00:04:09就是通过改善这三个方面,
00:04:13在计算 RMS Norm 的方式上
00:04:16运用了一些代数技巧。
00:04:19仅此而已。
00:04:21在数学层面上,技巧主要如下。
00:04:25主要围绕前两个命题。
00:04:29一个是无权重归一化。
00:04:31大家可以在这里看到。
00:04:33另一个是延迟归一化。
00:04:35这是第二个命题。
00:04:37而在较新的模型架构中,
00:04:39有时 RMS Norm 会连续出现两次。
00:04:44例如在 Gemma 4 中就是这样。
00:04:48因此抵消预归一化(Pre-normalization)也是可行的。
00:04:52所有这些都在论文中给出了完整的代数证明。
00:04:58第一个命题是这样的,
00:05:00它将缩放增益和权重
00:05:04折叠进同一个矩阵中。
00:05:06大家可以看到带有星号的 W。
00:05:09那是离线计算好的,
00:05:12类似于 FlashAttention 中
00:05:14在侧边预先计算一些内容,
00:05:16这样就不需要时刻在内存之间
00:05:18进行频繁的数据通信。
00:05:20这是完成的其中一步,
00:05:24即这种权重折叠。
00:05:26另一步是延迟矩阵乘法中的
00:05:32标量除法,以便它们能够并行执行。
00:05:35在常规情况下,你必须先计算一次,
00:05:38然后等待,再进行下一次计算。
00:05:41而在这里,思路是将它进行拆分,
00:05:44使其能够实现并行化处理。
00:05:48第三个命题算是它的一个变体,
00:05:51如果存在两个归一化层,
00:05:56由于它具有缩放不变性,
00:05:58就可以省去其中一个,且依然能正常工作。
00:06:01这适用于支持此类架构和实现的
00:06:04最新模型。
00:06:10因此为了在实际中实现这一点,
00:06:13特别是第二个命题。
00:06:16比如对于第一个命题,实现起来很简单。
00:06:20有一个叫 transformer tricks 的代码库,
00:06:22你可以直接应用到任何模型上,都能起作用。
00:06:25但要实现第二个,就需要编写一些 CUDA 算子(Kernel)。
00:06:29所以做起来没那么简单直接。
00:06:31为了做到这一点,
00:06:35我在实现它时做过一次实验。
00:06:42整体看起来还行,就像是:
00:06:46“提示词是:Transformer 架构”
00:06:48“彻底改变了自然语言处理,因为……”
00:06:51然后应该会得到某种符合预期的输出。
00:06:54但在我得到的输出中,
00:06:56我看到了这种重复和一步的滞后。
00:06:59正如大家在这里看到的,单词“because”再次出现了。
00:07:01GPU 流(Streams)当时出了一些状况,
00:07:07我试图找出到底出了什么问题。
00:07:10我遇到了这种一步的滞后,以及某种程度上
00:07:15来自于上一步的输出。
00:07:18在对这一切进行调试的过程中,
00:07:21我发现在构建类似系统的过程中,
00:07:26就像我解释的命题二,即把这两个操作延迟处理,
00:07:32在 CUDA 中你可以做两件事。
00:07:35你可以利用 Tensor Core 来完成矩阵乘法的一部分,
00:07:39也可以利用 CUDA Core 来运行诸如逐元素操作、
00:07:44归约(Reduction)、平方根等等。
00:07:47所以思路是让它们并行运行,从而获得
00:07:52我在论文中解释的优势,来实际验证这个概念。
00:07:58理想中的运行流程原本是这样的。
00:08:00如果采用顺序执行的方式,
00:08:03当向量单元在计算 RMS 和缩放时,就会存在这段空闲等待时间,
00:08:10紧接着才是矩阵乘法。
00:08:12所以论文中提出的 flash norm 技术,其基本构想
00:08:16就是让这两项计算并行进行。
00:08:19由矩阵单元计算矩阵乘法,同时由向量单元计算 RMS。
00:08:23通过这种方式来节省运行时间。
00:08:26但是,单纯用 Python 是无法实现这种操作的。
00:08:28你必须深入到更底层的层面。
00:08:30于是我编写了像这样的 CUDA 代码来实现。
00:08:35当时来看,整体逻辑似乎没有什么问题。
00:08:42然而,我很快意识到自己漏掉了一个细节。
00:08:47问题在于,最后需要将两个流合并(join)的地方,被我写成了隐式合并。
00:08:57当我对这段代码进行测试时,单元测试顺利通过了。
00:09:01包括困惑度测试等指标看起来也很正常,因为短文本生成差异并不明显。
00:09:08但在长文本生成测试中,这个隐患就暴露出来了。
00:09:11起初我完全摸不着头脑,不知道原因出在哪里。
00:09:14根源就在于使用了隐式合并,导致其中一个流实际上还没有完成计算。
00:09:24从而引发了竞态条件(race condition),读取到了未完成矩阵乘法的历史残留数据。
00:09:31为了解决这个问题,核心思路是必须采用显式的流合并。
00:09:40明确等待其中一项操作彻底结束,从而确保在合并时不会读取到陈旧的历史数据。
00:09:48这就是在探索 CUDA 流(CUDA streams)过程中得出的重要认识。
00:09:54回顾一下我最初的处理方式:
00:09:57流的合并是隐式的,
00:10:00导致后缩放(post scale)步骤读取到了旧缓冲区中的脏数值。
00:10:06修正后的方案如下:你需要明确标记矩阵乘法结束的位置,
00:10:14接着标记 RMS 计算结束的位置,
00:10:17然后让后缩放步骤先等待第一个流,
00:10:21再等待第二个流完成。
00:10:24这样就成功修复了 Bug,让论文中的理论得以正常运行,模型输出也重回正轨。
00:10:33这可以说是学术研究视角下比较硬核有趣的一部分。
00:10:38但我同时也希望把这些想法付诸实践,对吧?
00:10:40去部署它、测试它,看看如何让它在更偏生产环境的场景中跑起来。
00:10:47大家也可以去阅读论文,查看里面所有的测试数据。
00:10:50虽然大部分测试是基于 LLaMA 模型完成的,但该方法同样适用于其他模型架构。
00:10:56例如,针对这篇论文中提到的第一项提议——权重折叠(weight folding),你可以直接调用仓库中的代码来实现。
00:11:10就像闪电转换一样,只需执行 flashify 指令即可自动完成。
00:11:13不过对于我提到的第二点,如果你想亲自尝试,就需要编写一些算子内核(kernel)代码,
00:11:19正如我刚才在示例中展示的那样。
00:11:22这里展示的是基于 LLaMA 模型的一些实验结果,大家可以参考其中包含的各种细节。
00:11:29比如仅使用延迟归一化会怎样,使用完整的融合算子又会怎样。
00:11:36因此,这里做了大量的底层探索实验,来验证论文中的各项设想。
00:11:41这些是在不同严谨度和细节层面上取得的实验数据。
00:11:48即便只是简单的权重折叠,也能带来一定程度的性能提升。
00:11:54而且这种方法能够无缝无缝兼容你在日常模型开发中使用的各种工具。
00:11:59这意味着你完全不需要去重新造轮子,或者从零开始构建。
00:12:05它能够完美支持 torch.compile,本质上就像是生成了一个新的检查点(checkpoint),仅此而已。
00:12:13FlashAttention 也是在不同的网络层级上运用了类似的技巧。
00:12:16同时,它也能够完美支持量化模型。
00:12:18所以将其应用到实际项目中非常酷,你可以轻松获得一个带有这种全新归一化层的模型。
00:12:27如果大家想要获取具体的实现细节和运行代码,可以关注 transformer-tricks 这个代码库。
00:12:34里面包含了各种代数优化技巧,以及我提到的这篇论文的实现。
00:12:41此外,在 Hugging Face 的模型仓库中,我也上传了应用该技术处理过的一些模型。
00:12:50大家可以通过 Hugging Face 链接直接获取这些模型并进行测试。
00:12:54利用这些 Hugging Face 模型,你还可以直接将它们部署到生产环境中。
00:13:00当我思考如何落地时,我意识到,既然科学研究阶段已经完成,并且已经在 Hugging Face 上发布了模型链接,
00:13:11那么利用 Superlink 的推理引擎来部署任意 Hugging Face 模型,无疑是一种非常便捷高效的选择。
00:13:19我们在黑客松活动中就尝试过这种做法,参赛者们会带来他们自己微调好的自定义 Hugging Face 模型或检查点。
00:13:27即使你只是想尝试某种改进模型的代数技巧版本、验证自己的研究设想,
00:13:37也完全可以直接上手测试,将其部署在集群上,完全不必去操心模型部署过程中那些繁琐的胶水代码(glue code)。
00:13:48这真的很棒。关键在于,当集群基础设施和模型推理引擎都完全开源时,
00:13:58你就能非常方便地去验证那些更具创新力的前沿研究设想,比如进行算子内核修改或尝试 flash norm 等技术;
00:14:12而如果是在无法掌控推理底层的租用端点(endpoint)上,做这些修改就会困难得多。
00:14:18因此,你需要的是既具备可移植性和灵活性、能支持你自由探索的工具,
00:14:23同时又具备生产就绪的稳定度,以便能够进行大规模测试。
00:14:27例如,你可以结合 Sy 来将该模型与其他模型进行组合。
00:14:33正如左上方所示,你可以将这些 flashified 优化模型与其他模型结合,来执行智能体(agentic)任务,
00:14:43从而构建出端到端的大型应用场景。
00:14:46Sy 的工作方式在于通过其生产级集群来协助你完成模型部署。
00:14:52如果想了解更多相关细节,大家也可以去查看 Sy 的代码库。
00:14:57此外,它还配备了更智能的队列机制,这对于使用较小规模模型的场景尤为实用;
00:15:03因为在进行 flash norm 开发时,我主要使用的是较小型的 LLaMA 模型,以及 Hugging Face 上的小型 Agent 模型。
00:15:11因此,能够有一种方式来部署小型模型,并让它们共享同一块 GPU 运行,不仅节省了昂贵的 GPU 硬件开销,
00:15:24还能实现不同模型(尤其是小型模型)之间的灵活切换,这一点非常具有实用价值。
00:15:30你还可以通过 API 对模型配置以及整个集群进行控制,
00:15:35这也极大地方便了开源研究,无需专门安排运维人员来提供技术支撑。
00:15:40所以这一点同样非常棒。
00:15:43同时你拥有对云端的完全控制权,这对于倡导开源权重、开源模型和开源生态来说非常关键。
00:15:50另外,Sy 还提供了一个包含多种模型的目录,除了我刚才提到的这些,
00:15:57大家也可以去浏览探索。
00:15:59如果你正在开发相关应用,里面还包含了重排序(re-ranking)和嵌入(embedding)模型。
00:16:03到这里,我的研究历程分享就告一段落了。我共同撰写了这篇关于提升 Transformer 性能技术的论文,
00:16:15同时也找到了一条将该技术推向生产落地、进行实际测试,并深入探索开源模型应用的方法。
00:16:24如果有任何疑问或想要贡献代码,欢迎随时在 LinkedIn 上与我取得联系。
00:16:30我提到的很多内容,比如部分代码已经以 PR 的形式提交到了 vLLM 或 Hugging Face 社区。
00:16:36大家可以在各大平台看到它们。
00:16:38大家也可以去阅读一下论文。
00:16:40屏幕上展示的就是 arXiv 的论文链接。
00:16:43上面还有 Sy 的代码库地址以及我的 LinkedIn 个人主页。
00:16:47非常感谢大家的聆听!
00:16:58如果有任何问题,欢迎随时来找我提问。
00:16:59我们稍后
00:17:00就在附近。
00:17:13谢谢大家。
커뮤니티 글
아직 글이 없습니다. 이 영상에 대한 첫 번째 글을 작성해 보세요!
이 영상에 대해 글쓰기