StreamSparkAI

概览/核心摘要 (Executive Summary)

本讲座（Stanford CS336, Spring 2025）深入探讨了从零开始构建语言模型所需的PyTorch核心组件和资源核算方法。核心目标是让学习者理解模型构建的底层机制，并掌握内存与计算效率的量化分析。讲座首先通过估算大规模模型（如700亿参数模型）训练时间和单机（8卡H100）可训练最大模型规模（约400亿参数，不含激活）来强调资源核算的重要性，指出效率直接关系到成本。

讲座详细介绍了PyTorch中的张量（Tensor）及其不同浮点数表示（float32, float16, bfloat16, fp8）对内存占用的影响，强调bfloat16在动态范围和内存效率上的优势。接着，讨论了计算资源（CPU/GPU）、张量操作（视图、复制、einops库）及其计算成本（FLOPs）。核心结论是矩阵乘法（matmul）主导计算量，其FLOPs约为2 * M * K * N。讲座引入了模型FLOPs利用率（MFU）作为衡量硬件效率的指标。关于梯度计算，指出反向传播的计算量约是前向传播的两倍，因此训练一个参数的总FLOPs约为6 * num_tokens * num_params。最后，概述了模型参数初始化、自定义模型构建、数据加载、优化器（以Adagrad为例说明其状态内存需求，Adam优化器通常需要额外存储两倍参数量的状态，即参数、梯度、优化器状态共需约16字节/参数，不含激活）、训练循环和混合精度训练等实践环节。

PyTorch基础与资源核算动机

Speaker 1首先回顾了上节课关于语言模型概览及从零构建原因的内容，并提及了Tokenization。本讲座聚焦于实际模型构建，涵盖PyTorch原语、模型、优化器和训练循环，并特别关注效率，即内存和计算资源的使用。

资源评估的启发性问题

为了引出资源核算的重要性，讲座提出了两个可通过“餐巾纸计算”（napkin math）解答的问题：

1. 训练时长估算：在1024张H100上训练一个700亿参数的稠密Transformer模型，使用15万亿Token，需要多长时间？

   • 计算方法：
     • 总FLOPs ≈ 6 * num_parameters * num_tokens
     • H100单卡FLOPs/秒（假设MFU为0.5）
     • 计算集群总FLOPs/天
     • 总FLOPs / 集群每日FLOPs ≈ 144天
   • “这个6倍参数量乘以Token数的公式来源，将在本次讲座中讨论。”

2. 最大可训练模型规模估算：在8张H100（每张80GB HBM显存）上，使用AdamW优化器（不采用特别技巧），能训练的最大模型参数量是多少？

   • 计算方法：
     • 每个参数所需字节数（参数、梯度、优化器状态）：16字节 (讲座后续会解释来源)
     • 总可用显存 / 每参数字节数 ≈ 400亿参数
   • “这是一个非常粗略的估算，因为它没有考虑激活值（activations）的显存占用，激活值大小取决于批量大小（batch size）和序列长度（sequence length）。”

Speaker 1强调：“效率是关键（efficiency is the name of the game）。要做到高效，你必须确切知道你实际消耗了多少FLOPs，因为当这些数字变得巨大时，它们直接转化为美元，而你希望这个成本尽可能小。”

学习目标

• 机制 (Mechanics)：PyTorch及其底层工作原理。
• 思维模式 (Mindset)：资源核算。
• 直觉 (Intuitions)：目前主要关注机制和思维模式，关于模型性能的直觉暂为宏观层面。

内存核算 (Memory Accounting)

张量 (Tensors)

• 深度学习中存储一切（参数、梯度、优化器状态、数据、激活值）的基础构建块。
• 创建方式多样，包括未初始化创建。

浮点数表示与内存占用

• float32 (fp32 / 单精度):

  • 32位：1位符号，8位指数，23位尾数。
  • 每个元素占 4字节。
  • 被认为是计算领域的“黄金标准”，在机器学习中常被称为“全精度”（尽管科学计算领域可能使用float64或更高精度）。
  • 内存计算：num_elements * element_size_in_bytes。
  • 示例：一个 4x8 的 float32 张量占用 32 * 4 = 128 字节。
  • GPT-3中一个前馈层权重矩阵（2048 x 8192）约占用 2.3GB。

• float16 (fp16 / 半精度):

  • 16位：1位符号，5位指数，10位尾数。
  • 内存减半，计算通常也更快。
  • 缺点：动态范围（dynamic range）有限，容易出现上溢（overflow）或下溢（underflow），例如 1e-8 在 float16 中会变成0。不适用于大型模型训练或需要高精度表示的场景。

• bfloat16 (bf16 / Brain Float):

  • 由Google于2018年开发，专为深度学习设计。
  • 16位：1位符号，8位指数，7位尾数。
  • 特点：与 float16 占用相同内存，但拥有与 float32 相同的动态范围，牺牲了尾数精度。
  • 深度学习对动态范围的需求高于精度，因此 bf16 表现良好。1e-8 在 bf16 中不会是0。
  • “bf16基本上是进行计算时通常会使用的数据类型，因为它对于前馈和反向传播计算来说足够好。”
  • 注意：存储优化器状态和参数本身通常仍推荐使用 float32 以免训练不稳定。

• fp8 (8位浮点数):

  • 由Nvidia于2022年开发。
  • 位数极少，表示非常粗略。存在两种变体，侧重于不同方面（更高分辨率或更大动态范围）。
  • H100 GPU支持 fp8。

• 混合精度训练 (Mixed Precision Training):

  • 根据模型不同部分的需求选择最低可用精度。
  • 例如：参数和优化器状态使用 float32，前向传播和反向传播中的矩阵乘法使用 bf16，某些如Attention的敏感部分可能仍用 float32。
  • Speaker 1观点：“目前来看，在深度学习中你可能不想使用float16。”

计算核算与硬件 (Compute Accounting & Hardware)

CPU 与 GPU 数据传输

• PyTorch张量默认在CPU上创建和存储。
• 使用GPU进行计算必须显式将张量转移到GPU：tensor.to('cuda') 或创建时指定 device='cuda'。
• 数据从CPU RAM传输到GPU HBM（High Bandwidth Memory）需要时间。
• “在PyTorch中处理张量时，应始终清楚它驻留在哪里。”

硬件概览

• 示例中使用的GPU：NVIDIA H100，拥有80GB HBM。

张量操作与视图 (Tensor Operations & Views)

• 张量的本质：PyTorch中的张量是指向已分配内存区域的指针，并包含元数据（如stride）来解释如何索引该内存。
• 视图 (Views)：许多操作（如索引、tensor.view()、tensor.transpose()）不会创建数据的副本，而是创建指向相同底层存储的新“视图”。
  • 修改视图会影响原始张量，反之亦然。
  • tensor.storage().data_ptr() 可检查两个张量是否共享底层存储。
• 连续性 (Contiguity)：
  • 如果张量在内存中是按其逻辑顺序连续存储的，则为连续张量。
  • 转置等操作可能导致非连续张量。
  • 对非连续张量执行某些视图操作可能会失败。此时需先调用 tensor.contiguous()，这会创建一个数据的副本使其连续。
  • tensor.reshape() 行为类似 tensor.contiguous().view()。
• 建议：视图操作本身开销很小（不分配新内存），可多加使用以提高代码可读性。但需注意 contiguous() 或 reshape() 可能触发数据复制。

创建新张量的操作

• 元素级操作（element-wise operations）通常会创建新张量。
• torch.triu() (上三角矩阵) 是一个元素级操作，常用于创建因果注意力掩码 (causal attention mask)。

矩阵乘法 (Matrix Multiplications / Matmuls)

• 深度学习的核心运算。
• 支持批处理：例如，一个 (batch_size, seq_len, input_dim) 的张量与一个 (input_dim, output_dim) 的权重矩阵相乘，PyTorch会自动处理为在每个batch和sequence position上进行独立的矩阵乘法，结果为 (batch_size, seq_len, output_dim)。

einops 库

• 受爱因斯坦求和约定启发，用于以更清晰、不易出错的方式处理张量维度操作。
• 动机：避免使用难以理解的负数索引（如 transpose(-2, -1)）。
• 核心思想：为维度命名。
  • jax typing 库可用于在类型注解中以字符串形式声明维度名称（如 Float[Tensor, "batch seq_len hidden_dim"]），但PyTorch本身不强制执行。
• einsum (from einops):
  • 用于执行带维度名称的矩阵乘法或更广义的张量缩并。
  • 示例：einsum(x, y, 'b s1 h, b s2 h -> b s1 s2') 表示将两个张量 x (batch, seq1, hidden) 和 y (batch, seq2, hidden) 沿 hidden 维度相乘并求和，保留 batch, seq1, seq2 维度。
  • ... 可用于表示任意数量的前导广播维度。
  • Speaker 1观点：“一旦你习惯了它，它会比使用-2、-1索引好得多。” 并且 torch.compile 可以有效地编译 einsum 操作。
• reduce (from einops): 用于按名称聚合维度。
• rearrange (from einops): 用于重排、拆分或合并维度，类似更强大的 view。
  • 示例：将 batch seq (heads hidden_dim) 拆分为 batch seq heads hidden_dim。

张量操作的计算成本 (FLOPs)

• FLOP (Floating Point Operation)：指单个浮点运算（如加法、乘法）。衡量计算量。
• FLOPs/sec (Floating Point Operations per second)：衡量硬件速度。本课程避免使用大写S的FLOPs，而明确写为 FLOPs/sec。
• 数量级参考：
  • GPT-3训练约消耗 3.23e23 FLOPs。
  • GPT-4训练（推测）约消耗 2e25 FLOPs。
  • 美国曾有行政命令要求超过 1e26 FLOPs 的基础模型需向政府报告（后被撤销）。欧盟AI法案中有 1e25 FLOPs 的阈值。
• 硬件性能：
  • A100峰值性能（fp16/bf16）：312 TFLOPs/sec。
  • H100峰值性能（fp16/bf16）：1979 TFLOPs/sec (带稀疏性)，约一半即 ~990 TFLOPs/sec (不带稀疏性)。
  • fp32 在H100上的性能远低于此。
  • 稀疏性说明：NVIDIA宣传的稀疏性加速特指“结构化稀疏性（structured sparsity），即每4个元素中有2个为0”，Speaker 1评论：“没人用它，这是市场部门的说法。”
• 矩阵乘法的FLOPs：对于 (M, K) @ (K, N) -> (M, N)，FLOPs ≈ 2 * M * K * N (每个输出元素需要K次乘法和K-1次加法，近似为2K)。
  • “你应该记住，如果你在做矩阵乘法，FLOPs数量是三个维度乘积的两倍。”
• 其他操作的FLOPs通常与张量大小成线性关系，远小于大型矩阵乘法。因此，性能分析常聚焦于矩阵乘法。

壁钟时间与模型FLOPs利用率 (MFU)

• MFU (Model FLOPs Utilization) = (实际有效FLOPs / 实际耗时) / 硬件理论峰值FLOPs/sec。
  • 衡量模型在特定硬件上实际利用硬件计算能力的程度。
  • MFU > 0.5 通常被认为是好的。
  • MFU通常在矩阵乘法占主导时较高。
• 实验数据 (H100, 32x32x32 matmul):
  • float32: 耗时0.16s, 实际FLOPs/sec 5.4e13, H100 float32 理论峰值 67 TFLOPs/sec, MFU ≈ 0.8。
  • bfloat16: 耗时0.03s, 实际FLOPs/sec更高。但相对于bfloat16的理论峰值（约990 TFLOPs/sec），MFU可能反而较低（示例中为0.26），表明理论峰值有时过于乐观。
• 建议：始终对代码进行基准测试，不要假设能达到理论性能。

梯度与反向传播的计算成本

• 两层线性网络示例：Y = W2 @ (W1 @ X)
  • 前向传播FLOPs：2 * num_tokens * (params_W1 + params_W2)。对于简单线性模型，可简化为 2 * num_tokens * total_params。
  • 反向传播FLOPs：
    • 计算 dL/dW2 涉及类矩阵乘法，FLOPs ≈ 2 * B * D * K (假设W2是D*K)。
    • 计算 dL/dH1 (H1是W1的输出) 涉及类矩阵乘法，FLOPs ≈ 2 * B * D * K。
    • 对W1重复此过程。
    • 核心结论：反向传播计算量大约是前向传播的两倍。
    • “前向传播是2倍参数量，反向传播是4倍参数量。” (这里的参数量指的应是与数据点数相乘后的总计算量基数)
  • 总FLOPs (训练一步)：前向 + 反向 ≈ (2 + 4) * num_tokens * num_params = 6 * num_tokens * num_params。
    • 这解释了讲座开头估算训练时长时使用的 6x 因子。
    • 此规则对许多模型（包括Transformer，当序列长度不过大时）大致成立。

构建模型 (Building Models)

参数 (Parameters)

• 在PyTorch中，模型参数通常存储为 torch.nn.Parameter 对象，它们是会自动记录梯度的张量。

参数初始化 (Parameter Initialization)

• 问题：朴素的高斯初始化（如 torch.randn）可能导致激活值随网络深度增加而爆炸或消失，因为方差会累积（输出标准差约与 sqrt(hidden_dimension) 成正比）。
• 解决方案 (Kaiming/Xavier 初始化思想)：按 1 / sqrt(input_dimension) 缩放初始权重。
  • 例如：W = nn.Parameter(torch.randn(in_dim, out_dim) / math.sqrt(in_dim))
  • 这有助于使输出激活值的方差保持在1附近。
• 额外技巧：截断正态分布（truncated normal），例如将值限制在 [-3, 3] 标准差范围内，以避免极端值。

自定义模型示例：Cruncher (深度线性网络)

• 一个包含 num_layers 个线性层（矩阵乘法）的简单模型。
• 使用 nn.Module 和 nn.ModuleList 构建。
• 参数量计算：若每层为 D x D 矩阵，共 L 层，加一个 D 维输出头，则总参数为 L * D*D + D。
• 模型和数据都需移至GPU (model.to('cuda'), data.to('cuda'))。

随机性管理 (Randomness)

• 随机性来源：参数初始化、Dropout、数据打乱顺序等。
• 最佳实践：始终为每个随机源传递固定的随机种子 (random_seed) 以保证可复现性。
  • torch.manual_seed(), numpy.random.seed(), random.seed()。
• “确定性是调试时的朋友。”

训练循环组件 (Training Loop Components)

数据加载 (Data Loading)

• 语言模型数据通常是Token ID序列（整数）。
• 内存映射 (Memory Mapping)：对于非常大的数据集（如Llama的2.8TB数据），一次性加载到内存不可行。可使用 numpy.memmap，它允许像操作内存数组一样操作磁盘上的文件，数据按需加载。
• torch.utils.data.Dataset 和 torch.utils.data.DataLoader 用于高效批处理数据。

优化器 (Optimizers)

• 常见优化器回顾：
  • SGD (Stochastic Gradient Descent)：沿负梯度方向更新。
  • Momentum：引入梯度的一阶动量（指数移动平均），加速收敛并抑制震荡。
  • Adagrad (Adaptive Gradient Algorithm)：为每个参数维护一个梯度的平方和，并用此来调整学习率（梯度大的参数学习率减小）。
  • RMSProp (Root Mean Square Propagation)：Adagrad的改进，使用梯度的平方的指数移动平均，避免学习率过早衰减。
  • Adam (Adaptive Moment Estimation)：结合了Momentum（一阶矩估计）和RMSProp（二阶矩估计）。是目前广泛使用的优化器。
• 实现自定义优化器 (以Adagrad为例)：
  • 继承 torch.optim.Optimizer。
  • 在 step() 方法中：
    1. 遍历 self.param_groups 中的参数。
    2. 访问参数的梯度 (param.grad)。
    3. 更新优化器状态（self.state[param]），例如Adagrad中累积梯度平方和 (g_squared_sum)。
    4. 根据优化算法更新参数值 (param.data.add_())。
    5. （可选）在 step() 结束时 zero_grad()（或在训练循环中显式调用）。
• 优化器状态的内存需求：
  • Adagrad：为每个参数额外存储一个浮点数（梯度平方和）。
  • Adam：为每个参数额外存储两个浮点数（一阶矩和二阶矩）。
  • 因此，使用Adam时，除了参数本身和梯度，优化器状态大约需要两倍于参数量的内存。
  • 总计（不含激活）：参数(P) + 梯度(P) + Adam状态(2P) = 4P。若为float32 (4字节/值)，则每个参数约需 16字节。

完整模型资源需求总结 (以简单线性网络为例)

• 参数 (Parameters)：D*D * num_layers + D (假设fp32，每个参数4字节)
• 激活 (Activations)：batch_size * D * num_layers (假设fp32，每个激活值4字节)。
  • “为什么需要存储激活值？朴素地看，因为在反向传播时，计算前一层的梯度依赖于后一层的激活值。如果更聪明些，可以不必存储所有激活值，可以通过重新计算来获得，这是一种称为激活检查点（activation checkpointing）的技术。”
• 梯度 (Gradients)：与参数量相同。
• 优化器状态 (Optimizer States)：Adagrad为1倍参数量，Adam为2倍参数量。
• 总内存 = 4 bytes * (num_params + num_activations + num_gradients + num_optimizer_states)
• 总FLOPs ≈ 6 * num_tokens * num_params

模型检查点 (Checkpointing)

• 语言模型训练耗时长，易中断。
• 应定期保存模型状态（model.state_dict()）和优化器状态（optimizer.state_dict()）以及当前迭代次数到磁盘，以便恢复训练。

混合精度训练进阶 (Mixed Precision Training)

• 权衡：高精度（准确、稳定，但昂贵） vs. 低精度（便宜，但可能不稳定）。
• 建议：
  • 默认使用 float32。
  • 尽可能尝试 bf16 甚至 fp8。
  • 一种常见策略：前向/反向传播使用低精度（如 bf16），参数更新和主权重保持 float32。
• PyTorch提供 torch.cuda.amp (Automatic Mixed Precision) 工具简化此过程。
• 前沿研究：探索全程使用 fp8 进行训练，挑战在于数值稳定性控制。模型设计与硬件特性协同发展，例如NVIDIA芯片对低精度（如 int4）的支持可能驱动新的模型架构。
• 训练 vs. 推理：训练对精度要求更高。模型训练完成后，在推理时可以采用更激进的量化策略（如 int8, int4) 以获取性能提升。

总结与展望

讲座系统梳理了从张量到完整训练循环的PyTorch原语，并重点讲解了内存和FLOPs的核算方法。通过作业一（Assignment 1）对Transformer模型进行类似的分析，将有助于巩固这些概念。