原文：https://zhuanlan.zhihu.com/p/10091011992
为什么会有这篇文章：虽然工作内容不是infra，但是我比较喜欢研究训练方法，魔改训练框架造轮子。正好最近看到OpenRLHF用ray管理VLLM的方案，感觉很有意思，遂研究了一下，发现VLLM的TP切分和Megatron是一套逻辑，用torch的rpc也可以代替ray的远程调用，所以打算用Megatron+TorchRPC+VLLM实现一套类似的框架，后期再把VLLM原地换掉直接megatron推理。在开始这个大工程之前，正好有机会写下这篇文章，就算是开工仪式了。
本文的主要内容
本文主要是从编程的角度，对LLM训练框架所涉及的一些前置编程知识进行讲解，并且会举一些应用技巧，对应到当前的LLM训练框架，辅助理解训练框架的代码逻辑。举个例子，下面是一段megatron初始化多卡通信组的代码：
rank = torch.distributed.get_rank()
for ranks in rank_generator.get_ranks('dp'): # 迭代生成所有数据并行ranks的列表
group = torch.distributed.new_group(
ranks, timeout=timeout, pg_options=get_nccl_options('dp', nccl_comm_cfgs)
) # 根据数据并行ranks创建通信组
group_gloo = torch.distributed.new_group(ranks, timeout=timeout, backend="gloo")
if rank in ranks: # 如果当前rank属于这个数据并行ranks，则保存创建的通信组
_DATA_PARALLEL_GROUP = group
_DATA_PARALLEL_GROUP_GLOO = group_gloo
_DATA_PARALLEL_GLOBAL_RANKS = ranks
读完本文你会了解：
-
什么是rank、什么是world_size，什么是通信组group
-
为什么这段代码是先建通信组，再根据rank in ranks决定是否保存？判断rank in ranks 了再创建group行不行？
-
在子进程和子线程中创建通信组有什么区别和要注意的地方。
-
backend="gloo"：什么是backend，gloo backend 是干什么的，为什么不用nccl backend。
再比如下面是一段deepspeed在参数上注册的回调函数：
def create_reduce_and_remove_grad_hooks(self):
self.grad_accs = []
for i, param_group in enumerate(self.bit16_groups): # 遍历混精优化器的参数
for param in param_group:
if param.requires_grad:
def wrapper(param, i):
param_tmp = param.expand_as(param) # 在原始的参数上建立一个视图
grad_acc = param_tmp.grad_fn.next_functions[0][0] # 获得这个这个视图的前一个算子
def reduce_partition_and_remove_grads(*notneeded):
self.reduce_ready_partitions_and_remove_grads(param, i)
self._grad_acc_hooks.append(grad_acc.register_hook(reduce_partition_and_remove_grads)) # 在视图的前一个算子上注册回调函数
self.grad_accs.append(grad_acc)
wrapper(param, i)
读完本文你会了解到：
-
register_hook是干什么，什么时候触发。
-
通常都是直接在参数上注册hook，为什么这里要在参数的视图的前一个算子注册hook：param.expand_as(param).grad_fn.next_functions[0][0]
下面开始正文。
训练框架是干什么的
目前LLM训练有两大主流框架：Deepspeed和Megatron-LM。前者的主要提出和维护者是微软的工程师，后者是英伟达的工程师。两个框架从底层原理到设计语言可以说是大相径庭。训练框架的主要目标有2：一是在有限的GPU中尽可能地塞入一个大号模型，二是高效地利用多GPU进行训练。完成第一个目标主要依赖的是模型切分，或者更笼统地说是降低单卡显存占用。完成第二个目标依赖的是异步、高重叠度、高带宽的数据通信。Deepspeed在降低显存这方面应用的技术主要有Zero-1、2、3， 序列并行、CPU Offload，高效通信方面主要是依赖register_hook回调函数的异步通信、多cuda事件流、GPU计算重叠、连续锁页缓存。Megatron在降低显存方面的主要技术有Distributed Optimizer、Tensor Model Parallel、Pipeline Model Parallel、序列并行，在高效通信方面主要的技术有P2P通信、重叠流水线并行、梯度缓存。在设计语言上，DeepSpeed属于外挂框架，框架并不介入模型前向的计算图，因此对模型结构一般没有特殊要求，核心代码通过大量回调函数和torch派生类封装，并不暴露给用户。Megatron则是属于内嵌框架，直接改变模型计算图，因此限制模型结构必须是类Transformer的结构，代码全部暴露在外，不怎么依赖回调函数。外挂框架的好处是兼容性好，对新手友好，缺点是启动慢、计算速度略低，适合数据量不大、不关注训练加速技术，专注于模型和数据迭代的人。内嵌框架的优点是启动、训练效率高一点，对于想魔改底层的人更友好，但是对于想轻微修改训练逻辑的不太友好，适合大规模训练和想要了解、改动并行训练代码的人。
Hello World
我们以hello world来开场：
import torch
import os
if __name__ == '__main__':
rank = int(os.getenv('RANK','0'))
world_size = int(os.getenv('WORLD_SIZE','1'))
torch.distributed.init_process_group(rank=rank, world_size=world_size, backend='nccl')
print(f'Hello World from rank:{torch.distributed.get_rank()} out of {torch.distributed.get_world_size()} worlds.')
假设代码命名为t1.py，则可以用 torchrun --nproc-per-node 2 t1.py 来启动。（如果不做特殊说明，后面都用这组参数启动。后面给出的demo大多数没有GPU也可以跑，只需要将张量的创建位置改为cpu，以及通信相关的后端全部使用gloo后端即可）
正常情况下，可以看到程序打印出
下面对代码稍作解释。
启动命令
torchrun是安装torch后自带的命令，作用是帮助我们以多进程方式启动指定脚本。
在分布式环境中，我们首先要区分多机和多卡，多机是指多台运行训练任务的服务器，多卡是指一台机器上有多个显卡。我们一般称一台服务器为node，有几台服务器就有几个node。
那么 --nproc-per-node 从字面意思上就能看出表示每台服务器上启动多少个进程。一般来说，我们启动的进程数与服务器上的gpu数相同，也就是8卡机 nproc-per-node 设为8。当然一台8卡机你可以只启动2个进程，也可以启动100个进程，这个并不是硬限制。但是保持进程数和卡数相同可以简化我们的编程逻辑，让每个进程只负责一张卡上的计算。
除了nproc-per-node，启动命令还有几个常见参数：
-
--master_addr和--master_port：当启动的是多机环境时，需要用这两个参数指定主机的IP地址和端口号。这两个参数在所有master和slaver机器上都是一样的，不需要修改。
-
--nnodes：当多机启动分布式时，使用这个参数，表明总共有多少台机器。master会根据这个参数配置来等待slaver链接，直到slaver数量凑够nnodes才会启动。这个参数在每台机器上都是一样的，不需要修改。
-
--node_rank：表示当前node是全部node中的第几个，这个参数需要根据启动的机器进行更改，master的rank必须是0。
上述启动参数要放在torchrun之后，脚本路径之前。写在脚本路径之前的参数不会被传递给脚本运行的环境，所以脚本中不要处理这些参数。
rank & world size
在启动的训练脚本中，需要先初始化分布式环境，也就是下面这几行：
rank = int(os.getenv('RANK','0'))
world_size = int(os.getenv('WORLD_SIZE','1'))
torch.distributed.init_process_group(rank=rank, world_size=world_size, backend='nccl')
如果是使用torchrun启动的脚本，torch会帮你在环境变量中写入RANK和WORLD_SIZE这两个变量，在脚本中可以读出来。它们分别表示当前进程的序号和总进程数。这个进程序号是全局的进程序号，也就是说如果每台服务器启动8个进程，总共10台服务器参与训练，那么第十台机器的8个进程rank分别是72、73、74、75、76、77、78、79。world_size也是全局的进程数，而不是单台服务器上的进程数。
获得rank和world_size后，就可以用torch.distributed.init_process_group初始化分布式环境。
在初始化之后的任何地方，都可以用torch.distributed.get_rank()和torch.distributed.get_world_size()来获取当前进程序号和总进程数。
另外需要提的一点是，torch并没有local_rank 这个概念。这是一些训练框架自己定义的，通常用来表示当前进程是这台服务器上的第几个进程。
后端 backend
在初始化命令中，我们还指定了backend参数，这个参数表示的是分布式环境默认使用的通信后端是什么，一般可以选择 nccl、gloo和mpi后端。gpu to gpu的通信选择nccl后端。cpu to cpu的通信选择gloo后端，一般不太会用到mpi后端。
nccl通信会比gloo通信快，所以应该尽量使用nccl进行通信。但是在有些时候，比如读取数据的阶段，如果多个进程之间需要通信，一般使用用gloo通信。因为这时，我们并不希望这些还没有开始正向传播的张量过早出现在gpu上，避免过多占用显存。在后面讲group的部分，我们会提到怎么让训练进程使用nccl后端，数据读取进程使用gloo后端。
训练脚本开始阶段的一些小细节
set device
分布式训练中，在我们创建gpu张量时，经常使用torch.cuda.current_device()获取当前rank使用的cuda设备，然后直接在这个设备上创建。使用这个函数的前提是之前通过set_device显式设定过默认设备。另外还有一些算子会要求用户必须执行过set_device。
所以一个好习惯是，在执行主要的训练逻辑之前，尽早设置默认的cuda设备：
devices = torch.cuda.device_count()
torch.cuda.set_device(rank % devices)
如果不手动设置，current_device默认会返回cuda:0。
固定随机种子
分布式训练时随机种子是一个容易忽略的事，但却是很重要的事。在脚本初始化阶段设置一个固定的随机种子有两个好处，一个是让实验可复现，另一个就是让不同rank的模型以相同的方式初始化，不然还要用通信操作同步一下各个模型来保证模型初始化状态的一致性。
def set_random_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
通信算子
通信算子指的是在分布式训练中，能够在不同rank之间进行数据交换的算子。torch提供了很多通信相关的算子，最常用可以划分为5类：规约、聚集、广播、点对点通信和其他。
规约
规约最典型的就是torch.distributed.all_reduce，这个算子能够对不同rank之间的数据进行求和、求均值、最大值等操作，特点是不论有多少个rank，最终结果只有一个tensor。例如下面这段代码，作用是对所有rank之间的"tensor"求和：
import torch
import os
if __name__ == '__main__':
rank = int(os.getenv('RANK','0'))
world_size = int(os.getenv('WORLD_SIZE','1'))
torch.distributed.init_process_group(rank=rank, world_size=world_size, backend='nccl')
devices = torch.cuda.device_count()
torch.cuda.set_device(rank % devices)
tensor = torch.tensor([rank +。1], dtype=torch.long, device='cuda')
torch.distributed.all_reduce(tensor)
print(f'rank:{rank} {tensor}')
rank0的tensor=1，rank1的tensor=2，求和结果为3，因此能看到这样的输出：
最低0.47元/天 解锁文章
扫一扫
6847
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
