CS336 on Ganko Space

CS336 Lecture 5-6: GPUs & Kernels, Triton

Mon, 08 Jun 2026 00:00:00 +0000

为什么这两讲至关重要

在已经学过的课程中，L1-L4 讲的是"模型长什么样"，L7-L8 讲的是"多卡怎么协作"。L5-L6 卡在两者之间——单卡内部到底在发生什么。

如果你不理解 GPU 的硬件结构，你就无法理解：

为什么 Flash Attention 是过去五年 LLM 领域最重要的系统优化
为什么大矩阵乘法是"compute-bound"而 Attention 是"memory-bound"
为什么你写的 PyTorch 代码有时 GPU 利用率只有 30%
ZeRO-3 的通信开销到底卡在哪里（看懂 L7-8 那块没懂的）

Percy 在 L6 的原话很直白：“You cannot optimize what you don’t measure, and you cannot measure what you don’t understand.”

Lecture 5：GPU 架构解剖（Tatsu）

5.1 从一张 H100 说起

H100 的基本参数：

HBM（显存）：80GB，带宽 ~3.35 TB/s
SRAM（片上共享内存）：每 SM 约 256KB，总带宽 ~几十 TB/s（比 HBM 快一个数量级）
计算能力：BF16 下 ~990 TFLOPS

关键不在于这些数字有多大，而在于数字之间的比例。

CS336 Lecture 7-8：Parallelism

Thu, 09 Apr 2026 00:00:00 +0000

本篇主题

这两节Lecture都是关于并行训练的，主要内容包括：

理解训练超大模型时系统层面的复杂性
掌握不同的 parallelization paradigms，以及为什么人们通常会同时使用多种并行方式
了解大规模训练任务通常是如何进行的

Lecture 7 - Parallelism 1

单块GPU无法满足SCALING的需求，必须使用多块GPU进行训练,所以我们需要Multi-GPU、Multi-Machine的并行训练方法，如下图所示：

首先讲了一些Basic Concepts：

All Reduce：通信操作，所有参与的进程都将输入数据进行规约（如求和、最大值等）后，将结果分发给所有进程。常用于分布式训练中的梯度同步。
Broadcast：通信操作，数据从一个进程发送到所有其他进程。常用于分布式训练中的分发模型参数或者初始化数据。
All Gather：通信操作，所有参与的进程将各自的数据发送给所有其他进程，最终每个进程都获得所有数据的集合。常用于分布式训练中的收集模型输出或者中间结果。和All Reduce的区别在于，All Reduce会对数据进行规约操作（如求和），而All Gather只是简单地收集数据，不进行任何计算。
Reduce Scatter：通信操作，所有参与的进程将各自的数据发送给所有其他进程，并对数据进行规约操作（如求和），最终每个进程都获得规约后的结果的一部分。常用于分布式训练中的分布式梯度更新。

了解了初始知识之后，可以正式进入核心部分，不同的并行方式：

Data Parallelism：最早的并行方式，模型复制到每个GPU上，每个GPU处理不同的数据batch，计算梯度后进行同步更新。优点是实现简单；缺点是通信开销大，尤其是模型参数较大时。假设计算一个 SGD：我们会把 B 大小下的 batch 分给 M 个不同的机器，然后交换梯度去同步计算。这种情况对于 Compute scaling，每个 GPU 计算 B/M 个数据（不错！）；对于 Communication overhead，每个 batch 需要转移两次梯度（发送和接受）；对于 Memory scaling，完全没有，每个 GPU 都需要复制一次模型参数。早期的data parallelism问题是只缓解了计算压力，内存压力没有缓解，需要每个GPU都复制一份模型参数，通信压力也很大。 Zero[HTTPS://arxiv.org/pdf/1910.02054]可以解决这个问题，核心的思想是将模型的state（参数和优化器状态）分布在不同的GPU上，每个GPU只存储模型的一部分，这样可以显著减少每个GPU的内存占用，同时通过通信操作来同步更新模型参数。 ZeRO分三个阶段：

ZeRO-1：主要聚焦于 optimizer state sharding。把 optimizer state（first + second moments）分给每个 GPU，每个 GPU 都有 parameters + gradients，负责更新一部分 params。每个 GPU 根据分配到的 batch 子集计算完整的梯度，此时只拥有局部梯度利用 Reduce-Scatter 将所有 GPU 的局部梯度汇总并分散到每个设备上，现在每个 GPU 只持有全局梯度的一部分每个 GPU 利用局部梯度和 optimizer state 更新该部分参数利用 AllGather 将所有 GPU 的部分参数收集并分发给所有设备，确保每个 GPU 拥有完整的参数。这种方法的通信开销并没有增加，而且 memory 减少了接近四倍（优化器状态是fp32，参数是fp16）。

CS336 Lecture 2: Resource Accounting

Mon, 06 Apr 2026 00:00:00 +0000

本讲主线

这节课围绕一个核心问题展开：如何精确估算训练成本。主要分为三块：

内存占用
FLOPs 估算
模型与训练配置

这和 Lecture 1 提到的“效率优先”直接相关。

例题：70B 模型训练时长

问题：使用 1024 张 H100，在 15T tokens 上训练一个 70B 模型，需要多久？

根据 Transformer 常用训练量估算公式（$6 \times$ 参数量 $\times$ token 数）：

$$ \text{总 FLOPs} = 6 \times 70 \times 10^9 \times 15 \times 10^{12} = 6.3 \times 10^{24}\,\text{FLOPs} $$

单卡 H100（BF16）理论性能约为 $2 \times 10^{15}$ FLOPs/s。若考虑大模型训练实际效率约为 30%，则单卡有效 FLOPs/s 为：

$$ 0.3 \times 2 \times 10^{15} = 6 \times 10^{14}\,\text{FLOPs/s} $$

总时长：

$$ \text{总时间} = \frac{\text{总 FLOPs}}{\text{每秒有效 FLOPs} \times \text{GPU 数量}} = \frac{6.3 \times 10^{24}}{6 \times 10^{14} \times 1024} $$

约等于 119 天。

CS336 Lecture 3-4: Architectures, Hyperparameters, and MoE

Mon, 06 Apr 2026 00:00:00 +0000

本讲主线

老师推荐了 The Illustrated Transformer 作为辅助材料。

这节课基于 Tatsu H 对 2017-2025 年经典模型的统计，讨论两个问题：

不同模型在架构与超参数上有哪些共性与分歧？
这些演化背后的工程动机是什么？

架构与超参数观察

Pre-Norm vs Post-Norm 几乎所有现代 LLM（BERT 这类早期模型除外）都采用 Pre-Norm。核心原因是训练更稳定，梯度路径更顺畅。近年还出现了 double norm 等变体（如 Grok、Gemma 2），但尚未完全普及。
LayerNorm vs RMSNorm RMSNorm 不计算均值，只做均方根归一化，计算更简单。表面上看节省 FLOPs 不多，但在真实系统中，norm 的时间开销往往受 data movement 影响明显，因此 RMSNorm 的工程收益更像是在缓解内存墙，而不仅是减少算术量。
激活函数演化 从早期 ReLU，到 GPT 系列常见的 GeLU，再到近年广泛采用的 SwiGLU（门控激活变体）。门控结构通过逐元素调制，提升了表达能力与训练效率的折中表现。
串行层与并行层 传统 Transformer block 以串行堆叠为主。并行 block（如部分 GPT-J 设计）存在，但总体采用度不高。
位置编码路线 Sine/Absolute/Relative 都在不同阶段被主流模型采用。近年 RoPE（rotary position embeddings）几乎成为事实标准之一。
FFN 维度比（ffn dim : model dim） 传统经验常用 4 倍宽度；在 GLU 系列激活流行后，很多模型会选在 2.5-3.5 区间（常见约 8/3），以平衡参数规模与表达能力。

CS336-写在所有之前

Mon, 06 Apr 2026 00:00:00 +0000

CS336 是一门很好的课，但是很多内容都比较基础，或者在网上已经有很多资源了。

2026/04/06:我学到了Lecture4，目前来看，CS336的精华内容还是在Lab上面，如果想要系统学习LLM的训练细节，建议直接看Lab的内容，Lecture部分可以作为辅助材料来理解一些概念和背景知识。

2026/04/09：Lecture4之后因为对并行训练的兴趣，直接跳到了Lecture7-8，感觉内容非常有价值，值得细看一下lecture视频或者讲义，精华在于不同并行方式的介绍和分析，在Lecture8里有一些简单的通信操作代码示例，帮助理解。

CS336 Lecture 1: Overview and Tokenization

Fri, 03 Apr 2026 00:00:00 +0000

课程导读

这门课的讲义基于 notebook 组织，不同章节通过函数模块化展开。

开场老师提出了一个很尖锐的问题：研究者正在与底层技术逐步脱节。他给了一个时间线：

八年前，研究者会自己补充数据并训练模型；
六年前，研究者还会下载模型后进行 fine-tune；
现在，很多人直接向闭源模型（GPT-4/Claude/Gemini）提问。

这虽然有些夸张，但确实反映了趋势：模型能力提升后，很多基础环节被“封装”了。老师强调：

“Full understanding of this technology is necessary for fundamental research”

随着模型规模持续增大，训练中的计算压力正在从 Attention 逐步转向 FFN。

课程收获

Mechanics：系统如何工作（如 Transformer 结构、GPU 并行方式）
Mindset：如何最大化利用硬件与内存，如何看待 scaling laws
Intuition：哪些数据和建模决策更可能带来更好结果

老师也指出一个常见误解：只要堆算力，模型就会自动变好。他给出的观点是：准确率 = 效率 × 资源。

在数据和资源固定时，效率决定了上限。

老师将近年的 LLM 工作大致分成两类：

闭源路线：以 GPT 系列为代表，最早系统性拥抱 Scale，但细节封闭。
开放权重路线：以 Qwen 等 open-weight 模型为代表，拥抱 Scale 的同时提升开放度，但常见情况是只公开部分训练细节，数据与失败案例仍不完整。

Tokenization

为什么需要 tokenization？因为 LLM 在 token 序列上建模概率分布，我们需要把原始字符串编码为 token，并保证可逆解码。

🔗 Tokenization 可视化工具: tiktokenizer.vercel.app

常见 tokenization 方法：