[megatron代码阅读] 2. TP和PP实现
训练并行实现#
TensorParallel#
张量并行代码路径, 代码路径: megatron/core/tensor_parallel
主要包含Linear / VocabEmbedding / cross_entropy 三部分.
Linear#
参数初始化
如果是从checkpoint热启, perform_initialization需要打开这个配置
1.set_tensor_model_parallel_attributes
: 设置weight的三个属性: is_parallel / partition_dim / stride
2.调用传入的init_method, 初始化weight. 这里注意要使用同一个随机种子. 如果是expert网络, 每个expert要用自己独立的rng_tracker
3.如果启用expert_parallel, 设置allreduce属性为false. 否则为true
列并行
sequence_parallel回忆, 为了节省显存,拆分了layernorm后的激活存储. 在进入TP前通过allGather获取到完整的激活, 经过TP后再通过reduceScatter分离到各张卡.
grad_accumulation_fusion
代码: https://github.com/NVIDIA/apex/blob/master/csrc/megatron/fused_weight_gradient_dense_cuda.cu
主要作用是在显存受限,无法一次性更新大batch_size的时候, 通过mini-batch来累加多个小批量的梯度到weight.main_grad
, 这里fused意思就是在main_grad上原地更新, 最后在用main_grad来更新大batch里的weight.
LinearWithGradAccumulationAndAsyncCommunication
forward:
- 输入: 如果开sp, torch.distributed.all_gather_into_tensor input
- Matmul(input, weight) + bias
backward:
- wgrad_deferral_limit: TODO: 没弄懂用于降低pipeline flush延迟的含义
- 如果开sp, grad_input bp对应的集合通信操作是dist_reduce_scatter_func
- 没开sp, grad_input bp对应的是all_reduce_func
- 如果开了grad_accumulation_fusion & sp, 需要先all_gather input, 也就是X, 因为求grad_weight的时候需要matmul input的转置.
- 这里的异步优化指的是先进行 grad_input的异步集合通信, 在此同时计算grad_weight, 算完grad_weight后再等grad_input通信完成, 这样就能overlap一部分通信耗时.
行并行
forward
- 开sp, input不做处理,因为直接输入切分好的input, 经过线性层后ouput reduceScatter到对应的节点
- 关sp, input需要先进行ReduceScatter对输入x做切分, 经过线性层后output allReduce结果.
Backward: 没有集合通信
PipelineParallel#
核心配置参数有两个:
-
pipeline_model_parallel_size
: pp切分数, transformer_layer实际被切分为多少个group -
virtual_pipeline_model_parallel_size
: 举例 tensor_model_parallel_size=1, pipeline_model_parallel_size=4, virtual_pipeline_model_parallel_size=2, 一共有16个transform_layer的情况下, 模型被切分为:GPU 0: [1, 2] [9, 10] GPU 1: [3, 4] [11, 12] GPU 2: [5, 6] [13, 14] GPU 3: [7, 8] [15, 16]
一共8个stage, 每个stage有2个layer. PP原理回忆
PP代码逻辑位置 megatron/core/pipeline_parallel
train_step->get_forward_backward_func->forward_backward_pipelining_with_interleaving
![image-20250126160014500](https://img2023.cnblogs.com/blog/1439743/202501/1439743-20250126160024948-212841132.png)
P2P通信#
有两种方式 batch_isend_irecv
与 _p2p_ops
, 后者即send和recv独立作为一个通信操作
batch_isend_irecv
: 将send_prev/recv_prev/send_next/recv_next可以异步并发执行
p2p通信步骤:
- 传输tensor_shape, int64类型 (类似sequence压缩通信方式, 先传长度)
- 对所有的pp group进行遍历, 如果需要recv_prev / recv_next, 先创建空tensor用于结果存储 (这里是否能优化)
- 根据是否batch传输, 分别进行并行/串行的方式通信.
- 等待通信完成, 进行cuda流同步
1F1B(非交错式)#
![image-20250126165021086](https://img2023.cnblogs.com/blog/1439743/202501/1439743-20250126165027410-1287729087.png)
缺点: 无法支持 p2p通信耗时的overlap
Warmup
num_warmup_microbatchs = min(microbatch, pp_world_size - pp_rank - 1), 比如device1的warmup就是 4 - 0 - 1 = 3, 前3个microbatch warmup的时候, 整体pipeline处于串行的执行状态.
步骤: recv_forward->forward_step->send_forward 再到下一层PP, 直到warmup步骤全部走完.
Steady
在稳态状态下就是1F1B描述的情况. 交替进行fp和bp
以device3刚进入steady状态为例:
forward_step
: warmup执行了microbatch1, steady执行的第一个forward是 batch2send_forward_recv_backward
: 向device4发batch2的fp结果, 同时等device4返回batch1的bp结果. 这里是同步通信, 需要等bp执行完成, 这时候并没有跑到batch3的fp上.backward_step
: 执行batch1的bpsend_backward_recv_forward
: 把batch1的bp结果发给device2, 同时接受device2的batch3 fp结果, 用来执行下一轮的batch3 fp.
5,6,7,8 图上描述的状态和代码是完全一致的, 但1,2,3,4不完全一致.
Cooldown
和warmup刚好是相反的逻辑.根据warmup microbatchs的个数, 等待bp执行完成.
1F1B with interleaving#
虚拟流水线的主要目的是让microbatch_size更小更多, 从而减少气泡。方法是让一个device虚拟成
初始化
- warmup_batch数计算方法, 如下代码:
total_num_microbatches = num_microbatches * num_model_chunks #模型分块数(virtual pipeline size) * microbatch
all_warmup_microbatches = False
if forward_only:
num_warmup_microbatches = total_num_microbatches
else:
# 这里*2的原因是 为了充分利用设备资源,会使用双倍缓冲技术。这意味着每个设备会同时处理两个microbatches,一个在前向传播,另一个在后向传播。因此,热身阶段的microbatches数量需要乘以2,以覆盖前向和后向传播。
num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
# microbatch_group_size_per_vp_stage 默认值 = pipeline_parallel_size, 用于
num_warmup_microbatches += (num_model_chunks - 1) * config.microbatch_group_size_per_vp_stage
if num_warmup_microbatches >= total_num_microbatches:
num_warmup_microbatches = total_num_microbatches
all_warmup_microbatches = True
num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches
- 设置schedule_table, 为了方便计算, 将microbatch+chunk重映射成了virtual_microbatch_id
# PP2 N3M5 with VP2 is constructed as below:
# virtual_microbatch_id | 0 1 2 3 4 5 6 7 8 9
# microbatch_id | 0 1 2 0 1 2 3 4 3 4
# model_chunk_id | 0 0 0 1 1 1 0 0 1 1
根据chunk_id, 还能判断出这个virtual_id是这个device上的第一个chunk还是最后一个chunk
recv_tensor_from_previous_stage
: 先判断当前stage是否为leading stage(forward第一个, backward最后一个), 如果virtual_microbatch_id < (pipeline_parallel_size - 1)
, 说明当前stage没有任何前置需要接受的tensor, 否则说明他和之前的最后一个stage连在一起. 以PP=4举个例子:
# 0 1 2 3 ... 这里的microbatch 0的下一个stage是 device0的microbatch3
# 0 1 2 3 ...
# 0 1 2 3 ...
# 0 1 2 3 ...
warmup
注意配置项: overlap_p2p_comm_warmup_flush
: 在打开这个开关后支持overlap warmup和flush阶段前向计算和通信, 后面看代码默认这个开关打开, warmup步骤:
- 根据microbatch id判断是不是leading_stage, 如果不是的话需要等上一个循环发出的异步接受前向结果的handle.
- 异步通信结果保存在fwd_recv_buffer, 异步发出预取下个循环的recv_forward请求
- 进行该stage的forward_step
- 把output_tensor 异步发出 send_forward, 等上一个循环的send_next_wait_handle完成.
- 把通信完的fwd_recv_buffer 赋值给input_tensor用于下个循环的forward
- 在warmup的最后, 触发异步等待recv_backward的请求. 方便衔接steady阶段
steady
循环num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches
次, 步骤:
- 等warmup的recv_prev 异步执行完, 收到forward结果到buffer里
- Forward_step
send_forward_recv_forward
: 同时接受previous stage的forward结果, 同时把next stage的forward输入发出.- Wait recv_next传回来的grad
- Backward_step
send_backward_recv_backward
: 反向往之前的stage发grad- 等上一个batch的backward send_prev发完, 相当于一个buffer切换过程.
整个流程像下面这个流水线示意图.
![image-20250206201325086](https://img2023.cnblogs.com/blog/1439743/202502/1439743-20250206202215958-2146954135.png)
cooldown
与warmup刚好完全相反, 只有backward的计算和通信操作.
注意在每个阶段完成的时候都回将通信用到的output_tensor重新释放回显存池, 用来缓解显存压力.
参考:#
对VPP的进一步优化: https://zhuanlan.zhihu.com/p/681363624
作者:sunstrikes
出处:https://www.cnblogs.com/sunstrikes/p/18701639
版权:本作品采用「署名-非商业性使用-相同方式共享 4.0 国际」许可协议进行许可。
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· DeepSeek-R1本地部署如何选择适合你的版本?看这里
· 开源的 DeepSeek-R1「GitHub 热点速览」
· 传国玉玺易主,ai.com竟然跳转到国产AI
· 揭秘 Sdcb Chats 如何解析 DeepSeek-R1 思维链
· 自己如何在本地电脑从零搭建DeepSeek!手把手教学,快来看看! (建议收藏)