Notes

TL;DR 本讲将分布式通信概念落地为可执行代码，通过 PyTorch 和 NCCL 展示了集合通信的实际实现，并手写了三种并行策略的简化版本。从 All-Reduce 的基准测试到 MLP 的切分实现，揭示了通信开销与计算模式的核心逻辑。 ...

分布式通信原语 All reduce：根进程 $P_r$ 拥有一个大小为 $K$ 的消息缓冲区 $B$。操作结束后，通信组内的所有进程 $P_i$ 都拥有了 $B$ 的完全一致的副本。 Scatter：根进程 $P_r$ 拥有一个大向量 $V$，将其均匀切分为 $N$ 个块，记为 $v_0, v_1, \dots, v_{N-1}$。操作结束后，$P_i$ 仅获得对应的分块 $v_i$。 Gather：Scatter 的逆操作。每个进程 $P_i$ 拥有一个数据块 $v_i$。操作结束后，根进程 $P_r$ 将这些块按 Rank 顺序拼接，在其内存中形成完整向量 $V = [v_0, v_1, \dots, v_{N-1}]$。 All-Gather：每个进程 $P_i$ 拥有数据块 $v_i$。操作相当于先执行 Gather，再执行 Broadcast。操作结束后，每一个进程 $P_i$ 都拥有完整的向量 $V = [v_0, v_1, \dots, v_{N-1}]$。 Reduce：每个进程 $P_i$ 拥有一个向量 $X_i$。操作结束后，根进程 $P_r$ 获得结果 $R = X_0 \oplus X_1 \oplus \dots \oplus X_{N-1}$。 All-Reduce：逻辑上等价于 Reduce + Broadcast。操作结束后，所有进程 $P_i$ 都获得了完全相同的规约结果 $R = \sum_{j=0}^{N-1} X_j$。 Reduce-Scatter：逻辑上等价于先对全局数据做 Reduce 得到结果 $R$，然后将 $R$ Scatter 给各个进程。操作结束后，$P_i$ 获得结果向量的第 $i$ 个分块。 LLM 训练中不同的并行策略 朴素数据并行 DP 数据并行中每个 GPU 都有一个完整的模型副本，输入 Batch 会被切分为多个 Mini Batch 并喂给不同的 GPU，每个 GPU 正常计算出梯度后，通过 All-Reduce 计算出全局梯度，再对参数进行更新。 ...

TL;DR 本讲是 CS336 系列笔记的第六讲。本节从算子融合的必要性出发，横向对比了手写 CUDA、使用 Triton 以及 PyTorch 2.0 编译技术三种实现方式。重点解析了 Triton 如何通过“块级（Block-level）”抽象简化显存管理，并以 GeLU、Softmax 和 Matmul 为例，展示了利用共享内存（SRAM）和分块（Tiling）技术打破访存墙的关键技巧。 ...

TL;DR 本讲是 CS336 系列笔记的第五讲。本讲从 GPU 的设计理念出发，梳理了 SM/SP 计算架构与内存层级体系，并结合 Roofline 模型，重点解析了访存合并、算子融合、重计算及分块（Tiling）等核心优化策略。此外，还应用上述策略，粗略推导了 FlashAttention 如何利用 Online Softmax 技术打破显存带宽瓶颈，实现 IO 感知的极致加速。 ...

TL;DR 本讲 CS336 系列笔记的第四讲。本讲梳理了 MoE 架构利用稀疏激活实现“高效扩参”的核心机制，并结合 DeepSeek 系列模型的演进路线，重点解析了细粒度专家、共享专家及无辅助损失负载均衡等策略，如何解决大规模训练中的路由坍塌与稳定性难题。 ...

TL;DR 本讲梳理了现代 LLM 架构设计的“事实标准”（Post-Norm + RMSNorm + RoPE），并从系统视角解析了 GQA 与滑动窗口机制如何通过优化 KV Cache 访存，解决推理阶段的算术强度劣化问题。 ...

本文记录了 CS336 Lab 1 的实验笔记，整个 Lab 的工作量很大，主要内容包括从头实现 BPE Tokenizer、一系列算子和基于 Transformer 的语言模型，并在此基础上进行大量的调优和消融实验。做完这个 Lab 可以对分词器的实现细节有高细粒度的理解，也能积攒对 Transformer 模型各组件的直观认识。 ...

本文是 CS336 系列笔记的第二讲，PyTorch 与资源计算。本讲比较简单，从两个问题出发，讨论如何计算计算量，并顺带讲一些 PyTorch 的入门概念。本讲的重点不在于 PyTorch 基础，而在于培养“资源计算”的思维模式。 ...

本文是 CS336 第一讲的学习笔记，主要介绍开设这门课程的背景和动机，并对课程主要内容做了概览。还介绍了不同的分词器基本原理及其优缺点。 引入 为什么要学习这门课程 现象：研究人员与底层技术越来越远。八年前，他们需要自己实现和训练模型；六年前，他们需要下载一个模型，并进行微调；现如今，他们仅仅修改模型的提示词。 ...

本文为 EfficientML Lab 2 实验笔记，包含 K-Means 量化、K-Means QAT、线性量化等内容，难度不大，内容丰富。 Part 1: K-Means Quantization Qustion 1 第一个问题是实现 K-means 量化的核心算法，其中 K-means 本身是调库实现的。 ...