https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/4548253a09264c5f53afa1b9eba10ef0_29.png

  1. 图像-标题对:从网络(如谷歌图片)爬取的带有描述性文本的图像。

  2. 交错图文数据:网页、PDF、学术论文中的图文混排内容。这是数据量最大但解析最困难的一类。

  3. 教科书与练习题:包含图表、公式的学术材料,对提升模型推理能力至关重要。

  4. 合成数据(Chart QA/Table QA):使用代码或语言模型自动生成的图表、表格及其相关问题。这是可无限扩展的高质量数据源。

  5. 文档布局理解数据:对网页或文档截图进行人工标注,用于训练模型理解UI元素、布局结构。

  6. OCR(光学字符识别)数据:将PDF页面作为图像输入,原始文本/LaTeX作为输出,训练模型“阅读”文档。

为了评估视觉语言模型的能力,业界也建立了一系列基准测试,它们与上述训练数据紧密对应:

  • MMMU:涵盖大学级别多学科问题的综合基准。

  • Chart QA / DocVQA:测试图表和文档理解能力。

  • TextVQA:测试图像中文本的阅读和理解能力。

  • OCR基准:测试光学字符识别精度。

  • 视觉细节理解:测试高分辨率图像下的细粒度识别能力。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/4548253a09264c5f53afa1b9eba10ef0_31.png


总结 🎯

本节课中我们一起学习了视觉语言模型的核心内容。

我们首先了解了视觉语言模型的基本架构,即通过一个视觉编码器将图像转换为语言模型可处理的序列。接着,我们深入探讨了两种主流的图像编码器:基于VQ-VAE的编码器基于CLIP的编码器。VQ-VAE通过向量量化将图像离散化为词元序列,支持图像生成;而CLIP通过对比学习得到连续的图像向量序列,语义对齐更好,但不支持直接图像生成。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/4548253a09264c5f53afa1b9eba10ef0_33.png

最后,我们认识到对于视觉语言模型乃至所有大模型而言,高质量、多样化的训练数据是性能提升的关键,并介绍了几类重要的训练数据源和评估基准。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/4548253a09264c5f53afa1b9eba10ef0_35.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/4548253a09264c5f53afa1b9eba10ef0_37.png

视觉语言模型是迈向多模态通用人工智能的重要一步,它将视觉感知与语言推理相结合,极大地扩展了模型的理解和应用边界。

15:缩放定律 (第二部分)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_1.png

在本节课中,我们将探讨一个在大型语言模型训练中至关重要但通常被视为商业机密的概念:缩放定律。我们将学习两种核心的缩放定律,它们能帮助我们科学地预测模型性能与模型规模、训练计算量之间的关系。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_3.png


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_5.png

缩放定律概述

训练大型语言模型的计算成本极高。例如,GPT-4是一个拥有1.6万亿参数的模型,在约100万亿个令牌上进行训练。在投入如此巨大的资源之前,我们需要进行预测:对于一个给定规模的模型,投入多少计算量能获得怎样的性能?这就是缩放定律要解决的问题。它通过在较小规模的模型上进行实验,总结规律,并外推预测大规模模型的行为。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_7.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_9.png


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_11.png

容量缩放定律:模型能记住多少知识?

上一节我们介绍了缩放定律的基本概念,本节中我们来看看第一种具体的定律:容量缩放定律。它回答一个核心问题:一个拥有 M 个参数的模型,最多能记住多少条事实性知识?

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_13.png

如何定义和测量“知识”?

为了精确研究,我们需要一个可控的实验环境。我们不会在真实的互联网数据上进行,因为几乎无法精确统计模型记住了多少知识。

我们采用一种简化的“传记”数据。每条知识被定义为一个 (名称,属性,值) 三元组。例如:(哈佛大学,所在地,马萨诸塞州剑桥市)。我们创建一批合成传记数据,每条包含6个固定属性(如生日、大学等),属性值随机生成。这样,知识的总量和结构是完全已知且可控的。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_15.png

核心测量方法

  1. 用纯传记数据训练一个语言模型。

  2. 训练后,向模型提问(例如,“某人的生日是什么?”),并测量其回答的准确率(精确匹配)。

  3. 模型能正确回答的知识数量,即为其记忆的知识量。

信息理论边界

模型记忆知识的方式可以很“聪明”。例如,如果所有人的生日只集中在10个日期,模型可以建立一个映射表,而不是为每个人单独存储日期字符串。因此,我们需要一个基准来衡量模型记忆效率的上限:即信息理论最优编码所需的最少比特数。

对于我们的传记数据结构,存在一个数学公式可以计算这个最优比特数 B_opt。它取决于多个因素,例如知识条目数 N、属性值的多样性、以及模型达到的损失值 L。公式大致形式如下(具体系数取决于数据特性):

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_17.png

B_opt ≈ C1 * N * L + C2 * N * log(D) + C3 * T * log(L)

其中,D 代表属性值的可能取值数量,T 与文本长度相关,L 是交叉熵损失。这个 B_opt 值代表了存储这些知识所需信息量的理论下限,我们将用它来衡量模型实际记忆的“知识量”。

实验结果与惊人发现

我们在不同规模(从2600万到10亿参数)的Transformer模型上进行实验,并充分训练它们(约1000个数据轮次),直至性能不再提升。

以下是关键发现:

  • 2比特/参数定律:无论模型架构细节如何(层数、头数、MLP与注意力层的比例),所有充分训练的Transformer模型,其有效记忆容量都趋近于 每参数2比特。这意味着一个 M 个参数的模型,最多能记忆约 2M 比特的最优编码信息。

  • 量化几乎无损:使用FP8甚至INT8格式训练模型,依然遵循相同的2比特/参数定律。这表明低精度训练是可行且高效的。

  • 对GPT-4的启示:GPT-4约有1.6万亿参数,其理论记忆容量约为 3.2万亿比特。而估算所有人类教科书知识总量大约在 200亿比特 左右。因此,仅从记忆所有教科书知识的角度看,一个约100亿参数的模型就足够了。GPT-4的规模远超此需求。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_19.png

数据质量的关键影响

然而,上述理想定律的前提是训练数据纯净(全是需要记忆的知识)。如果数据中混杂了大量“垃圾”信息(无用或随机的互联网文本),模型的训练效率会急剧下降。

实验表明,如果信号(有用知识)与噪声(垃圾信息)的比例为 1:7,那么要达到接近最优的容量,所需的训练数据量(或等效训练步数)将增加约 64倍。因此,高质量的数据过滤对于高效训练至关重要

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_21.png


计算缩放定律:如何最优分配计算资源?

上一节我们了解了模型容量的上限,本节中我们来看看在有限计算预算下,如何最有效地训练模型。这就是由DeepMind提出的 Chinchilla 缩放定律(或称计算缩放定律)。

它研究模型规模(参数数量 N)、训练数据量(令牌数 D)和最终模型损失 L 三者之间的关系。其核心结论是:在固定的总计算预算 C ≈ 6ND(用于训练前向和反向传播)下,模型参数数量和训练数据量应该成比例地增长

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_23.png

一个经验性的最优比例是:每10亿参数,大约需要20亿个训练令牌。例如,一个70亿参数的模型,最优训练数据量应在1.4万亿令牌左右。

对当前实践的批判与启示

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_25.png

许多现有模型的训练并未遵循此定律:

  • 训练不足:例如,一些大型模型(如700亿参数)的训练令牌数可能和较小模型(如70亿参数)一样多。根据Chinchilla定律,这意味著大型模型训练不足,没有发挥其全部潜力。

  • 规模效率悖论:一个反直觉但关键的推论是:扩大模型规模有时能提高训练效率。如果一个模型处于容量临界点(刚好能记住所有数据),其参数必须被极度压缩利用,优化过程会变慢。而一个规模更大的模型(“过度参数化”)有更宽松的优化空间,可能用更少的训练步数就能达到相同的性能。因此,在计算受限时,选择比理论最优容量更大的模型,可能是更快的训练策略。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_27.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_29.png

总结

本节课我们一起学习了两种核心的缩放定律:

  1. 容量缩放定律:揭示了Transformer模型记忆能力的理论上限(约每参数2比特),并强调了高质量训练数据过滤的极端重要性。

  2. 计算缩放定律:指出了在固定计算预算下,模型规模与训练数据量应平衡缩放,并解释了为何在实践中常常会使用“过度参数化”的模型以加速训练。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/8697c78d15511a0826ffca2d847dc936_31.png

这些定律是指导大型语言模型研发的科学基础,帮助我们在投入海量资源前进行性能预测和资源规划。尽管真实的工业级缩放定律更为复杂且保密,但其核心思想均源于此类基础研究。

16:扩展技术 II - 混合专家模型 (Mixture of Experts)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_1.png

在本节课中,我们将学习一种称为混合专家模型 (Mixture of Experts, MoE) 的特殊架构。这种架构在当前非常重要,因为所有大规模模型实际上都使用它来加速推理时间。我们将了解MoE的直观原理、其为何有用,以及成功训练MoE层所需的一些计算技巧。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_5.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_7.png

回顾与动机

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_9.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_11.png

上一节我们介绍了扩展语言模型的好处:它能提升模型容量,并且令人惊讶的是,在总计算量(FLOPs)方面,更大的模型收敛得更快。因此,从训练成本角度看,我们希望模型尽可能大。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_13.png

然而,这些大模型与Hessian矩阵模型相比有一个缺点:推理成本要高得多。例如,一个7B参数的模型推理成本远低于一个70B参数的模型。在模型部署时,我们每天需要为数百万甚至数亿用户提供服务,这是一个巨大的推理开销。因此,我们希望在不牺牲模型容量或训练优势的前提下,尽可能降低推理成本。

接下来的课程将介绍一些降低语言模型推理成本的技术。本节课重点讨论混合专家模型架构,它主要用于替换或高效实现Transformer中的MLP层。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_15.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_17.png

Transformer架构回顾

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_19.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_21.png

一个标准的Transformer块主要由两种类型的层构成:自注意力层和前馈网络层。MoE层将作为MLP层的一个更高效的替代品,用于加速推理。下一讲我们将学习使注意力层推理更快的技巧。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_23.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_24.png

目前,所有大型语言模型都使用了MoE架构,例如GPT-4、GPT-3.5、Mixtral和Google的Gemini 2等。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_26.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_27.png

MoE的直观理解

在Transformer中,MLP层更像是知识的存储单元,存储事实性知识,而自注意力层则负责逻辑推理。人类的实际知识是非常稀疏的。当我们向模型提问时,模型通常只需要提取与问题相关的一小部分知识。

MoE的核心理念是将MLP层划分为多个“专家”块,每个专家可能专注于某一类知识。在推理时,根据输入内容,模型只激活并使用少数相关的专家进行计算,而忽略其他专家。这就像为每个问题只调用相关的知识库。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_29.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_31.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_32.png

混合专家层定义

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_34.png

一个MoE层包含以下组件:

  • 输入:一个维度为 d 的向量 x

  • 路由器:一个线性层,将输入 x 映射到维度为 M 的向量,其中 M 是专家数量。

  • 门控值:对路由器输出应用Softmax,得到一个概率分布 s(x)

  • 专家选择:根据 s(x) 选择概率最高的前 K 个专家(通常 K=2)。

  • 专家网络:每个专家本身是一个标准的MLP层。

  • 输出计算:最终输出是所选专家输出的加权和,权重是所选专家门控值的重新归一化结果。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_36.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_37.png

公式表示如下:

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_39.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_41.png

  1. 路由器输出:g(x) = W_router * x

  2. 门控概率:s(x) = Softmax(g(x))

  3. 选择函数:I_k(x) 表示第 k 个被选中的专家索引。

  4. 重新归一化权重:s'(x)_i = s(x)_i / (sum_{j in I_k(x)} s(x)_j),仅对选中的专家 i 有效。

  5. 最终输出:MoE(x) = sum_{i in I_k(x)} s'(x)_i * Expert_i(x)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_43.png

在Transformer中,MoE层替换了原有的MLP层,并对序列中的每个令牌独立应用。所有令牌共享同一个MoE层,但不同令牌可能激活不同的专家子集。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_45.png

MoE的参数与效率

对于MoE模型,需要区分两个概念:

  • 总参数量:模型中所有参数的总和,包括所有专家的参数。

  • 有效参数量:每个令牌前向传播时实际激活的参数数量,约为 K * (单个专家参数量)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_47.png

MoE的关键优势在于,其知识存储的缩放律(约1.5比特/参数)是基于总参数量的,而推理成本仅与有效参数量相关。这意味着你可以用一个总参数量巨大(从而知识容量大)的模型,但每个令牌的推理只涉及其中一小部分,实现了效率与容量的平衡。

MoE的实现挑战

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_49.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_51.png

简单地用PyTorch实现上述公式无法利用稀疏性获得加速,因为框架仍会计算所有专家。为了实现高效计算,需要采用“快速编解码”策略。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_53.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_55.png

以下是实现思路:

  1. 编码:对于一个批次中的令牌序列,根据路由器结果,将使用相同专家的令牌聚集到一起,形成一个规整的张量。

  2. 专家计算:将每个聚集后的张量发送给对应的专家进行批量计算。

  3. 解码:将专家计算后的结果根据原始令牌顺序重新组装回去。

为了处理专家间负载可能不均衡的问题,通常会设置一个“容量因子”(如1.1-1.25),为每个专家分配略高于平均预期的计算缓冲区,以避免令牌被丢弃。

专家并行化

MoE架构天然适合一种称为“专家并行化”的分布式训练策略。我们可以将不同的专家放置在不同的GPU上。在前向和反向传播时,只需将令牌数据路由到对应的GPU上进行专家计算,专家之间无需频繁通信。这比传统的张量并行或流水线并行效率更高,是训练超大规模MoE模型的关键。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_57.png

负载均衡损失

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_59.png

在训练MoE时,我们希望所有专家都能被相对均匀地使用,以避免某些专家过载而其他专家未被充分训练。为此,我们会在损失函数中添加一个“负载均衡损失”。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_61.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_63.png

负载均衡损失公式

L_balance = sum_i (f_i * p_i)

其中:

  • f_i 是实际使用专家 i 的令牌比例。

  • p_i 是所有令牌的门控概率在专家 i 上的总和。

最小化这个损失会鼓励路由分布和专家使用频率都趋向均匀。这是稳定训练MoE模型的重要技巧。

总结

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/ec99548739435bcc246c778158ac54a1_65.png

本节课我们一起学习了混合专家模型。我们了解了MoE如何通过为每个输入动态选择少数专家来大幅降低推理成本,同时保持庞大的总参数量以存储海量知识。我们还探讨了MoE的关键实现技术,如快速编解码和专家并行化,以及训练稳定所需的负载均衡损失。MoE是目前构建超大规模高效语言模型不可或缺的核心架构之一。

17:注意力机制扩展(第三部分)- 注意力优化 🚀

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_1.png

在本节课中,我们将学习如何更高效地计算注意力单元。我们将重点介绍两种方法:Flash Attention 和 Multi-Query Attention。这些技术能显著加速注意力计算并大幅降低内存使用量,尤其是在处理长序列时。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_4.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_6.png


注意力机制回顾

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_8.png

上一节我们介绍了如何通过稀疏化来加速MLP层的计算。本节中,我们来看看如何优化注意力层的计算。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_10.png

首先,让我们回顾一下注意力层的基本结构。多头注意力层由多个注意力头组成。给定输入向量 V₁Vₙ(其中 n 是序列长度,每个向量的维度是 d),注意力层的输出需要经过一系列计算。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_12.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_13.png

注意力公式 可以表示为:

输出_i = Σ_j (softmax( (Q_i * K_j^T) / sqrt(d_k) ) * V_j)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_15.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_17.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_19.png

其中,QKV 分别是查询、键和值矩阵。本质上,你需要计算一个 n × n 的注意力分数矩阵,然后对每一行进行 softmax 归一化,最后根据归一化后的权重对值向量进行加权求和。

你可以将 V_i 视为一个查询。对于一个查询,你需要计算它与所有键的内积,找出最相似的键。然后,该查询的输出将是相似区域值的加权平均。这就像根据当前词与文档中其他词的相似性来理解其上下文。

主要问题是,这个 n × n 的注意力矩阵在计算和内存使用上都是 O(n²) 的。当序列长度 n 很大时(例如 100k),这个矩阵会变得极其庞大,内存消耗会非常高。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_21.png


注意力层的计算与内存瓶颈

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_23.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_25.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_27.png

在标准的 Transformer 块中,包含注意力层、层归一化、残差连接和 MLP(或 MoE)层。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_29.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_30.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_31.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_33.png

计算时间分析

  • MLP层:计算时间约为 n × d²

  • 多头注意力层:计算时间约为 n² × d + n × d²。其中 n × d² 部分与 MLP 层类似,但 n² × d 项在序列很长时会占主导。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_35.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_36.png

内存使用分析

  • MLP层:只需存储每个 token 的激活值,内存使用约为 n × d

  • 多头注意力层:需要存储 n × n × M 的注意力矩阵(其中 M 是注意力头数,通常与 d 成线性关系)。因此,内存使用约为 n² × d,这在序列很长时会成为主要瓶颈。

所以,注意力层在长序列处理上,无论是计算时间还是内存使用,都面临巨大挑战。我们的目标是学习一种名为 Flash Attention 的技术,它能将注意力层的内存使用从 O(n²) 有效降低到几乎 O(n)


Flash Attention 的核心思想

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_38.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_40.png

Flash Attention 的关键思想是:不将整个 n × n 注意力矩阵一次性存储在内存中,而是只存储其中的一部分,并通过更巧妙的计算和重计算来避免存储全部中间结果

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_42.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_44.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_46.png

为了理解其原理,我们先看一个注意力头中单行的计算。问题是:能否在不将所有中间值存入内存的情况下,计算这一行的输出(即加权和)?

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_48.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_49.png

一个简单但低效的方法是顺序计算加权和,只维护一个累加器,而不存储所有中间分数。但这会带来两个问题:1) 顺序循环导致计算慢;2) 数值稳定性问题(softmax 中的指数运算可能导致浮点数溢出)。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_51.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_52.png

以下是解决数值稳定性问题的关键技巧:在计算 softmax 时减去最大值

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_54.png

稳定 softmax 计算伪代码

# 输入: 向量 x = [x1, x2, ..., xn]
m = max(x) # 找到最大值
exp_x_shifted = exp(x - m) # 每个元素减去最大值后求指数
sum_exp = sum(exp_x_shifted)
softmax_x = exp_x_shifted / sum_exp

这样做确保了指数运算的结果在 (0, 1] 范围内,避免了溢出,并且不改变最终的 softmax 结果。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_56.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_58.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_60.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_62.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_63.png

然而,上述方法需要遍历数据两次(一次求最大值,一次计算加权和),或者需要将数据存入内存。Flash Attention 的精妙之处在于引入了 “运行最大值” 的概念,使得只需遍历数据一次即可同时完成稳定化的加权和计算。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_65.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_66.png

Flash Attention 核心算法思想(简化版)

初始化 O = 0向量, l = 0, m = -inf
for i in 1 to n:
    xi = 计算当前查询与第i个键的内积
    m_new = max(m, xi)
    # 根据新的运行最大值更新累加器
    l = l * exp(m - m_new) + exp(xi - m_new)
    O = O * exp(m - m_new) + Vi * exp(xi - m_new)
    m = m_new
最终输出 = O / l

这个算法只存储了运行最大值 m、归一化因子 l 和输出累加器 O,内存使用很低,且只遍历数据一次。但它仍然是一个顺序循环,没有利用硬件擅长的快速矩阵乘法。


Flash Attention V2:结合分块与快速矩阵乘法

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_68.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_70.png

最终的 Flash Attention V2 算法结合了所有优点:

  1. 分块处理:将长的序列分成较小的块(例如 96x96)。

  2. 块内快速计算:在每个块内部,使用标准的、高效的矩阵乘法来计算注意力分数和 softmax。因为块很小,所以即使存储中间矩阵,内存开销也有限,同时能利用硬件对特定尺寸矩阵运算的优化。

  3. 块间顺序更新:像核心算法思想一样,在块与块之间使用运行最大值和累加器进行顺序更新,确保数值稳定性并避免存储全局大矩阵。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_72.png

高级别流程

将 Q, K, V 矩阵分成多个块
初始化全局累加器 O_global, l_global, m_global
for 每个块 j:
    # 1. 使用快速矩阵乘法计算当前块内的注意力分数 S_block
    # 2. 在块内计算 softmax (使用块内最大值进行稳定化)
    # 3. 计算当前块的输出 O_block
    # 4. 用运行最大值方法,将 O_block 与全局累加器 O_global, l_global, m_global 合并
最终输出 = O_global / l_global

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_74.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_75.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_77.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_79.png

这个算法在数学上被证明能产生与原始注意力计算完全相同的精确结果。它极大地减少了内存占用(从 O(n²)O(n)),同时通过分块利用了硬件的计算效率。目前,Flash Attention 已集成在 Hugging Face Transformers 库中,是加速推理和训练的关键技术。


其他注意力优化方法简介

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_81.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_83.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_85.png

除了 Flash Attention,还有其他用于加速注意力计算的方法:

  • Multi-Query Attention (MQA):让多个注意力头共享同一个键(K)和值(V)矩阵,显著减少计算量。例如,Mistral 模型就使用了 MQA。

  • Sliding Window Attention:每个 token 只关注其附近一个窗口内的 token,而不是整个序列,将计算复杂度从 O(n²) 降到 O(n × w),其中 w 是窗口大小。

  • Blockwise Attention:将序列分成块,主要进行块内的注意力计算,减少长距离依赖的计算开销。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_87.png

总结

本节课中,我们一起深入探讨了如何优化 Transformer 中注意力机制的计算。

  • 我们回顾了注意力层在长序列下面临的 O(n²) 计算和内存瓶颈。

  • 我们详细剖析了 Flash Attention 的核心思想:通过避免存储完整的注意力矩阵,并利用运行最大值、分块计算和顺序更新等技巧,在保证数值稳定的同时,将内存占用降至 O(n),并保持了计算效率。

  • 我们还简要了解了其他如 Multi-Query Attention 等优化方法。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/f493c94c2d33e2fa48c94b8ea8760dd3_89.png

掌握这些优化技术对于理解和部署高效的大语言模型至关重要。

18:分布式训练

在本节课中,我们将学习生成式AI中一个至关重要的主题:分布式训练。随着语言模型和生成模型规模的不断扩大,传统的单GPU训练方式已无法满足需求。分布式训练通过在多块GPU上协同工作,解决了大模型训练中的内存和计算瓶颈。我们将从基础概念入手,逐步探讨几种关键的分布式训练技术。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_1.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_5.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_7.png


数据并行:最基础的分布式形式

上一节我们介绍了单GPU训练的简单流程。本节中我们来看看当拥有多块GPU时,最直接的加速方法:数据并行。

在数据并行中,每块GPU都拥有完整的模型副本,但处理不同的数据批次。以下是其核心步骤:

  1. 数据分发:将训练数据批次分割,每块GPU获得一个子批次。

  2. 独立前向与反向传播:每块GPU使用自己的数据独立完成前向传播和损失计算,并计算梯度。

  3. 梯度同步:所有GPU计算出的梯度通过通信操作(如 all_reduce)进行求和或平均。

  4. 参数更新:每块GPU使用同步后的梯度统一更新其模型参数。

数据并行的核心优势是通信开销小,通常只需在梯度同步时进行通信。在PyTorch中,这可以通过一行代码实现:

model = torch.nn.DataParallel(model)

然而,数据并行要求每块GPU都能在内存中容纳完整的模型参数、梯度和优化器状态。对于当今动辄数十亿参数的大模型,这变得不可能。


内存挑战:为何需要更高级的技术

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_9.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_10.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_12.png

上一节我们了解了数据并行的局限性。本节中我们来看看驱动更高级分布式训练技术的核心挑战:GPU内存限制。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/3aced6959747b3a5105a4f8c95d93c73_14.png

训练一个大语言模型时,主要的内存消耗来自三个方面:

  1. 模型参数:存储模型权重,通常使用 FP16BF16 精度以节省内存。一个70亿参数的模型约需 14 GB 内存(FP16)。

  2. 优化器状态:对于像Adam这样的优化器,需要存储动量和二阶动量。为了数值稳定性,这些状态通常以 FP32 精度存储。对于一个70亿参数的模型,优化器状态约需 56 GB 内存。

  3. 激活值(Activations)与梯度:在前向传播中产生的激活值,以及在反向传播中计算的梯度,也会消耗大量内存,其大小与模型参数和批次大小相关。

简单相加,仅存储一个70亿参数模型的权重和优化器状态就可能超过 70 GB,这已经接近甚至超过单块高端GPU(如H100的80GB)的内存容量,尚未计算激活值和梯度所需的内存。因此,无法在单GPU上训练此类模型。


优化器状态分片:减少内存占用的第一步

上一节我们明确了内存是主要瓶颈。本节中我们来看看第一种高级技术:优化器状态分片(也称为DeepSpeed ZeRO阶段1)。

这项技术的核心思想是:不再在每块GPU上存储完整的优化器状态,而是将其均匀分片,每块GPU只存储其中一份

以下是其工作流程:

  1. 模型复制:每块GPU仍然存储完整的模型参数副本。

  2. 优化器状态分片:将优化器状态(如Adam的动量和二阶动量)分割成N份(N为GPU数量),每块GPU存储其中一份。

  3. 梯度计算与同步:每块GPU独立计算其数据子批次对应的梯度。然后,通过一个 reduce_scatter 通信操作:

    • Reduce(规约):将所有GPU上对应同一参数分片的梯度进行求和。

    • Scatter(散射):将求和后的梯度分片发送回负责该参数分片优化器状态的GPU。

  4. 参数更新:每块GPU使用接收到的梯度分片和本地存储的优化器状态分片,更新其负责的那部分模型参数

  5. 参数同步:更新完成后,通过通信(如 all_gather)将更新后的参数分片同步到所有GPU,确保每块GPU上的完整模型参数保持一致。

通过分片优化器状态,我们将其内存占用从 56 GB 降低到了约 56/N GB。对于8块GPU,这意味着一块70亿参数模型所需的内存从约70GB降至约 14(参数) + 7(优化器分片) + 梯度内存 ≈ 28 GB 左右,使得训练成为可能。


完全分片数据并行:进一步分片模型参数

上一节我们通过分片优化器状态节省了大量内存。本节中我们来看看更激进的方案:完全分片数据并行,它同时分片模型参数、梯度和优化器状态。

FSDP是优化器状态分片的自然延伸。其核心思想是:

  • 不仅将优化器状态分片,也将模型参数本身进行分片存储。

  • 每块GPU只存储整个模型的一个参数子集以及对应的优化器状态分片。

其工作流程更为复杂:

  1. 前向传播:当计算需要某些不在本地的参数时,通过 all_gather 通信从其他GPU收集这些参数,在本地临时组装成完整层进行计算,计算后释放这些临时参数。

  2. 反向传播:类似地,在计算梯度时,再次通过 all_gather 组装所需参数。计算出的梯度根据其对应的参数分片进行分片。

  3. 梯度同步与更新:使用类似优化器分片中的 reduce_scatter 操作同步梯度分片。每块GPU然后使用本地存储的参数分片、梯度分片和优化器状态分片进行更新。

  4. 参数同步:更新后的参数分片可能需要同步(取决于实现)。

FSDP的优势在于:

  • 内存效率极高:理论上可以训练模型的大小与GPU总数成线性比例。

  • 支持更高精度:可以在本地以 FP32 精度维护参数分片和优化器状态,进行更精确的更新,同时在前向/反向传播时使用 FP16/BF16 以节省内存和加速计算。

在PyTorch中,FSDP可以通过一个包装器使用:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model)

张量并行:将计算图本身进行分割

上一节介绍的FSDP通过在垂直方向(参数)上分片来节省内存。本节中我们来看看另一种维度:张量并行,它通过水平分割单个层的计算来分布内存和计算负载。

张量并行的核心是将一个大型神经网络层(如线性层)的矩阵运算分布到多个GPU上。以线性层 Y = XA 为例(X 输入,A 权重矩阵)。

主要有两种分割方式:

  1. 列并行(Column Parallel):将权重矩阵 A 按列分割。每块GPU持有 A 的一部分列。计算时,每块GPU计算 X 与本地权重子矩阵的乘积,得到输出子部分,然后通过 all_gather 通信将所有GPU的输出子部分拼接成完整的 Y

    • 前向传播公式Y = [X * A1, X * A2, ...],然后 all_gather

    • 内存:每块GPU存储的 A 的大小变为原来的 1/N

  2. 行并行(Row Parallel):将权重矩阵 A 按行分割。此时,需要先将输入 X 通过 all_gather 广播到所有GPU(或更高效地,按行分割 Xall_gather 结果)。每块GPU计算本地 X 子部分与本地 A 子矩阵的乘积,然后通过 reduce_sum 通信对所有GPU的结果求和,得到完整的 Y

    • 前向传播公式:每块GPU计算 Yi = Xi * Ai,然后对所有 Yi 进行 reduce_sum

在实践中,一个线性层的前向传播可能采用列并行,而为了高效地进行反向传播,其对应的梯度计算会自然地采用行并行。像Megatron-LM这样的库提供了封装好的并行线性层,用户只需用它们替换标准线性层即可构建张量并行模型。

张量并行的主要挑战是通信频繁,因为几乎每一层的前后向传播都需要GPU间通信。它通常用于模型实在太大,无法用FSDP装入单个GPU内存的情况。


混合并行与总结

在实际的超大规模模型训练中(如训练GPT-4或Llama),通常会混合使用多种并行策略:

  • 数据并行:在不同的GPU组上处理不同的数据批次。用于扩大有效批次大小。

  • 张量并行:在单个GPU组内,将大型模型层拆分到多个GPU上。用于解决单层参数过大的问题。

  • 流水线并行:将模型的不同层组放置在不同的GPU上。一个批次的数据像流水线一样依次经过这些GPU。用于解决模型深度过大的问题。

  • 专家并行:用于混合专家模型,将不同的专家分布到不同的GPU上。

系统会根据硬件拓扑(如NVLink连接)和模型结构,将数千块GPU划分成不同的并行组,以最大化计算效率和最小化通信开销。

本节课中我们一起学习了分布式训练的核心概念。我们从最简单的数据并行开始,揭示了其内存瓶颈,进而深入探讨了优化器状态分片、完全分片数据并行和张量并行等高级技术。这些技术通过巧妙地分割模型参数、优化器状态和计算图,使得训练拥有数千亿参数的大模型成为可能。理解这些原理是从事前沿大模型开发和优化的基础。

20:大语言模型的长上下文处理(第二部分)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_1.png

在本节课中,我们将探讨大语言模型中长上下文处理的核心技术。我们将了解为何长上下文至关重要,以及如何通过序列并行化、改进的注意力机制和位置编码外推等技术来高效地训练和使用具有长上下文能力的模型。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_3.png

为什么长上下文很重要?

上一节我们介绍了长上下文的基本概念。本节中,我们来看看长上下文为何成为现代大语言模型的核心能力。

目前,大语言模型最重要的应用之一是作为智能体,阅读和理解企业文档,例如公司政策或工具使用说明。模型需要能够从这些文档中提取信息并可靠地回答相关问题。

例如,如果你向一个支持100万令牌上下文的模型输入一本“书”,然后询问“要完成XX任务,我应该使用公司的哪个特定工具?”,模型需要能够从这海量信息中准确找到答案。这对于新员工快速熟悉公司架构和工具非常有价值,也是当前语言模型商业化的主要方向。

因此,长上下文能力是当前大语言模型产品的核心竞争力,仅仅支持4K上下文的模型已无法满足市场需求。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_5.png


长上下文训练的挑战与策略

我们了解到长上下文很重要,但直接训练支持长上下文的模型并非易事。本节将探讨其中的挑战和主流训练策略。

回顾之前的内容,标准的语言模型训练通常分为三个阶段。在第一阶段(预训练),出于速度和效率考虑,上下文长度通常限制在4K。如果一开始就使用过长的上下文(如16K以上),模型在训练初期会感到“困惑”,因为它不知道在预测下一个词时应该关注前文中的哪个部分,导致损失难以下降。

较短的上下文强制模型专注于局部一致性,使训练更容易。因此,长上下文的扩展通常在第二阶段进行,此时模型已经理解了语言的基本规律。第二阶段通常是长短上下文混合训练,以保持模型在短上下文任务上的性能。第三阶段(后训练)则保持长上下文进行微调。


关键技术:序列并行化

要训练支持长上下文的模型,首先需要解决计算和内存的挑战。本节介绍一个关键的基础技术:序列并行化。

对于一个大模型(例如70亿参数),如果使用8K上下文,在不进行任何并行优化的情况下,一个微批次就可能占满一张H100 GPU(80GB)的内存。当上下文长度扩展到128K时,内存消耗可能增加16倍以上,单卡根本无法容纳。

序列并行化是一种高效的解决方案。其核心思想是沿着序列维度对训练样本进行切分。

假设我们有一个长度为128K的序列。我们可以将其切分成多个片段,例如:

  • GPU 0 处理第 1 到 8K 个令牌。

  • GPU 1 处理第 8K 到 16K 个令牌。

  • 以此类推。

对于MLP层,每个令牌的计算是独立的,因此序列并行化不需要额外的通信。主要的通信开销发生在注意力层,因为每个令牌在计算注意力时可能需要关注其他GPU上的令牌。为了减少这种通信开销,需要结合使用改进的注意力机制,如滑动窗口注意力、稀疏注意力或环状注意力。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_7.png


改进的注意力机制

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_9.png

序列并行化解决了内存问题,但注意力计算本身仍是瓶颈。本节我们看看几种能提升长上下文处理效率的注意力变体。

以下是几种常见的用于长上下文的注意力机制:

  1. 滑动窗口注意力

    • 公式/概念:每个令牌 i 只关注其前面一个固定窗口大小 W(例如2047)内的令牌,即关注 [i-W, i-1] 范围内的令牌。

    • 优点:极大减少了计算和通信量,因为注意力范围是局部的。

    • 缺点:纯局部注意力可能难以捕捉长距离依赖,尽管信息可以通过多层网络间接传递。

  2. 稀疏注意力

    • 公式/概念:在不同层使用不同的注意力模式。例如,第一层使用局部滑动窗口注意力;第二层则让令牌 i 关注 i-100, i-200, i-300 等间隔较远的令牌。

    • 优点:这种“跳跃式”连接非常适合信息检索任务。低层汇总局部窗口信息,高层则在这些汇总信息间进行搜索。

    • 缺点:实现更复杂,需要精心设计稀疏模式。

  3. 环状注意力

    • 概念:一种更极致的优化,确保通信只发生在相邻GPU之间,类似于滑动窗口,但设计上可能更高效。

结合序列并行化和这些稀疏注意力机制,可以有效地训练和运行支持长上下文的模型。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_11.png


位置编码的外推

解决了计算问题后,另一个关键挑战是位置编码。模型需要知道令牌的顺序。本节探讨如何将位置编码扩展到训练时未见过的超长上下文。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_13.png

标准的旋转位置编码(RoPE)可以理解为:将词嵌入向量的每两个维度视为一个复数,然后根据令牌的位置 pos 和该维度对应的旋转速度 θ_k 进行旋转。

  • 公式:对于位置 pos 上词嵌入向量的第 k 个二维分量 (x_k, y_k),旋转后变为:

    (x_k * cos(pos * θ_k) - y_k * sin(pos * θ_k), x_k * sin(pos * θ_k) + y_k * cos(pos * θ_k))

  • 直观理解:不同维度以不同速度旋转。快速旋转的维度擅长捕捉局部位置信息(周期短),慢速旋转的维度擅长区分全局位置(周期长)。

在预训练阶段,θ_k 的设置通常使最慢旋转维度的周期覆盖训练时的最大上下文长度(如4K或10K)。当我们需要将上下文长度扩展到远大于这个周期(如100K)时,直接使用原来的 θ_k 会导致周期外的位置无法被正确区分(因为 cossin 函数是周期性的)。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_15.png

因此,需要进行位置编码外推。一种在实践中有效的策略是非均匀缩放

  • 方法:不直接缩放所有位置的旋转角度,而是定义一个临界位置 N。对于位置 pos < N 的令牌,使用原始的旋转速度(保持短上下文性能)。对于 pos >= N 的令牌,让旋转速度随着位置增加而逐渐变慢。

  • 公式(概念性)θ_k’(pos) = θ_k * f(pos),其中 f(pos)pos < N 时为1,在 pos >= N 时是一个缓慢衰减的函数。

  • 优点:最大程度地保留了模型在短上下文上的性能,为长上下文微调提供了良好的起点。具体的 f(pos) 函数形式可以通过超参数搜索确定。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_17.png


生成长上下文训练数据

拥有了训练长上下文模型的技术,我们还需要合适的数据。本节介绍如何构建用于训练和评估的长上下文数据。

高质量、包含长距离依赖的自然文本数据很难大量获取。因此,目前主流的方法是合成数据生成

核心思路是“藏宝于海”:

  1. 从一个短的文档-问答对 (D, Q, A) 开始,其中 A 的答案明确依赖于文档 D

  2. 用大量其他无关或相似的文档将目标文档 D 包围起来,形成一个超长的合成文档。

  3. 将问题 Q 放在这个长文档的末尾,并要求模型给出答案 A

  4. 为了增加难度,甚至可以将目标文档 D 拆分成多个片段,分散插入到长文档的不同位置,要求模型进行信息聚合。

通过这种方式,可以大规模生成用于训练和评测模型长上下文理解与信息检索能力的数据集。


总结

本节课我们一起学习了实现大语言模型长上下文能力的核心技术。

我们首先了解了长上下文对于智能体应用的重要性。接着,探讨了分阶段训练的策略,即在模型掌握语言基础后再进行长上下文扩展。然后,我们深入研究了实现长上下文训练的关键技术:序列并行化 用于解决内存瓶颈,滑动窗口/稀疏注意力 用于降低计算和通信开销,以及位置编码的外推(特别是非均匀缩放策略)用于让模型理解超长序列中的位置信息。最后,我们介绍了通过合成数据生成来构建训练和评估数据集的方法。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/200a7f3602a0fb54c34b23be94b7ac8a_19.png

掌握这些知识,有助于理解当前主流大语言模型如何突破上下文长度的限制,以及在这一前沿领域进行探索和创新的可能方向。

21:视频生成模型 🎬

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_1.png

在本节课中,我们将学习如何将图像生成的扩散模型技术扩展到视频生成领域。我们将深入探讨其背后的数学原理,特别是扩散模型的理论基础,并了解OpenAI的Sora模型是如何应用这些原理的。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_3.png

扩散模型的数学基础

上一节我们提到了扩散模型在视频生成中的应用。为了理解其工作原理,本节中我们来看看扩散模型背后的核心数学原理。

扩散模型包含一个前向过程和一个反向过程。前向过程将一个给定的分布转换为高斯分布,而反向过程则将高斯分布转换回原始分布。其数学基础是所谓的Wasserstein梯度流

从数学角度看,概率分布可以视为一个度量空间中的点。这个空间的标准度量是Wasserstein度量。两个分布P和Q之间的Wasserstein距离(W₂)定义为:

W₂(P, Q) = inf_{(X, Y)} E[|X - Y|²]^(1/2)

其中,(X, Y) 是满足X服从P、Y服从Q的联合分布。这个度量衡量的是将分布P“移动”到分布Q所需的最小“工作量”。

扩散过程的前向过程,本质上就是沿着Wasserstein度量空间中的梯度流,将原始分布移动到高斯分布。这条路径是连接两个分布的最短路径。反向过程则是这条路径的逆过程。

在实践中,我们通过离散化(例如100个时间步)来近似这条连续的路径,类似于用梯度下降法近似梯度流。

前向过程与随机微分方程

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_5.png

上一节我们介绍了Wasserstein梯度流的概念。本节中我们来看看如何用具体的方程来描述这个过程。

Wasserstein梯度流有一个简洁的数学表征,即一个随机微分方程。假设我们想从初始分布P₀移动到最终的高斯分布P_∞。在时间t,变量X_t的分布为P_t,并且X_t满足以下关系:

X_t = e^{-t} * X_0 + sqrt(1 - e^{-2t}) * Z

其中,X₀ 服从初始分布P₀,Z 服从标准高斯分布。这意味着,前向过程本质上就是不断向原始数据添加噪声。

这个X_t作为时间t的函数,满足以下随机微分方程:

dX_t = -X_t dt + sqrt(2) dB_t

这里,dB_t是布朗运动。方程中的 -X_t dt 项使X_t收缩,而 sqrt(2) dB_t 项则不断添加噪声。随着时间t增大,X_t的分布逐渐趋近于高斯分布。

反向过程与福克-普朗克方程

上一节我们描述了前向过程的随机微分方程。为了生成数据,我们需要一个反向过程。本节中我们来看看如何推导反向过程。

我们的目标是从高斯噪声(X_T)开始,通过一个反向过程恢复出原始数据(X₀)。我们定义一个新的变量Y_t = X_{T-t}。那么Y_0 = X_T(近似高斯分布),Y_T = X₀(目标分布)。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_7.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_8.png

根据福克-普朗克方程,我们可以推导出Y_t满足的随机微分方程:

dY_t = [Y_t + 2 * ∇ log p_{T-t}(Y_t)] dt + sqrt(2) dB_t

其中,p_{T-t}(·) 是在时间 T-t 时变量的概率密度函数。与简单的前向过程相比,反向方程中多出了一项 2 * ∇ log p_{T-t}(Y_t),即概率密度函数对数的梯度,这被称为得分函数

因此,运行反向过程、从噪声生成数据的关键,就在于估计这个得分函数。

得分匹配与神经网络

上一节我们得出结论,反向过程的核心是估计得分函数。本节中我们来看看如何利用神经网络来实现这一点。

为了运行反向的随机微分方程,我们需要计算 ∇ log p_t(y)。我们可以训练一个神经网络 s(y, t) 来近似这个得分函数。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_10.png

一个重要的数学结论是:训练神经网络去预测添加到数据中的噪声,在数学上等价于最小化得分匹配的目标函数。也就是说,如果我们给一个带噪声的图像,让神经网络预测所添加的噪声,那么训练好的网络输出就是得分函数的一个近似。

因此,整个扩散模型的流程在数学上是严谨的:

  1. 前向过程:通过添加噪声将数据分布变为高斯分布。

  2. 训练:训练神经网络来预测给定噪声数据所对应的噪声。

  3. 反向过程:从高斯噪声开始,使用训练好的神经网络来近似得分函数,运行反向SDE,逐步生成数据样本。

这个方法适用于任何分布,只要我们能最小化得分匹配目标。这就是为什么它可以被推广到图像、视频、音频等各种生成任务。

从图像到视频的扩展

上一节我们明确了扩散模型是分布无关的通用方法。本节中我们来看看如何将其直接应用于视频生成。

视频可以看作是一系列图像的序列。因此,视频的分布就是图像序列的联合分布。从原理上讲,我们可以简单地将视频数据(所有帧的像素)视为一个非常高维的随机变量X₀,然后直接应用扩散模型。

然而,这里存在一个挑战:扩散模型的收敛速度和样本复杂度与随机变量的维度D大致成正比。对于视频,维度D = (帧高) × (帧宽) × (通道数) × (帧数)。当帧数很多时,维度会非常高,导致需要更精确的得分估计和更多的训练数据。

所以,虽然理论上可以直接应用,但为了高效地生成长视频,我们需要一些技术来降低处理维度。

潜在扩散与维度缩减

上一节我们提到了视频生成中的维度挑战。本节中我们来看看一种常用的解决方案:潜在扩散模型。

为了降低数据维度,我们可以使用自动编码器。首先,用一个编码器E将高维数据x映射到一个低维的潜在空间z = E(x),其中z的维度远小于x。同时,存在一个解码器D,可以(近似地)从z重建回x,即 D(E(x)) ≈ x。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_12.png

理想情况下,我们应该在低维的潜在空间z上运行扩散过程(即对z加噪和去噪)。这被称为潜在扩散。

但在一些实际实现(如早期的潜在扩散模型)中,噪声仍然被添加到原始数据x上,而不是潜在表示z上。从数学角度看,这导致过程不再是原始Wasserstein度量空间上的梯度流,而是在编码器诱导出的新度量空间上的梯度流。

这个差异意味着,在反向过程中,除了要估计得分函数,理论上还需要考虑一个由诱导度量带来的额外项(与局部曲率相关)。为了获得更高质量的生成结果,需要在训练目标中加入对这个项的近似。

Sora的架构:扩散Transformer

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_14.png

上一节我们讨论了潜在扩散的思想。本节中我们来看看OpenAI Sora模型采用的具体架构:扩散Transformer。

Sora的核心是一个基于Transformer架构的扩散模型,称为扩散Transformer。其工作流程如下:

  1. 编码:使用预训练的编码器将视频帧(或图像块)映射到低维潜在空间。

  2. 标记化:将潜在表示展平为一序列的标记,作为Transformer的输入。

  3. Transformer处理:序列标记通过多个Transformer块。每个块包含多头注意力层和前馈网络。

  4. 条件注入:在Transformer块中,通过交叉注意力等方式融入文本描述条件,指导视频生成。

  5. 预测:Transformer的输出被重新整形,用于预测噪声(得分函数)。

  6. 解码:去噪后的潜在表示通过解码器恢复为像素空间的视频。

以下是架构的简化表示:

[视频帧] -> [编码器] -> [潜在表示] -> [展平/标记化] -> [Transformer块 + 文本条件] -> [预测噪声] -> [去噪] -> [解码器] -> [生成视频]

关键技术:三维位置编码

上一节介绍了扩散Transformer的整体架构。本节中我们来看一个支持可变长度视频生成的关键技术:三维位置编码。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_16.png

视频数据具有三维结构:宽度(X)、高度(Y)和时间(T)。为了将视频块输入Transformer,我们需要告诉模型每个块在三维空间中的位置。

Sora采用了类似Google NaViT 模型的思想,使用分解式的三维位置编码。对于一个位于坐标 (x, y, t) 的视频块,其位置编码是三个独立的一维位置编码之和:

位置编码(x, y, t) = PE_x(x) + PE_y(y) + PE_t(t)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_18.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_20.png

这种方法的优势在于:

  • 支持原生分辨率:视频可以保持原有的宽高比和帧率,无需强制缩放到统一尺寸。

  • 支持可变长度:可以处理不同时长和不同尺寸的视频。

  • 灵活性高:为模型提供了明确的空间和时间结构信息。

总结与展望

本节课中我们一起学习了视频生成扩散模型的原理与实现。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_22.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_23.png

我们从扩散模型的数学基础出发,理解了其核心是Wasserstein梯度流,并通过随机微分方程和福克-普朗克方程描述了前向和反向过程。生成的关键在于使用神经网络进行得分匹配。

我们看到,该理论具有普适性,可直接应用于视频分布。为了应对高维挑战,采用了潜在扩散技术来降低维度。OpenAI的Sora模型基于扩散Transformer架构,并利用三维位置编码来处理可变长度的视频数据。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/af0747308152623b6530f4cffe43a1d7_25.png

当前,这类模型的主要限制在于巨大的计算成本,无论是训练还是推理。然而,由于其坚实的数学基础,性能主要遵循扩展定律:更多的数据和更大的模型将直接带来更好的结果。这为未来的发展提供了清晰的方向。

21:视频生成模型 🎬

在本节课中,我们将学习如何将图像生成中的扩散模型技术扩展到视频生成领域。我们将深入探讨其背后的数学原理,特别是瓦瑟斯坦梯度流,并了解OpenAI的Sora模型所采用的核心架构。

数学基础:扩散模型与瓦瑟斯坦梯度流 🔢

上一节我们介绍了扩散模型的基本概念,本节中我们来看看其背后的数学原理。扩散模型的核心在于其前向过程和反向过程。

从数学角度看,扩散过程实际上是在一个特殊的度量空间——瓦瑟斯坦空间中,沿着梯度流路径移动概率分布。我们可以将一个概率分布视为这个度量空间中的一个点。

瓦瑟斯坦距离 的公式定义如下:

$$

W_2(P, Q) = \inf_{\gamma \in \Gamma(P, Q)} \sqrt{\mathbb{E}_{(x, y) \sim \gamma} [|x - y|^2]}

$$

其中, Γ ( P , Q ) \Gamma(P, Q) Γ(P,Q) 是所有边缘分布分别为 P P P Q Q Q 的联合分布 γ \gamma γ 的集合。

扩散模型的前向过程,就是将原始数据分布 P 0 P_0 P0 沿着该空间中的最短路径(即梯度流)转化为高斯分布 P ∞ P_\infty P。反向过程则是这条路径的逆过程。

前向过程:随机微分方程视角 📈

理解了梯度流的视角后,我们来看看如何用具体的方程来描述这个过程。前向过程可以通过一个随机微分方程来刻画。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_1.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_5.png

假设 X t X_t Xt 是一个遵循分布 P t P_t Pt 的随机变量。那么,存在一个随机过程满足以下随机微分方程:

$$

dX_t = -X_t dt + \sqrt{2} dB_t

$$

其中 B t B_t Bt 是标准布朗运动。这个方程描述了如何通过“收缩”信号并添加噪声,将分布逐渐推向高斯分布。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_7.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_9.png

从采样角度看,在时间 t t t,变量 X t X_t Xt 的分布等同于以下过程的结果:

$$

X_t \stackrel{d}{=} e^{-t} X_0 + \sqrt{1 - e^{-2t}} Z

$$

其中 X 0 ∼ P 0 X_0 \sim P_0 X0P0 Z ∼ N ( 0 , I ) Z \sim \mathcal{N}(0, I) ZN(0,I)。这正是我们在扩散模型前向过程中看到的“不断加噪”操作。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/a1ee31da10bda008d2935bd850b36d20_11.png

反向过程与福克-普朗克方程 🔄

为了从噪声中生成数据,我们需要运行反向过程。这涉及到随机微分方程的反转,而福克-普朗克方程为此提供了数学工具。

定义一个新的过程 Y t = X T − t Y_t = X_{T-t} Yt=XTt,其中 T T T 是一个足够大的时间。那么 Y t Y_t Yt 从近似高斯分布( Y 0 ≈ X T Y_0 \approx X_T Y0XT)出发,最终应能恢复原始分布( Y T = X 0 Y_T = X_0 YT=X0)。

利用福克-普朗克方程,可以推导出 Y t Y_t Yt 满足的反向随机微分方程为:

$$

dY_t = [Y_t + 2 \nabla_{y} \log p_t(Y_t)] dt + \sqrt{2} dB_t

$$

其中 p t ( ⋅ ) p_t(\cdot) pt() X t X_t Xt 的概率密度函数。与正向方程相比,反向方程多出了一项 ∇ log ⁡ p t ( y ) \nabla \log p_t(y) logpt(y),即概率密度对数(得分函数)的梯度。

得分匹配与神经网络训练 🧠

反向过程的核心在于估计得分函数 ∇ log ⁡ p t ( x ) \nabla \log p_t(x) logpt(x)。在实践中,我们训练一个神经网络来近似这个函数。

训练目标是最小化得分匹配损失。一个关键结论是,对于上述特定的前向过程,训练神经网络根据含噪输入预测所添加的噪声,在数学上等价于学习得分函数。因此,扩散模型的训练可以归结为一个去噪任务。

这个过程是数学上严格的,没有近似。只要我们能完美地学习到得分函数,运行反向SDE就能精确地采样出目标数据分布。这解释了扩散模型为何是一种通用的生成方法,可应用于图像、视频、音频等任何数据分布。

扩展到视频生成 🎥

现在,我们将上述原理应用到视频生成。视频本质上是一个图像序列的概率分布。

我们可以将一段视频视为一个高维随机变量 X ∈ R H × W × C × T X \in \mathbb{R}^{H \times W \times C \times T} XRH×W×C×T,其中 T T T 是帧数。扩散模型的理论保证,只要我们能学习该视频分布的得分函数,就能生成视频。

然而,维度 D = H × W × C × T D = H \times W \times C \times T D=H×W×C×T 可能非常大。扩散模型的收敛速度和样本复杂度通常与维度 D D D 成正比,这意味着生成长视频需要极高的计算成本和数据量。

潜在扩散与维度约减 🗜️

为了应对高维度的挑战,常用的技术是潜在扩散。其核心思想是先用一个编码器将高维数据(如图像/视频帧)压缩到一个低维潜在空间,然后在潜在空间中进行扩散过程。

设编码器为 E E E,解码器为 D D D,满足 D ( E ( x ) ) ≈ x D(E(x)) \approx x D(E(x))x。理想情况下,我们应在潜在变量 z = E ( x ) z = E(x) z=E(x) 上运行扩散过程。但许多实际实现(如Stable Diffusion的早期版本)仍然在原始像素空间加噪,这相当于在由编码器诱导的度量空间中进行梯度流,而非标准的瓦瑟斯坦梯度流。

这引入了额外的几何项(与诱导度量相关)。为了获得高质量的生成结果,在训练目标中需要考虑这个项,例如除了预测噪声外,还可能预测一个与局部协方差相关的项 Σ \Sigma Σ。最新的潜在扩散模型已开始纳入这些修正。

Sora的架构:扩散Transformer (DiT) 🏗️

OpenAI的Sora模型基于扩散Transformer架构。其核心思想是将扩散模型中的U-Net主干替换为Transformer。

以下是该架构的关键步骤:

  1. 输入处理:视频经过编码器被映射到潜在空间。潜在表示被展平为一序列的令牌(tokens),作为Transformer的输入。

  2. 位置编码:Sora采用了类似Google“原生分辨率ViT”的技术,使用3D位置编码。一个令牌的位置编码是其空间坐标 ( x , y ) (x, y) (x,y) 和时间坐标 t t t 的位置编码之和: P E = P E x ( x ) + P E y ( y ) + P E t ( t ) PE = PE_x(x) + PE_y(y) + PE_t(t) PE=PEx(x)+PEy(y)+PEt(t)。这支持可变分辨率、可变长度的视频输入,无需将视频裁剪成固定尺寸。

  3. Transformer块:序列令牌通过一系列Transformer块。这些块集成了多头自注意力机制和前馈网络。

  4. 条件注入:生成的条件信息(如文本描述)被注入到每个Transformer块中,通常通过交叉注意力或直接添加到隐藏状态来实现。

  5. 输出与训练:Transformer的输出被重新变换回潜在空间的形状,用于预测噪声(得分函数)以及可能的方差项 Σ \Sigma Σ。训练目标是最小化去噪得分匹配损失。

总结 📝

本节课我们一起学习了视频生成扩散模型的数学基础与核心架构。

我们首先回顾了扩散模型的本质,即瓦瑟斯坦空间中的梯度流,并通过随机微分方程和福克-普朗克方程理解了其严格的反向过程。我们认识到,训练神经网络进行去噪等价于学习得分函数,这使得扩散模型成为适用于任何数据分布的通用生成框架。

接着,我们将该框架应用于视频生成,指出了高维度带来的挑战,并介绍了潜在扩散技术作为解决方案。最后,我们剖析了OpenAI Sora模型所使用的扩散Transformer架构,其关键创新在于使用了支持可变长视频的3D位置编码,并将Transformer的强大表达能力与扩散模型的严格生成理论相结合。

视频生成的突破主要得益于扩散模型的数学保证和计算规模的扩大,而非魔法般的算法创新。随着模型规模和数据量的持续增长,遵循缩放定律,我们有望看到更加强大和通用的生成模型。

23:遗留主题 - 流水线并行与μP初始化

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_1.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_5.png

在本节课中,我们将学习两个之前课程中未深入讨论的重要遗留主题:用于分布式优化的流水线引擎,以及一种用于初始化Transformer网络的新方法——μP。这些技术是当前训练超大规模模型(如万亿参数模型)的关键。

流水线并行引擎

上一节我们介绍了专家并行,它主要解决了MLP层的内存问题。本节中我们来看看另一种关键的并行化技术——流水线并行,它能够处理模型深度(层数)的扩展问题。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_7.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_8.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_10.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_12.png

什么是流水线并行?

流水线并行的核心思想非常简单:将神经网络按层切分,并将不同的层放置在不同的GPU上。对于一个由多个Transformer块顺序堆叠的网络,最自然的并行方式就是将每一层放在一个独立的GPU节点上。

例如,如果你的网络定义为一个nn.Sequential模块:

net = nn.Sequential(layer1, layer2, layer3, ..., layerL)

那么流水线引擎在精神上会自动将layer1放在GPU1,layer2放在GPU2,依此类推。这种方法理论上可以扩展到无限深度,只要你有足够多的GPU。

朴素实现的挑战

然而,这种按层切分的方式带来了一个核心问题:计算是顺序的,而非并行的。GPU1必须先完成第1层的计算,才能将结果传给GPU2进行第2层的计算。这导致了大量的GPU空闲时间,计算效率极低。如果有100层,计算时间可能比单GPU单层模型慢100倍。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_14.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_15.png

高效的流水线引擎

如何解决GPU空闲问题,让流水线引擎更高效?核心思路是利用空闲时间。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_17.png

观察计算图,当一个GPU完成当前数据(例如x1)在当前层的计算后,在等待下一层GPU接收并开始计算时,它处于空闲状态。高效的流水线引擎会利用这个空闲时间,开始处理下一个数据(例如x2)在当前层的计算。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_19.png

以下是实现高效流水线的关键机制:

  1. 状态卸载与加载:在完成x1的前向计算后,GPU立即将其激活状态(用于后续反向传播)卸载(offload)到CPU内存。由于卸载到CPU是非阻塞操作,GPU可以立刻开始计算x2的前向传播。

  2. 调度与重叠:通过精心调度,让不同数据样本(x1x2x3…)在不同层上的计算相互重叠,形成一种“流水线”效果,从而填满GPU的空闲时间。

  3. 周期性同步:在实践中,为了鲁棒性和内存管理,流水线引擎通常会设置一个“微批次”(micro-batch)边界。在完成一定数量的重叠计算后,会进行一次同步,清空中间状态,然后开始新的计算周期。这比完全无休止的重叠更稳定,也更容易处理故障恢复。

流水线并行的优势与局限

流水线并行非常适合Transformer这类结构,因为每一层的计算量和张量形状大致相同,易于负载均衡。它与专家并行、数据并行结合,构成了训练超大模型的“3D并行”范式。

  • 流水线并行:跨层切分,解决模型深度问题。

  • 数据并行:跨数据批次切分,增加吞吐量。

  • 模型并行(如专家并行):跨层内组件切分,解决层宽度问题。

对于注意力层,如果单层仍然无法放入内存,还可以使用序列并行等技术。

重要提示:流水线并行的实现(如Hugging Face的Pipeline Optimizer)对用户几乎是透明的。你只需要将模型定义为顺序模块,优化器会自动处理层间通信和调度,无需像张量并行那样手动指定复杂的张量切分。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_21.png


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_23.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_24.png

μP 初始化

接下来,我们探讨另一个主题:μP(Maximal Update Parametrization)。这是一种有争议但被OpenAI成功用于训练MOE模型的技术。它主要调整了Transformer网络的初始化和学习率调度策略。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_26.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_28.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_29.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_31.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_33.png

为什么需要μP?

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_35.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_36.png

要理解μP,我们先思考一个根本问题:当模型尺寸(例如宽度N)趋向无穷大时,我们常用的优化器(如Adam)和固定学习率还能稳定工作吗?

考虑一个简单的单层线性网络:y = Wx + b。假设使用标准初始化:Wb的元素服从N(0, 1/N)

  • 初始化输出:在N很大时,输出y的尺度是O(1),这是合理的。

  • 一次更新后:Adam等优化器的参数更新量级通常是O(1)(经过梯度归一化)。对于偏置b,更新ΔbO(1)。那么更新后的输出包含一项Δb * x。由于Δb与梯度相关,而梯度又与权重W相关,这项的期望尺度可能是O(√N)。这意味着仅一次更新,网络输出就可能爆炸(或变得极大),导致训练不稳定。

如果为了避免爆炸而将学习率设为O(1/√N),那么更新权重所需的有效步数将变为O(√N)。当N很大时,这意味着需要极多的步骤才能有意义地更新网络,训练效率极低。

μP的设计目标

μP旨在寻找一个参数化(初始化+学习率缩放)方案,使得在模型宽度N → ∞的极限下,满足两个核心条件:

  1. 网络可快速更新:每个参数都能在常数步数内被有效更新。

  2. 更新过程稳定:单次参数更新不会导致网络输出发生剧烈(O(√N))变化。

这需要在“更新幅度不能太小”(目标1)和“更新幅度不能太大”(目标2)之间找到一个精妙的平衡点。

μP的解决方案

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_38.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_40.png

μP通过层间差异化的初始化方差和学习率缩放来实现这一目标。其推导基于理论分析,确保在前向传播和反向传播过程中,各层激活值和梯度的尺度保持一致且可控。

一个简化的示意是,对于深度为L的Transformer,μP可能会规定:

  • 输入嵌入层:保持标准初始化(如N(0, 1))和常数学习率。

  • 中间层:初始化方差需要按1/N1/(N√L)等因子缩放。

  • 输出层:初始化方差需要按1/N^2等因子缩放,并且其学习率可能需要按1/N缩放。

这些缩放因子确保了即使在超宽网络中,前向信号、反向梯度以及参数更新量的尺度都是O(1),从而满足上述两个设计目标。

μP的意义与争议

μP提供了一种原则性的方法来初始化和大规模训练Transformer,尤其是宽度极大的模型或MOE模型。OpenAI的成功经验表明,这种精细的缩放对训练稳定性至关重要。

然而,μP在学术界和工业界其他团队中并未被广泛验证为“唯一有效”的方法。一些替代方案,例如使用极大的权重衰减(weight decay)配合层归一化(LayerNorm),也能起到类似的稳定作用。但μP的价值在于它从一个理论极限(N → ∞)出发,给出了一个系统性的超参数设置框架。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_42.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_43.png

总结

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_45.png

本节课中我们一起学习了两个支撑当今超大规模生成式AI模型训练的关键技术:

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_47.png

  1. 流水线并行引擎:通过将模型按层切分到不同GPU,并利用巧妙的调度重叠不同数据样本的计算,极大地提高了深度模型的训练效率,是实现模型深度扩展的核心手段。

  2. μP初始化:通过理论推导出一套针对超宽网络的初始化与学习率缩放方案,旨在保证模型在宽度趋向无穷时,仍能保持快速且稳定的训练动态。这是OpenAI成功训练MOE等巨型模型的重要“秘方”之一。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_49.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_51.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_52.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e559adab3bfbcaf84a3abca30946f81b_54.png

理解这些底层优化技术,有助于我们更好地把握大模型训练的工程挑战和前沿方向。

24:遗留主题 - 流水线并行与μP初始化 🧠

在本节课中,我们将学习两个在前几讲中因时间关系未深入讨论的重要主题:用于分布式优化的流水线引擎,以及一种旨在提升大规模Transformer网络训练稳定性的初始化与学习率调度方法——μP。这些技术是当前训练超大规模语言模型(如万亿参数模型)的核心。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e58729540836aeb03b3de82193becdec_1.png


流水线并行引擎 🚂

上一节我们介绍了专家并行等模型并行技术。本节中,我们来看看另一种关键的并行范式——流水线并行。当模型过大,无法放入单个GPU内存时,除了数据并行和模型并行,我们还需要沿模型的深度方向进行切分。

为什么需要流水线并行?

以H100 GPU(约80GB内存)为例,其最大能容纳的模型大约是70亿参数、上下文长度8K的稠密模型。而当前业界关注的模型规模通常是这个的10倍甚至100倍以上。因此,必须采用更高级的并行技术。

实践中,为了扩展到极大模型,主要依赖两种技术组合:

  1. 将稠密模型转换为混合专家模型,以启用高效的专家并行

  2. 在上述基础上,应用流水线并行。这一步尤为关键,其实现效率直接决定了整体优化速度。

以下是对这两种并行方式的高层次对比:

# 专家并行:每个专家放置在不同的GPU上
expert_1 = nn.Linear(D, 2*D).to('gpu:0')
expert_2 = nn.Linear(D, 2*D).to('gpu:1')
# ... 计算时,根据路由将token发送到对应的专家GPU

# 流水线并行:将模型的连续层放置在不同的GPU上
model = nn.Sequential(layer1.to('gpu:0'), layer2.to('gpu:1'), layer3.to('gpu:2'))

流水线并行的核心思想与挑战

流水线并行的思想非常直观:将一个由L层组成的神经网络(如Transformer)按层切分,分别放入L个不同的GPU中。每个GPU仅负责其对应层的前向和反向计算。

然而,最朴素的实现方式效率极低。因为计算是严格顺序的:必须等第1层计算完,才能开始第2层的计算,以此类推。这导致在任一时刻,大部分GPU都处于空闲状态,计算时间将是单GPU的L倍,无法有效利用多GPU的算力。

高效的流水线调度

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e58729540836aeb03b3de82193becdec_3.png

为了解决GPU空闲问题,核心思路是让GPU在等待当前数据流中上一层的计算结果时,提前开始处理下一批数据。这需要通过巧妙的调度和内存管理来实现。

以下是实现高效流水线的关键机制:

  1. 计算与通信重叠:当一层完成其前向计算后,立即将输出的激活值卸载到CPU内存(这是一个非阻塞操作)。与此同时,该GPU可以立即开始处理下一批数据的前向计算。

  2. 调度填充:通过合理安排不同批次数据在不同层上的计算顺序,可以填充大部分空闲时间,形成类似“流水线”的高效运作模式。理想情况下,调度应使得所有GPU持续处于工作状态。

一个简化的流水线调度时间线示意图如下(其中F代表前向,B代表反向,数字代表批次序号):

时间 -> 
GPU0: F1 -> B1 -> F4 -> B4 ...
GPU1: 空闲 -> F1 -> B1 -> F4 ...
GPU2: 空闲 -> 空闲 -> F1 -> B1 ...

(注:实际调度更复杂,需考虑内存和依赖关系)

实践中的考量

在实际系统中(如OpenAI使用的引擎),流水线并行的实现还涉及更多细节:

  • 微批次:将一个大批次拆分成更小的微批次,以更细粒度地填充流水线。

  • 梯度累积:为了保持有效的批次大小,需要在多个微批次上累积梯度后再更新权重。

  • 周期同步点:定期设置同步点以清空流水线,这有助于内存管理和提高系统容错性。

  • 与其它并行方式结合:现代大规模训练通常采用 3D并行:流水线并行(层间)+ 数据并行(数据间)+ 专家并行/张量并行(层内)。

流水线并行的优势在于,对于像Transformer这样各层计算量均匀的模型,它几乎可以自动实现(将模型定义为nn.Sequential即可),无需像张量并行那样手动重写层内计算逻辑。


μP初始化:面向超大规模模型的稳定训练 🔬

上一节我们探讨了如何通过并行化来容纳大模型。本节中,我们来看看另一个关键问题:当模型参数数量N趋向于极大时,如何保证优化过程依然稳定高效?这就是μP方法要解决的核心问题。

标准初始化的问题

考虑一个简单的两层线性网络(忽略偏置和激活函数):

输出 = W2 * (W1 * 输入)

假设输入维度为D,隐藏层维度为N,输出维度为1。

标准初始化(如Kaiming初始化)将W1W2的元素初始化为均值为0、方差分别为1/D1/N的高斯分布。

当隐藏层维度N极大时,使用固定学习率(如0.1)的Adam优化器进行一步更新,可能会引发两个问题:

  1. 输出爆炸:更新步长相对于网络输出可能过大,导致一步更新后网络输出值急剧增大甚至溢出。

  2. 更新无效:为避免爆炸而过度调低学习率,又会导致参数更新幅度过小,需要极多步迭代才能对网络产生有意义的影响。

这两种情况都会导致训练不稳定或效率极低。

μP的设计原则

μP旨在寻找一个初始化与学习率调度的配置,使得当N → ∞时,满足以下两个条件:

  1. 网络输出稳定:单步更新不会导致网络输出发生剧烈(O(√N)量级)变化。

  2. 参数有效更新:经过常数步(而非O(N)步)的更新后,网络参数能够发生显著变化。

这需要在初始化尺度层间学习率之间进行精细的协调。

μP的解决方案(示意)

μP提出了一套系统性的缩放规则。其核心思想是:不同层参数的初始化方差和学习率应随网络宽度N进行不同的缩放。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/e58729540836aeb03b3de82193becdec_5.png

以一个简化版本为例(忽略层归一化等细节):

  • 输入层(W1):保持标准初始化(方差~1/D)和常数学习率。这使得嵌入层能在常数步内被有效更新。

  • 输出层(W2):为了抑制输出爆炸,需要将其初始化方差缩小为~1/N^2(而非标准的1/N)。同时,为了保证该层参数也能在常数步内被更新,其对应的学习率需要放大为O(N)

这样,虽然W2的初始值很小(O(1/N)),但乘以放大的学习率(O(N))后,单步更新量级为O(1),从而能在常数步内显著改变该层参数。同时,由于初始化很小,即使更新量级为O(1),也不会使最终输出爆炸。

意义与争议

μP从理论上为超宽网络的稳定训练提供了一种原则性方法。据报道,它是OpenAI成功训练大规模MoE模型的关键技术之一。

然而,该方法也存在争议。在OpenAI之外,很少有团队能完全复现其宣称的效果。这可能是因为实际系统还包含了未公开的细节调整,或者依赖特定的硬件和软件栈。其他公司可能采用不同的启发式方法,例如使用极大的权重衰减配合层归一化,也能达到稳定训练的目的。

尽管如此,理解μP背后的思想——即在网络宽度缩放时,协调初始化与学习率以保持优化动态的稳定性——对于从事大规模模型训练的研究者和工程师至关重要。


总结 📝

本节课中我们一起学习了两个支撑当今超大规模生成式AI模型训练的关键技术:

  1. 流水线并行引擎:通过将模型按层切分到多个GPU,并利用智能调度重叠计算与通信,极大提高了内存利用率和训练效率,是实现模型深度扩展的核心手段。

  2. μP初始化:通过精心设计参数初始化方差和层间学习率的缩放规则,旨在保证当模型宽度极大时,优化过程依然保持稳定和高效。这体现了对神经网络训练动态的深刻理解。

这些技术代表了分布式AI系统与优化理论前沿的结合,是构建下一代更大、更强AI模型的基础设施的重要组成部分。

25:语音识别模型 🎤

在本节课中,我们将学习语音识别模型。虽然语音识别在技术上并非纯粹的生成式模型,但它使用了与生成式模型相似的架构,并且在生成式机器学习中扮演着重要角色。我们将了解如何将语音信号转换为文本,以及两种主流的模型架构。


概述:语音数据的重要性

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_1.png

语音数据正变得越来越重要,因为它是训练语言模型的高质量数据源。例如,播客、辩论甚至总统演讲的文本记录,都是极其重要且高质量的数据。这些数据中没有广告,信息密度高。仅YouTube视频的字幕就可能包含超过10万亿个高质量的词元。然而,我们需要将语音转换为自然语言文本才能利用它们。本节课,我们将学习如何实现这一点。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_3.png

声音的数学表示:从波形到向量

上一节我们提到了语音数据的重要性,本节中我们来看看如何将连续的语音信号转换为计算机可以处理的数字形式。

声音在物理学上由两个基本特征描述:

  1. 振幅:决定声音的响度。

  2. 频率:决定声音的音高。频率越高,音调越高;频率越低,音调越低沉。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_5.png

任何复杂的声音波形都可以分解为一系列具有特定振幅和频率的简单正弦波的组合。这种分解方法称为傅里叶变换

一个在区间 [0, 2π] 上的连续函数 f(x) 可以表示为:

f(x) = Σ (α_j * e^(i * θ_j * x))

其中,α_j 是复数系数(代表振幅和相位),θ_j 是频率。

通过傅里叶变换,我们可以将时域中的声音信号转换到频域,用一个(理论上无限维的)系数向量 [α_0, α_1, α_2, ...] 来表示。在实际应用中,由于系数会衰减,我们可以找到一个有限的截断点 N,用前 N 个系数 [α_0, α_1, ..., α_N] 来足够精确地近似原始声音。

因此,对于一段较长的音频,我们可以先将其切分为多个时间窗口,然后将每个窗口内的声音信号转换为一个频域系数向量。最终,整个音频就被表示为一个向量序列,这与自然语言处理中词元序列的形式非常相似。


梅尔频谱:更优的频域表示

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_7.png

上一节我们介绍了通过傅里叶系数向量表示声音的方法,但这种方法存在一个问题:它生成的向量维度过高,且不符合人耳的感知特性。

原始的频谱系数在高频区域非常集中,振幅较大。然而,人耳对高频变化的敏感度低于中低频。例如,人耳很难区分600Hz和800Hz的声音,但对1000Hz附近的变化非常敏感。因此,我们需要一种更符合人耳感知特性的表示方法。

直接将高维向量截断或均匀分组(池化)会丢失重要信息。一个更聪明的方法是非均匀分组,即根据一个特定的函数将频率系数分组到不同的“桶”中,并对每个桶内的系数进行平均。这个函数近似于指数函数,被称为梅尔尺度

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_9.png

梅尔频谱的计算过程如下:

  1. 计算音频窗口的傅里叶系数。

  2. 根据梅尔尺度函数,将频率系数分组到多个频带(例如80个)中。

  3. 对每个频带内的系数进行加权平均,得到该频带的能量值。

这样,我们就将一个高维的傅里叶系数向量,压缩成了一个低维的梅尔频谱向量。每个向量元素代表一个特定频带在某个时间窗口内的平均能量。整个音频因此被表示为一个二维矩阵:一个维度是时间窗口序列,另一个维度是梅尔频带。

这种表示不仅是维度上的压缩,更是对声音信息的一种感知优化,为后续的模型处理奠定了更好的基础。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_11.png

模型架构一:Wave2Vec(无监督学习)

现在我们已经将音频转换为向量序列,接下来看看如何训练模型来理解这些向量。第一种主流架构是 Wave2Vec,它主要采用无监督学习。

Wave2Vec 的训练数据大部分是未标注的纯音频。其核心思想是学习音频信号的良好表示,类似于 BERT 在文本领域所做的工作。

以下是 Wave2Vec 的无监督训练流程:

  1. 输入:音频经过梅尔频谱处理后的向量序列 Q = [q_1, q_2, ..., q_T]

  2. 掩码:随机掩码掉其中约15%的向量(类似于 BERT 的掩码语言模型)。

  3. 编码:将掩码后的序列输入一个 Transformer 编码器。

  4. 输出:Transformer 输出一个上下文向量序列 C = [c_1, c_2, ..., c_T]

  5. 对比损失:训练目标是让被掩码位置 t 的输出向量 c_t 尽可能接近其真实的向量 q_t,同时远离其他所有时间步的向量 q_{t'} (t’ ≠ t)。

其对比损失函数可以简化为:

Loss = -log[ exp(sim(c_t, q_t)) / Σ_{t'} exp(sim(c_t, q_{t'})) ]

其中 sim 是相似度函数,如余弦相似度。

为什么使用对比损失而非简单的预测损失?

因为相邻时间窗口的音频向量 q_tq_{t-1} 通常非常相似。简单的回归损失会让模型倾向于复制前一个向量,而无法学习到有区分度的特征。对比损失能迫使模型关注于每个时间窗口的独特特征,并忽略持续的背景噪音,这更接近人脑处理声音的方式。

训练完成后,Wave2Vec 得到了音频信号的高质量上下文表示。要用于语音识别等下游任务,只需在其顶部添加一个简单的线性分类层,并用少量有标注数据进行微调即可。这种架构特别擅长声音分类、说话人分割等任务。


模型架构二:Whisper(有监督编码器-解码器)

上一节我们学习了无监督的 Wave2Vec 模型,本节我们来看另一种主流的架构:Whisper。这是一个完全基于有监督学习的编码器-解码器模型,专门用于语音到文本的转录。

Whisper 模型的训练数据是成对的音频和文本字幕。但这里有一个关键挑战:字幕通常是整个音频段的概括,没有与音频时间轴精确对齐。我们不知道哪句文本对应哪段音频。

Whisper 巧妙地利用了解码器的自回归生成能力和交叉注意力机制来解决这个问题。

以下是 Whisper 的工作流程:

  1. 编码:音频被转换为梅尔频谱向量序列,然后送入一个编码器(可能包含CNN和Transformer)进行编码,得到音频的上下文表示。

  2. 解码:使用一个 Transformer 解码器来生成文本词元序列。

  3. 交叉注意力:解码器的核心创新在于其注意力机制。在预测下一个文本词元时,解码器不仅会关注之前已生成的所有文本词元(自注意力),还会通过交叉注意力层关注整个音频编码序列。

  4. 训练目标:模型的学习目标是最简单的下一个词元预测。给定一段音频和对应的文本字幕,模型需要根据音频信息和已生成的文本,预测出下一个正确的文本词元。

这类似于图像生成的扩散模型或DALL-E,只不过这里是条件于音频来生成文本。由于解码器在生成每个词时都能“听到”整个音频,理论上它应该能生成准确的转录文本。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_13.png

Whisper 模型因其出色的开箱即用转录能力和多语言支持而广受欢迎。它代表了当前大规模有监督训练在语音识别领域的成功应用。


总结与展望

本节课我们一起学习了语音识别模型的基础知识。

我们首先了解了语音作为高质量数据源的重要性。接着,探讨了如何将声音从连续的波形,通过傅里叶变换和梅尔频谱处理,转换为适合神经网络处理的向量序列。

然后,我们深入分析了两种主流的模型架构:

  • Wave2Vec:采用无监督对比学习,旨在学习音频信号本身的优良表示,适用于需要音频理解的分类任务。

  • Whisper:采用有监督的编码器-解码器架构,利用交叉注意力机制,直接条件于音频生成文本,是目前主流的高质量语音转录方案。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/371b02718c6791c24d1e433133c7943d_15.png

目前,语音识别模型仍有改进空间,例如更好地处理说话人分离、背景音过滤等。未来的方向可能是结合 Wave2Vec 的无监督学习能力与 Whisper 的强转录能力,以更少的有标注数据获得更强大、更鲁棒的模型。这是一个非常活跃且重要的研究领域。

26:语音识别模型

在本节课中,我们将学习语音识别模型。虽然语音识别在技术上并非纯粹的生成式模型,但它使用了与生成式模型相似的架构,并且在机器学习领域非常重要。我们将了解如何将语音信号转换为自然语言文本,并探讨两种关键的模型架构。

概述:语音数据的重要性

语音数据正变得越来越重要,因为它是训练语言模型的高质量数据来源。例如,播客、辩论甚至总统演讲的文本记录,都是极其重要且高质量的数据。这些数据不包含广告,信息密度高。粗略估计,仅YouTube转录文本就可能提供超过10万亿的高质量词元。然而,我们需要将语音转换为自然语言才能利用这些数据。本节课,我们将学习如何实现这一点。

从声音到向量:信号处理基础

上一节我们介绍了语音数据的重要性,本节中我们来看看如何将原始的语音信号转换为计算机可以处理的格式。

回想一下物理知识,声音由两个基本特征描述:

  • 振幅:决定声音的响度。

  • 频率:决定声音的音高。频率越高,音调越高;频率越低,声音越低沉。

任何复杂的声音波形都可以分解为一系列具有特定振幅和频率的简单正弦波的组合。这种分解方法称为傅里叶变换

通过傅里叶变换,我们可以将时域中的声音信号转换到频域。在频域中,一个声音片段(例如一个时间窗口内的声音)可以表示为一个向量,向量的每个维度对应一个特定频率分量的振幅。

然而,原始傅里叶系数向量维度可能非常高(例如对应高达100kHz的频率)。直接使用这样的高维向量作为模型输入效率低下。此外,人耳对不同频率的敏感度不同,对中频区间的变化更敏感。

因此,我们需要一种更智能的降维方法,而不是简单截断高频系数。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/81128a1ec29d029fd685c10365c031ae_1.png

梅尔频谱:更优的频域表示

上一节我们提到了直接使用傅里叶系数的局限性,本节中我们来看看如何更有效地表示频域信息。

解决方案是使用梅尔频谱。其核心思想是:根据一个类似指数函数的“梅尔尺度”将频率分组到不同的“频带”或“桶”中。

在梅尔尺度下:

  • 低频区域(如0-100Hz)被划分为较少的频带,分辨率较粗。

  • 高频区域(如8000-10000Hz)被划分为较多的频带,分辨率较细。

然后,我们将同一个频带内的多个傅里叶系数振幅进行平均(或采取其他聚合方式),用这个平均值作为该频带的代表值。这样,我们就将一个高维的傅里叶系数向量压缩成了一个低维的梅尔频谱向量。

这个过程可以总结为以下步骤:

  1. 对音频时间窗口进行傅里叶变换,得到频域系数。

  2. 根据梅尔尺度将频率分组到多个频带。

  3. 聚合(如平均)每个频带内的系数值,形成最终的梅尔频谱向量。

最终,一个完整的音频文件被处理成一个向量序列:[向量_窗口1, 向量_窗口2, ..., 向量_窗口T]。这与自然语言处理中的词元序列非常相似,因此可以输入给Transformer等序列模型。

语音识别模型架构(一):Wave2Vec 2.0

现在我们已经将语音转换为向量序列,接下来看看如何训练模型来理解这些向量。首先介绍一种以无监督学习为主的架构:Wave2Vec 2.0

Wave2Vec 2.0 的训练主要使用大量无标注的纯语音数据。其灵感来源于BERT的掩码语言模型(MLM)目标,但针对语音数据的特点进行了调整。

以下是其无监督训练的核心流程:

  1. 输入:经过梅尔频谱编码的语音向量序列 Q = [q1, q2, ..., qT]

  2. 掩码:随机掩码(例如遮盖15%)序列中的部分向量,用特殊的[MASK]向量替换。

  3. 编码:将掩码后的序列输入一个Transformer编码器。

  4. 输出:Transformer输出一个上下文向量序列 C = [c1, c2, ..., cT],其中每个 ct 对应输入位置 t 的编码。

  5. 对比损失:对于每个被掩码的位置 t,训练目标是:

    • 使输出 ct 与真实的、被掩码的输入向量 qt 尽可能接近。

    • 同时,使 ct 与序列中所有其他位置的向量 qt't' ≠ t)尽可能远离。

这被称为对比损失。为什么在语音中要使用对比损失,而不是像BERT那样直接预测被掩码的词元呢?

原因在于语音数据的连续性。相邻时间窗口的语音向量 qtqt-1 通常非常相似。如果使用简单的预测损失,模型可能学会简单地复制前一个向量,而无法捕捉细微的、有意义的改变(如音素变化)。对比损失迫使模型学习每个时间窗口的独特表征,从而更好地区分不同的声音片段,并忽略持续的背景噪音。

训练完成后,Wave2Vec 2.0 模型获得了一个强大的语音特征编码器。要用于语音识别(转写文字),只需要在编码器的输出上添加一个简单的线性分类层,并用少量有标注的(语音,文本)配对数据进行微调即可。这种架构特别适合语音分类任务,如说话人分割、语音活动检测等。

语音识别模型架构(二):Whisper

上一节我们介绍了无监督的Wave2Vec模型,本节中我们来看看另一种更近期、专注于有监督语音识别的架构:Whisper

Whisper 是一个基于Transformer的编码器-解码器模型,主要用于有监督的语音转写任务。其训练数据是大量的(音频,转录文本)对。注意,这些数据通常是非对齐的,即我们只知道一段音频的整体转录文本,但不知道文本中每个词具体对应音频的哪一部分。

Whisper 的核心思想是将语音识别构建为一个条件生成任务:给定音频输入,生成对应的转录文本。这类似于图像生成模型根据文本描述生成图像。

以下是Whisper模型的工作流程:

  1. 编码:音频信号通过一个编码器(包含卷积层和Transformer层)被处理成一个特征向量序列。这类似于之前得到的梅尔频谱序列的进一步抽象。

  2. 解码:一个Transformer解码器负责自回归地生成文本词元(转录结果)。

  3. 交叉注意力:解码器的关键机制是交叉注意力。在生成每一个文本词元时,解码器不仅会关注之前已生成的所有文本词元(自注意力),还会通过交叉注意力机制去“聆听”或“关注”编码器输出的整个音频特征序列。

这意味着,在生成“apple”这个词时,模型不仅考虑了上文“I eat an”,还同时考虑了整个音频上下文的信息,从而能更准确地预测当前词。

通过在大规模有监督数据上训练这个编码器-解码器架构(使用标准的自回归语言建模损失),Whisper学会了将音频内容准确地翻译成文本。它无需显式的对齐信息,而是通过注意力机制隐式地学习音频与文本之间的对应关系。

总结与展望

本节课中我们一起学习了语音识别模型的基础知识和两种主流架构。

我们首先了解了语音数据作为高质量训练语料的重要性。接着,学习了如何通过傅里叶变换和梅尔频谱将连续的语音信号转换为离散的向量序列,为神经网络处理做好准备。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/81128a1ec29d029fd685c10365c031ae_3.png

然后,我们深入探讨了两种模型:

  1. Wave2Vec 2.0:一种基于对比学习的无监督/自监督模型,擅长学习语音的通用表征,适用于多种语音任务,经过微调后可进行语音识别。

  2. Whisper:一种基于编码器-解码器架构的有监督模型,直接学习从音频到文本的端到端映射,在语音转写任务上表现出色。

目前,语音识别模型仍有改进空间,其质量尚不及最先进的语言模型。一个有趣的研究方向是结合Wave2Vec的无监督学习能力和Whisper的强大生成能力,以利用海量的无标注语音数据,同时提升有监督任务的性能。这将是获得更多高质量训练数据、推动生成式AI发展的重要一步。

27:CUDA编程入门教程 🚀

在本节课中,我们将学习CUDA编程的基础知识。CUDA编程对于希望从零开始训练大型语言模型至关重要,因为它能让我们直接控制GPU的计算和内存访问,从而实现远超现有库(如PyTorch)的性能优化。我们将了解为什么大公司需要编写自己的CUDA内核,并探索其背后的核心架构和编程模式。

概述

CUDA编程是使用C/C++在NVIDIA GPU上进行计算的过程。理解CUDA对于优化生成式AI模型(如大型语言模型)的训练至关重要。标准深度学习框架(如PyTorch)虽然方便,但其通用性设计使其在极致性能优化上存在瓶颈。为了处理万亿参数级别的模型并应对GPU固有的硬件错误,顶尖的AI公司(如OpenAI)需要编写高度定制化的CUDA内核,以控制内存访问、实现计算冗余校验,并融合操作以减少数据移动。

GPU架构与内存层次结构

上一节我们介绍了学习CUDA编程的必要性,本节中我们来看看CUDA编程所基于的GPU架构。

CUDA编程的核心在于理解GPU的内存层次结构,这直接决定了代码的性能。GPU的计算单元组织如下:

  • 线程:最基本的执行单元,类似于CPU中的寄存器。

  • 线程块:一组线程的集合,它们可以协作并共享一块高速的片上内存,称为共享内存

  • 网格:由多个线程块组成,负责执行一个完整的CUDA内核。

内存访问速度是关键瓶颈。以下是主要的内存类型:

  • 全局内存:即GPU的显存,容量大但访问速度慢。

  • 共享内存:每个线程块独有的小块高速内存,访问速度极快。

CUDA编程的主要目标就是尽可能多地将数据保留在共享内存中进行计算,从而减少对缓慢的全局内存的访问。

CUDA编程基础

理解了架构后,我们来看看如何编写一个基本的CUDA程序。

一个典型的CUDA程序包含两部分:

  1. 内核函数:在GPU上每个线程中执行的函数。使用 __global__ 关键字声明。

  2. 主机函数:在CPU上运行的函数,负责配置并启动内核。

以下是一个向量加法的内核函数示例:

__global__ void vectorAdd(float *A, float *B, float *C, int n) {
    int i = threadIdx.x;
    if (i < n) {
        C[i] = A[i] + B[i];
    }
}

这个内核假设只使用一个线程块。threadIdx.x 是CUDA内置变量,表示当前线程在线程块内的索引。

主机函数调用内核的语法如下:

// 定义执行配置:使用1个线程块,该块包含n个线程
dim3 threadsPerBlock(n);
dim3 blocksPerGrid(1);
// 启动内核
vectorAdd<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, n);

<<<blocksPerGrid, threadsPerBlock>>> 语法指定了网格和线程块的维度,告诉GPU如何组织线程来执行这个内核。

并行化与线程索引

上一节我们看到了一个简单的单线程块内核,本节中我们来看看如何利用所有线程实现真正的并行计算。

为了让所有线程协作处理整个向量,我们需要修改内核,使每个线程处理不同的数据片段。这通过结合 threadIdx.xblockDim.x(线程块内的线程总数)来实现。

以下是优化后的并行向量加法内核:

__global__ void parallelVectorAdd(float *A, float *B, float *C, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x; // 计算全局线程ID
    int stride = blockDim.x * gridDim.x; // 计算总线程数作为步长
    for (; i < n; i += stride) {
        C[i] = A[i] + B[i];
    }
}

这个内核的关键点在于:

  • blockIdx.x:当前线程块在网格中的索引。

  • blockDim.x:每个线程块中的线程数。

  • 通过 i = blockIdx.x * blockDim.x + threadIdx.x 计算出每个线程负责的全局起始索引。

  • 使用 stride 进行循环,使有限数量的线程能够处理任意大小的数组。

主机调用需要配置合适的网格和线程块大小:

int threadsPerBlock = 256;
int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock; // 向上取整
parallelVectorAdd<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, n);

利用共享内存优化:矩阵乘法示例

仅仅启动并行线程还不够,性能优化的核心在于智能地使用共享内存。我们以矩阵乘法为例。

一个朴素的矩阵乘法内核会频繁地从全局内存读取数据,速度很慢。优化的思路是将计算分块,先将数据块从全局内存加载到共享内存,然后在共享内存中进行高速计算。

以下是利用共享内存的矩阵乘法核心思想(伪代码表示):

__global__ void matrixMulShared(float *A, float *B, float *C, int width) {
    // 为子矩阵A和B声明共享内存
    __shared__ float sA[TILE_WIDTH][TILE_WIDTH];
    __shared__ float sB[TILE_WIDTH][TILE_WIDTH];

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    // 计算C中当前线程要处理的元素坐标
    int row = by * TILE_WIDTH + ty;
    int col = bx * TILE_WIDTH + tx;

    float sum = 0;
    // 循环遍历所有数据块
    for (int m = 0; m < width/TILE_WIDTH; ++m) {
        // 协作地将数据块从全局内存加载到共享内存
        sA[ty][tx] = A[row * width + (m * TILE_WIDTH + tx)];
        sB[ty][tx] = B[(m * TILE_WIDTH + ty) * width + col];
        __syncthreads(); // 等待块内所有线程完成加载

        // 在共享内存中进行子矩阵乘法计算
        for (int k = 0; k < TILE_WIDTH; ++k) {
            sum += sA[ty][k] * sB[k][tx];
        }
        __syncthreads(); // 等待计算完成,再进行下一轮数据加载
    }
    // 将结果写回全局内存
    C[row * width + col] = sum;
}

这种分块策略将全局内存访问量从 O(n^3) 显著降低到 O(n^2),因为每个数据块只需从全局内存加载一次到共享内存,然后在该块内的所有计算中重复使用。

内核融合

除了使用共享内存,另一个关键的优化技术是内核融合。在标准框架中,连续的运算(如矩阵乘法和ReLU激活)会分别启动独立的内核,每个内核都会将中间结果写回全局内存,再由下一个内核读回,造成了不必要的延迟和带宽消耗。

内核融合将多个连续操作合并到单个CUDA内核中执行。例如,将 Y = matmul(M, X)Z = relu(Y) 融合成一个内核 fused_matmul_relu。这样,中间结果 Y 可以保留在寄存器或共享内存中,直接用于ReLU计算,完全避免了写回和读取全局内存的开销。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/6d667130fec4a41e931978441be370e2_1.png

这种优化虽然增加了内核编写的复杂性,但通常能带来10%-20%的性能提升,对于大规模训练至关重要。这也是为什么追求极致性能的定制化AI系统(如Flash Attention)会将整个注意力机制实现为一个融合内核的原因。

总结

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/6d667130fec4a41e931978441be370e2_3.png

本节课中我们一起学习了CUDA编程的核心概念。我们了解到,为了极致优化生成式AI模型的训练性能,尤其是应对超大规模模型和硬件错误,直接进行CUDA级编程是必要的。关键要点包括:理解GPU的线程-块-网格层次结构和共享内存的重要性;掌握编写并行内核的基本方法;学会通过分块策略利用共享内存优化数据密集型运算(如矩阵乘法);以及认识内核融合技术对于减少冗余内存访问的巨大价值。这些知识是理解现代高性能AI系统底层实现的基础。

28:CUDA编程基础 🚀

在本节课中,我们将学习CUDA编程的基础知识。理解CUDA编程对于希望从底层优化和训练大型语言模型至关重要。我们将探讨为何需要CUDA编程、GPU的基本架构、如何编写简单的CUDA内核,以及如何利用共享内存和融合内核来提升计算效率。

概述 📋

CUDA编程是使用C/C++在NVIDIA GPU上进行编程的简称。在生成式AI领域,为了高效训练大型语言模型,我们常常需要绕过现有的高级库(如PyTorch),直接编写底层的CUDA代码以获得显著的性能提升。本节课将介绍CUDA编程的核心概念,包括线程、块、网格的结构,以及如何通过优化内存访问来加速计算。

GPU架构与CUDA编程模型 🏗️

上一节我们介绍了学习CUDA编程的必要性,本节中我们来看看GPU的基本架构和CUDA编程模型。

GPU的计算核心组织成一种层次结构。在最底层是线程,每个线程类似于CPU中的一个寄存器,执行最基本的计算单元。多个线程被组织成一个线程块。多个线程块进一步组成一个网格。一个CUDA内核调用通常对应整个网格。

这种结构设计主要与内存层次有关,而非纯粹的计算结构。在一个线程块内的所有线程共享一块快速的共享内存。这块内存比GPU的全局内存(显存)访问速度快得多。CUDA编程的核心挑战之一就是尽量减少从慢速的全局内存中读取数据的次数,尽可能多地利用快速的共享内存进行计算。

编写第一个CUDA内核:向量加法 ➕

理解了基本架构后,我们通过一个简单的例子来学习如何编写CUDA内核函数。

一个CUDA内核函数是一个在GPU上每个线程上并行执行的函数。我们以向量加法 C = A + B 为例。一个天真的实现是只使用一个线程,通过一个循环遍历所有元素进行加法。但这完全没有利用GPU的并行能力。

为了并行化,我们需要让多个线程同时工作。每个线程负责计算结果向量中不同位置的和。CUDA运行时提供了内置变量,如 threadIdx.xblockDim.x,它们分别表示当前线程在线程块内的索引和线程块中的线程总数。

以下是利用多线程进行向量加法的核心思路:

// 假设每个线程块有 blockDim.x 个线程
int i = threadIdx.x + blockIdx.x * blockDim.x;
int stride = blockDim.x * gridDim.x; // 总线程数
for (; i < N; i += stride) {
    C[i] = A[i] + B[i];
}

这段代码中,每个线程根据其唯一的全局索引 i 计算对应的向量元素。循环中的 stride 确保了即使向量长度 N 远大于总线程数,所有元素也能被覆盖到。通过这种方式,我们实现了计算的完全并行化。

优化关键:利用共享内存进行矩阵乘法 ✖️

上一节我们看到了如何并行化简单的向量操作,本节中我们来看看更复杂的操作——矩阵乘法,并学习如何使用共享内存进行优化。

在矩阵乘法 C = A * B 中,朴素的方法是每个线程计算输出矩阵 C 中的一个元素。这需要该线程读取 A 的一整行和 B 的一整列数据,这些数据都来自全局内存,访问速度很慢。

优化的关键在于利用线程块的共享内存。基本思想是将大矩阵分块。例如,将矩阵 AB 划分为多个 BxB 的子块。计算时,先将 AB 对应的子块从全局内存加载到共享内存中,然后线程块内的所有线程协作,在共享内存中完成子块的矩阵乘法计算。

以下是这个过程的简化描述:

  1. 将矩阵 A 的一个 BxB 子块加载到共享内存 As

  2. 将矩阵 B 的一个 BxB 子块加载到共享内存 Bs

  3. 线程块内所有线程同步,确保数据加载完成。

  4. 每个线程使用 AsBs 中的数据计算输出子块中自己负责的部分。

  5. 重复步骤1-4,遍历所有需要的子块对,并累加结果。

如果不使用共享内存,计算一个 BxB 输出子块需要 O(B^3) 次全局内存访问。而使用共享内存后,只需要 O(B^2) 次全局内存访问(用于加载子块),后续的 O(B^3) 次计算全部在快速的共享内存中进行,从而大幅提升性能。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/2f35edc6554a6991105a52ca3aacdb72_1.png

进阶技巧:融合内核 ⚡

我们了解了如何用共享内存优化单个操作。但在实际模型中,多个操作常常连续发生。本节介绍一种重要的优化技术:融合内核

考虑一个简单的操作序列:先进行矩阵乘法 Y = matmul(M, X),然后对结果 Y 应用ReLU激活函数。在标准的PyTorch写法中,这两个操作是独立的。计算流程如下:

  1. 计算 matmul(M, X),结果 Y 被写回全局内存。

  2. 从全局内存读取 Y,计算 relu(Y),结果再写回全局内存。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/2f35edc6554a6991105a52ca3aacdb72_3.png

这个过程在全局内存和计算单元之间产生了不必要的往返。融合内核的思想是将多个操作合并到单个CUDA内核中。对于这个例子,我们可以编写一个“MatMul + ReLU”融合内核:

  1. 线程计算 matmul 的部分结果。

  2. 该部分结果保留在寄存器或共享内存中,不写回全局内存

  3. 立即对该中间结果应用 ReLU 函数。

  4. 将最终结果写回全局内存。

通过避免中间结果的全局内存读写,融合内核可以显著减少内存带宽压力,通常能带来10%-20%的性能提升。像Flash Attention这样的先进优化,本质上就是将注意力机制中的矩阵乘、Softmax、缩放等多个步骤融合到一个精心设计的内核中。

总结 🎯

本节课我们一起学习了CUDA编程的基础知识。我们首先了解了为什么在训练超大语言模型时需要绕过PyTorch等框架进行底层CUDA编程。接着,我们探讨了GPU的线程-块-网格架构及其对应的内存层次(共享内存 vs. 全局内存)。通过向量加法和矩阵乘法的例子,我们学习了如何编写并行化的CUDA内核,并利用共享内存优化数据访问。最后,我们介绍了融合内核的概念,它将多个连续操作合并,以减少耗时的全局内存访问。

掌握这些基础概念对于理解现代大模型训练中的高性能计算优化至关重要。

29:状态空间模型 🧠

在本节课中,我们将要学习状态空间模型。这是一种结合了循环神经网络和Transformer架构思想的新型模型,旨在解决传统RNN在并行化和长程记忆方面的不足,同时保持线性计算复杂度。我们将从基础概念开始,逐步探讨其核心公式、优化方法以及实际应用。


https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_1.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_2.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_4.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_6.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_8.png

状态空间模型的基础

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_10.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_11.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_13.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_14.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_15.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_17.png

上一节我们介绍了生成式AI的背景,本节中我们来看看状态空间模型的基本概念。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_19.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_21.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_23.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_25.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_26.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_28.png

状态空间模型是由Albert Gu等人提出的一系列工作。它本质上是对Transformer出现之前的循环神经网络的一种升级。在Transformer架构之前,人们主要使用循环神经网络处理自然语言。

循环神经网络的回顾

循环神经网络的核心思想是维护一个状态。当输入一个序列时,RNN会根据输入逐步更新这个状态。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_30.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_31.png

其最简单的形式可以用以下公式描述:

h_t = f(h_{t-1}, x_t)

其中 h_t 是时刻 t 的隐藏状态,x_t 是时刻 t 的输入,f 是状态更新函数。

这种架构类似于人脑的工作方式:你逐词阅读,大脑状态随之逐步更新。

RNN的局限性

然而,传统的RNN存在两个主要问题:

  1. 缺乏并行化:RNN的更新是严格顺序的,这与现代GPU偏好并行计算的特点不兼容。Transformer则通过自注意力层和前馈层实现了高度并行计算。

  2. 缺乏长程记忆:RNN的状态是逐步更新的,某个时刻的标记(Token)难以直接回溯到很久之前的上下文去寻找答案。而Transformer的自注意力机制允许任何标记直接关注序列中任何位置的标记。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_33.png

正是这些缺点导致RNN一度被Transformer取代。但状态空间模型的出现,似乎让这种RNN结构重新焕发了活力。


状态空间模型的架构

上一节我们回顾了RNN的优缺点,本节中我们来看看状态空间模型是如何构建的。

状态空间模型可以看作是RNN与Transformer的一种结合。其核心思想是用一个类RNN的结构替换Transformer中的自注意力层,同时保留其前馈层。

一个基本的状态空间模型块结构如下:

  1. 一个状态空间层(替代自注意力层),用于处理序列信息。

  2. 一个前馈层(如MoE或MLP),用于进行逐标记的局部处理。

你可以将多个这样的块堆叠起来,形成一个深度模型,就像堆叠Transformer块一样。

核心公式:线性状态空间模型

状态空间模型的核心是一个线性化的RNN。它通过一个连续的或离散的线性过程来更新状态。

考虑一个离散时间的线性状态空间模型,其公式如下:

h_t = Ā * h_{t-1} + B̄ * x_t

y_t = C * h_t

其中:

  • x_t 是输入序列。

  • y_t 是输出序列。

  • h_t 是隐藏状态。

  • , , C 是可学习的参数矩阵。

通常是通过对连续时间公式进行离散化(如零阶保持法)得到的:

Ā = exp(Δ * A)

B̄ = (exp(Δ * A) - I) * (Δ * A)^{-1} * Δ * B

其中 Δ 是步长参数。

从状态更新到卷积

如果我们展开上述更新公式,会发现输出 y_t 实际上是输入 x 与一个卷积核 的卷积结果:

y = x * K̄

其中卷积核 的元素由 C * Ā^{k} * B̄ 决定(k 为时间步偏移量)。

这意味着,线性状态空间模型在数学上等价于一个(可能无限长的)卷积操作。


状态空间模型的优化:S4

上一节我们介绍了基础模型,本节中我们来看看如何优化它,使其更高效、更强大。

基础模型存在计算效率低和缺乏类似Transformer多头机制的问题。S4模型通过对角化技术来解决这些问题。

对角化与多头机制

在S4中,参数矩阵 A 被约束为对角矩阵。这大大减少了参数量(从 N×N 降至 N),并使得计算 的幂次变得非常简单(只需对每个对角线元素进行幂运算)。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/92b5331b966f21943acecf543713d871_35.png

更重要的是,这种对角化形式天然支持一种“多头”机制。我们可以将输入的每个特征维度(通道)视为独立的,并为每个通道配备一组独立的 A_i, B_i, C_i 参数。

以下是其工作原理:

  • 输入 x_t 是一个向量。

  • 对于该向量的第 i 个维度(通道),我们应用一个独立的状态空间模型:

    h_t^i = Ā_i * h_{t-1}^i + B̄_i * x_t^i

    y_t^i = C_i * h_t^i

  • 每个通道的卷积核 K̄_i 是不同的,这允许模型在不同的特征维度上捕获不同的时间模式。

例如,我们可以让某些通道的 Ā_i 接近1,使其关注长期历史(类似于求平均),而让另一些通道的 Ā_i 很小,使其只关注近期信息。这实现了某种与位置相关的编码功能。

尽管如此,S4的卷积核是上下文无关的,它们在训练后是固定的,不随输入内容变化。


状态空间模型的进化:Mamba (S6)

上一节我们介绍了S4,本节中我们来看看其关键进化——Mamba模型(在论文中常称为S6),它如何引入上下文感知能力。

S4的主要限制在于其卷积核是静态的。Mamba的核心改进是让参数 B, C 以及步长 Δ 成为输入 x_t 的函数,从而使模型能够根据输入内容动态调整其行为。

选择性机制

在Mamba中,我们有以下变化:

B_t = Linear_B(x_t)

C_t = Linear_C(x_t)

Δ_t = Linear_Δ(x_t)

其中 Linear_* 是简单的线性投影层。

这意味着:

  • B_tC_t 现在扮演着类似Transformer中“键”和“查询”的角色,它们基于当前上下文动态生成。

  • Δ_t 成为一个时间相关的缩放因子。

这个过程被称为选择性机制。模型可以根据当前输入 x_t(它已编码了之前的上下文信息)来决定当前标记的重要性。例如,对于关键词或实体名称,模型可以生成较大的 B_t,让该标记对隐藏状态产生更大影响;对于不重要的虚词,则可以忽略其影响。

计算效率与线性时间复杂度

Mamba的另一个巨大优势是其线性时间复杂度。与Transformer自注意力的 O(n²) 复杂度不同,状态空间模型按顺序处理序列,复杂度为 O(n)。这对于处理超长序列(如长文档、视频、音频)至关重要。

为了实现高速计算,Mamba使用了高度优化的CUDA内核,确保关键的中间状态(隐藏状态 h_t)始终驻留在GPU的高速缓存(SRAM)中,避免了与慢速显存(VRAM)的频繁数据交换。这是其性能远超朴素PyTorch实现的关键。

模型参数与性能权衡

状态空间层(尤其是经过对角化后)的参数数量远小于标准的自注意力层。这意味着,在总参数量固定的情况下,使用状态空间模型的网络可以将更多参数分配给前馈层(如MoE),而前馈层通常对模型的知识存储和复杂模式建模能力贡献更大。因此,在同等计算预算或参数量下,Mamba架构的模型可能表现更优。


总结与展望

本节课中我们一起学习了状态空间模型,从传统的RNN出发,探讨了其基本形式S4和先进的Mamba模型。

我们了解到:

  1. 状态空间模型 本质上是线性RNN,可视为一个卷积操作。

  2. S4模型 通过对角化和为每个通道配备独立动态,引入了高效的多头机制,但卷积核是静态的。

  3. Mamba模型 通过使关键参数依赖于输入,实现了选择性机制,从而能够动态聚焦于相关上下文。

  4. 该架构的核心优势在于线性计算复杂度高效的硬件利用,使其特别适合处理长序列数据。

状态空间模型并非要完全取代Transformer的自注意力机制,而是提供了一种高效的替代方案。在许多任务中,这种线性时间、上下文敏感的“总结”与“筛选”机制已经足够。对于需要精确回溯或多跳推理的复杂任务,未来可能会看到混合架构的出现,例如将Mamba与局部注意力相结合。无论如何,状态空间模型为生成式AI模型的设计开辟了一条富有前景的新路径。

30:PyTorch与Weights & Biases入门教程 🚀

在本节课中,我们将学习深度学习库PyTorch的基础知识,以及如何使用Weights & Biases(W&B)工具来记录和可视化训练过程。本教程旨在帮助初学者快速上手,为后续课程作业做好准备。

PyTorch基础:张量与操作 🔢

PyTorch的核心数据结构是张量(Tensor),可以将其理解为高维数组或矩阵。它与NumPy非常相似,但支持GPU加速和自动微分。

创建张量

以下是创建张量的几种基本方法:

  • 从列表创建:使用 torch.tensor() 函数。

    data = [[1, 2], [3, 4]]
    x = torch.tensor(data, dtype=torch.float32)
    
  • 从NumPy数组创建:使用 torch.from_numpy() 函数。

    import numpy as np
    np_array = np.ones((2, 2))
    x = torch.from_numpy(np_array)
    
  • 创建特殊张量:PyTorch提供了类似NumPy的函数来创建全零、全一或随机张量。

    zeros_tensor = torch.zeros((2, 3)) # 2x3的全零张量
    ones_tensor = torch.ones((2, 3))   # 2x3的全一张量
    rand_tensor = torch.rand((2, 3))   # 2x3的随机张量(值在0-1之间)
    

张量操作

创建张量后,可以进行各种数学运算。

  • 基本算术运算:支持加(+)、减(-)、乘(*)、除(/)和矩阵乘法(@torch.matmul)。

    a = torch.rand(2, 2)
    b = torch.rand(2, 2)
    c = a + b  # 逐元素相加
    d = a @ b  # 矩阵乘法
    
  • 形状操作view() 用于改变张量形状(要求内存连续),reshape() 更通用。transpose() 用于转置,squeeze()unsqueeze() 用于删除或添加维度。

    x = torch.rand(4, 5)
    y = x.view(10, 2)      # 改变形状为10x2
    z = x.transpose(0, 1)  # 转置,形状变为5x4
    
  • 拼接与堆叠torch.cat() 沿现有维度拼接张量,torch.stack() 沿新维度堆叠张量。

    a = torch.rand(2, 3)
    b = torch.rand(2, 3)
    c = torch.cat([a, b], dim=0) # 沿第0维(行)拼接,形状变为(4, 3)
    d = torch.stack([a, b], dim=0) # 沿新维度堆叠,形状变为(2, 2, 3)
    

上一节我们介绍了张量的基本操作,本节中我们来看看PyTorch如何构建计算图以实现自动微分。

自动微分(Autograd)

PyTorch的自动微分功能是其核心特性之一。要使用它,只需在创建张量时设置 requires_grad=True

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
z = y.sum()
z.backward() # 计算梯度
print(x.grad) # 输出梯度 tensor([2., 4.])

关键点:

  1. 调用 .backward() 方法后,叶子节点(用户直接创建的张量)的梯度会被计算并存储在 .grad 属性中。

  2. 非叶子节点(计算中间结果)的梯度默认不会保留,以节省内存。

  3. 在训练循环中,通常的模式是:计算损失 -> loss.backward() -> 优化器更新参数 (optimizer.step()) -> 清空梯度 (optimizer.zero_grad())。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_1.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_3.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_5.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_6.png

数据加载与处理 📂

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_8.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_9.png

在训练模型前,需要有效地加载和组织数据。PyTorch提供了 DatasetDataLoader 类来简化这一过程。

使用内置数据集

torchvision 库包含许多计算机视觉领域的常用数据集。

import torchvision
from torchvision import transforms

# 定义图像转换(如转为张量、归一化)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载Fashion-MNIST数据集
train_dataset = torchvision.datasets.FashionMNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

test_dataset = torchvision.datasets.FashionMNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform
)

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_11.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_13.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_15.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_17.png

Fashion-MNIST数据集包含10类服装的灰度图像,训练集60000张,测试集10000张,每张图像大小为28x28。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_19.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_21.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_23.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_24.png

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_25.png

创建自定义数据集

对于自己的数据,可以通过继承 torch.utils.data.Dataset 类来创建数据集。

from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, feature_files, annotation_files, transform=None):
        # 初始化,加载文件路径等元数据
        self.features = feature_files
        self.labels = annotation_files
        self.transform = transform

    def __len__(self):
        # 返回数据集大小
        return len(self.features)

    def __getitem__(self, idx):
        # 根据索引加载单个样本(如图像)和标签
        feature = load_feature(self.features[idx]) # 需实现load_feature函数
        label = self.labels[idx]

        if self.transform:
            feature = self.transform(feature)

        return feature, label

最佳实践:对于大型数据集,建议在 __getitem__ 方法中动态加载数据(如从磁盘读取图像),而不是在 __init__ 中全部加载到内存。这样便于进行随机数据增强。

使用DataLoader

DataLoader 负责从 Dataset 中按批次抽取数据,并支持多进程加载、数据打乱等功能。

from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,      # 训练时打乱数据顺序
    num_workers=2      # 使用2个子进程加载数据
)

test_loader = DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False,     # 测试时通常不需要打乱
    num_workers=2
)

# 迭代数据
for batch_idx, (images, labels) in enumerate(train_loader):
    # images形状: [64, 1, 28, 28]
    # labels形状: [64]
    # ... 进行训练 ...
    pass

构建神经网络模型 🧠

在PyTorch中,神经网络通常通过继承 torch.nn.Module 类来定义。

定义基础模型

你需要定义 __init__ 方法来初始化网络层,以及 forward 方法来定义前向传播过程。

import torch.nn as nn
import torch.nn.functional as F

class SimpleMLP(nn.Module):
    def __init__(self, input_size=784, hidden_size=128, num_classes=10):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # 将图像展平
        x = x.view(-1, 28*28)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

注意:你只需定义 forward 方法。反向传播(计算梯度)由PyTorch的 autograd 系统自动处理,只需调用 loss.backward()

自定义网络层

虽然不常用,但你可以通过继承 nn.Moduletorch.autograd.Function 来创建自定义层。

https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_27.png

class CustomLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input @ weight.t()
        if bias is not None:
            output += bias
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None
        # ... 计算梯度 ...
        return grad_input, grad_weight, grad_bias

<https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_29.png>

<https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_30.png>

<https://github.com/OpenDocCN/dsai-notes-pt2-zh/raw/master/docs/cmu-10423-genai/img/b89ecb09ae3ef6d1e615fc3fa91297f9_31.png>

# 使用自定义函数封装成模块
class CustomLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.randn(out_features))

    def forward(self, x):
        return CustomLinearFunction.apply(x, self.weight, self.bias)

更多推荐