并行训练系列：5. Megatron 之分布式环境初始化
训练脚本入口点： pretrain_gpt.py - GPT模型训练 pretrain_bert.py - BERT模型训练 pretrain_mamba.py - Mamba模型训练 pretrain_t5.py - T5模型训练 pretrain_vlm.py - 视觉语言模型训练 train_rl.py - 强化学习训练 例如 pretrain_gpt.py:if name == "main": -> 调用 megatron.training.pretrain 函数 123456789101112from megatron.training import pretrainif __name__ == "__main__": ...... pretrain( train_valid_test_datasets_provider, partial(model_provider, gpt_builder), ModelType.encoder_or_decoder, ...
并行训练系列：4. 张量并行（TP）
在 ZeRO 系列中，实现了 parameters, gradients 和 optimizer states 的切分，然而在 FWD/BWD 计算时，使用的依然是完整的模型参数（需要预先通过 All-Gather 收集）。能不能在计算时只使用模型参数的一部分呢？ 由此引出张量并行（Tensor Parallelism）。 TP 切分方式 在神经网络中，矩阵乘法常用以下形式表示：\(X\times W\)，其中\(X\)为 actication 的输入，\(W\)为模型参数，均以 tensors 的形式保存；tensors 将被沿着一个特定维度切分为 N 个 shards，分布到 N 个 GPU 上。矩阵可以按行或者按列切分，分别对应行并行、列并行。 row-linear forward 运用 scatter 操作，将输入矩阵切分为若干个列； 将每个权重矩阵切分为若干行，分别与输入矩阵的各列相乘，最后通过 all-reduce 操作相加。 backward f和g分别表示两个算子（每个算子包含一个 forward+backward 操作） g的...
并行训练系列：3. 数据并行下篇（DeepSeed-ZeRO）
上篇描述了朴素数据并行（DP）和分布式数据并行（DDP）策略。两者的总通信量相同，然而 DP 的单 Server 机制导致了通信负载不均的问题（Server 通信量和总 GPU 数量成线性关系），导致 DP 适用于单机多卡场景；DDP 通过 Ring-AllReduce 这一 NCCL 操作，使得通信量均衡分布到所有 GPU 上（单个 GPU 通信量不受到总 GPU 数量的影响），利于多机多卡场景。 解决通信量负载不均的问题后，开始考虑显存开销：每个 DP 分片都复制了一份完整的模型参数、梯度和优化器状态，造成极大的显存负担，如何优化呢？ 在优化显存前，需要先理清两个关键问题： 训练过程中 GPU 上存储了哪些内容？ 分别在哪些阶段需要使用？ 先看第一个问题。 显存开销计算 当训练一个神经网络时，需要在内存中存储以下内容：模型参数 W、模型梯度...
并行训练系列：2. 数据并行上篇（DP，DDP）
数据并行（DP） 思想：将模型复制到多个GPU上；在每个GPU上，对不同的micro batches执行前向/反向传播。 准备：在每个 GPU 上都拷贝一份完整的模型参数，将一个数据 batch 均分为多个 micro-batch，分别拷贝到不同 GPU 上（DP 分片）； 前向传播/反向传播：每个 GPU 上对不同的 micro-batch 分别执行前向/反向传播，计算得到梯度 G； 聚合并下发梯度（All-Reduce）： 每个 GPU 将本地梯度 G 发送给单个 GPU 做聚合（或者在不同的 GPU 上聚合梯度的某些部分）；再将平均后的梯度下发到所有 GPU； 本地参数更新：所有 GPU 利用下发的梯度完成本地参数更新。 实现 DP 的经典框架为参数服务器：计算 GPU 称为 Worker，梯度聚合 GPU 称为 Server。 Worker 和 Server 是逻辑概念，分别可以对应一块或多块物理意义上的 GPU（当前讨论的基本都是对应一张 GPU...
RL 系列：2. 从 Bellman 算子的角度解释策略迭代/价值迭代
看完策略迭代和价值迭代的定义后，我还是有点迷糊：为什么得到的策略一定会收敛呢？两种方法的本质区别是什么呢？ 于是搜罗了一些更“数学”的解释，来阐明这个问题。 Bellman 算子 定义状态空间：\(S={\{s_1, s_2, ..., s_n \}}\)；动作空间：\(A={\{a_1, a_2, ..., a_n \}}\) 贝尔曼期望方程： \[ V(s)=E_{a\sim\pi(a|s)}[r(s,a)+\gamma E_{s'\sim p(s'|s,a)}[V(s')]] \] 贝尔曼最优方程： \[ V(s)=\max_a(r(s,a)+\gamma E_{s'\sim p(s'|s,a)}[V(s')]) \] 根据以上方程，定义贝尔曼期望算子\(B_\pi\)和贝尔曼最优算子\(B_*\)，分别代表对当前价值函数集\(V\)利用贝尔曼方程更新的操作： \[ B_\pi V(s)=E_{a\sim\pi(a|s)}[r(s,a)+\gamma...
RL 系列：1. Markov 决策过程
本文旨在阐述基于 Markov 决策过程的强化学习原理，聚焦于关键公式（省去一些繁杂的数学推导），快速领会 RL 的核心理念。 RL 概述 强化学习的理念：Agent 根据 Environment 输出的 \(S_t\) 和 \(R_t\) 综合判断，产生当前步骤的动作 \(A_t\)，放入 Environment 执行；Environment 根据 \(A_t\) 执行的结果，输出下一个状态 \(S_{t+1}\) 和当前动作带来的奖励 \(R_{t+1}\)，返回给 Agent。 强化学习和监督学习的区别： RL 得到的时序数据不满足独立同分布； RL 无实时明确反馈当前 Action 是否正确，奖励信号延迟。（通过 rollout 过程从当前帧对 Action 进行采样，得到多个 observation， 每个 observe 对应一条 trajectory，即一个 State 和 Action 的序列 τ=(s0,a0,s1,a1,…)；通过最终 reward 来观测当前 trajectory 是否合适） Markov...
并发系列：2. Go 的长连接
TCP 连接的建立和关闭 绝大多数网络连接的建立都是基于 TCP 协议的，我们往往知道一个原则：建立 TCP 连接需要三次握手，其具体过程也是面试的一个常考点。那么“为什么 TCP 建立连接需要三次握手？”呢？这个问题很少深究。首先回顾一下建立连接的过程： TCP 连接是什么？ 连接：用于保证可靠性和流控制的信息，包括 Socket、序列号和窗口大小。其中：Socket 由互联网地址标志符和端口组成；窗口大小主要用来做流控制；最后的序列号用于追踪通信发起方发送的数据包序号，接收方可以通过序列号向发送方确认某个数据包的成功接收。 TCP 所有的协议状态如下图： 建立 TCP 连接的三次握手 抽象成通俗语言就是： 你能听到吗？ 我能听到，你听得到吗？ 我也能听到。 第一次握手：SYN 报文：客户端随机初始化序列号 client_isn，放进TCP⾸部序列号段，然后把SYN置1。把SYN报⽂发送给服务端，表示 发起连接，之后客户端处于...
Go-内存分配器
内存分配器 内存空间包含两个重要区域：栈区和堆区。不同编程语言使用不同的方法管理堆区的内存，C++ 等编程语言会由工程师主动申请和释放内存，Go 以及 Java 等编程语言会由工程师和编译器共同管理，堆中的对象由内存分配器分配并由垃圾收集器回收。 设计原理 内存管理一般包含三个不同的组件，分别是用户程序（Mutator）、分配器（Allocator）和收集器（Collector）；当用户程序申请内存时，它会通过内存分配器申请新内存，而分配器会负责从堆中初始化相应的内存区域。 分配方法 线性分配器 使用线性分配器时，只需要在内存中维护一个指向特定位置的指针；用户程序申请内存时，移动指针。然而，这不利于在内存被释放时重用内存。如下图中红色部分难以利用： 因为线性分配器需要与具有拷贝特性的垃圾回收算法配合，所以 C 和...
Go-调度器
Go-调度器 多个线程可以属于同一个进程并共享内存空间。因此它们也不需要内存管理单元处理上下文的切换，线程之间的通信是基于共享的内存进行的。 虽然线程比较轻量，但是在调度时也有比较大的额外开销。每个线程会都占用 1M 以上的内存空间，在切换线程时不止会消耗较多的内存，恢复寄存器中的内容还需要向操作系统申请或者销毁资源，每一次线程上下文的切换都需要消耗 ~1us 左右的时间1，但是 Go 调度器对 Goroutine 的上下文切换约为 ~0.2us，减少了 80% 的额外开销。 Go 语言的调度器通过使用与 CPU 数量相等的线程减少线程频繁切换的内存开销，同时在每一个线程上执行额外开销更低的 Goroutine 来降低操作系统和硬件的负载。 设计原理 包括以下几个版本： 单线程调度器：程序中只能存在一个活跃线程，由G-M模型组成； 单线程调度器 0.x...
并发系列：1. 自旋锁
首先从乐观锁和悲观锁的设计思想开始。 乐观锁和悲观锁 悲观锁 悲观锁总认为最坏的情况可能会出现，即数据很可能会被其他人所修改，所以悲观锁在持有数据的时候总会把资源或者数据锁住，其他线程想要请求这个资源时阻塞，直到悲观锁释放资源（读写操作均加锁）。传统的关系型数据库里边就用到了很多悲观锁机制，比如行锁，表锁等，读锁，写锁等，均在做操作之前先上锁。悲观锁的实现往往依靠数据库本身的锁功能实现。 Java 中的 Synchronized 和 ReentrantLock 等独占锁(排他锁)也是一种悲观锁思想的实现，因为 Synchronzied 和 ReetrantLock 不管是否持有资源，它都会尝试去加锁。 一个悲观锁的运用场景： select * from student where name="cxuan" for update 这条 sql 语句从 Student 表中选取 name =...
