[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#

核心配置参数有两个:

  1. pipeline_model_parallel_size: pp切分数, transformer_layer实际被切分为多少个group

  2. 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

P2P通信#

有两种方式 batch_isend_irecv_p2p_ops, 后者即send和recv独立作为一个通信操作

batch_isend_irecv: 将send_prev/recv_prev/send_next/recv_next可以异步并发执行

p2p通信步骤:

  1. 传输tensor_shape, int64类型 (类似sequence压缩通信方式, 先传长度)
  2. 对所有的pp group进行遍历, 如果需要recv_prev / recv_next, 先创建空tensor用于结果存储 (这里是否能优化)
  3. 根据是否batch传输, 分别进行并行/串行的方式通信.
  4. 等待通信完成, 进行cuda流同步

1F1B(非交错式)#

image-20250126165021086

缺点: 无法支持 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状态为例:

  1. forward_step: warmup执行了microbatch1, steady执行的第一个forward是 batch2
  2. send_forward_recv_backward: 向device4发batch2的fp结果, 同时等device4返回batch1的bp结果. 这里是同步通信, 需要等bp执行完成, 这时候并没有跑到batch3的fp上.
  3. backward_step: 执行batch1的bp
  4. send_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虚拟成 v 个device,从计算1个连续的layer段(有 x 个 layer)变成计算 v 个不连续的layer段(每段 layer 数量为 x/v). 比如之前1F1B时device1负责layer 1~4,device2负责 5~8,在 Interleaved 1F1B下device1负责layer 1~2 和 9~10,device2负责 3~4 和 11~12,这样可以让流水线中每个stage更小,因而下个stage的等待时间更短,气泡更小。需要注意的是,micro_batch_size需要是 pipeline_parallel_size的整数倍。

初始化
  1. 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

  1. 设置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

  1. 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步骤:

  1. 根据microbatch id判断是不是leading_stage, 如果不是的话需要等上一个循环发出的异步接受前向结果的handle.
  2. 异步通信结果保存在fwd_recv_buffer, 异步发出预取下个循环的recv_forward请求
  3. 进行该stage的forward_step
  4. 把output_tensor 异步发出 send_forward, 等上一个循环的send_next_wait_handle完成.
  5. 把通信完的fwd_recv_buffer 赋值给input_tensor用于下个循环的forward
  6. 在warmup的最后, 触发异步等待recv_backward的请求. 方便衔接steady阶段
steady

循环num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches次, 步骤:

  1. 等warmup的recv_prev 异步执行完, 收到forward结果到buffer里
  2. Forward_step
  3. send_forward_recv_forward: 同时接受previous stage的forward结果, 同时把next stage的forward输入发出.
  4. Wait recv_next传回来的grad
  5. Backward_step
  6. send_backward_recv_backward: 反向往之前的stage发grad
  7. 等上一个batch的backward send_prev发完, 相当于一个buffer切换过程.

整个流程像下面这个流水线示意图.

image-20250206201325086
cooldown

与warmup刚好完全相反, 只有backward的计算和通信操作.

注意在每个阶段完成的时候都回将通信用到的output_tensor重新释放回显存池, 用来缓解显存压力.

参考:#

对VPP的进一步优化: https://zhuanlan.zhihu.com/p/681363624

作者:sunstrikes

出处:https://www.cnblogs.com/sunstrikes/p/18701639

版权:本作品采用「署名-非商业性使用-相同方式共享 4.0 国际」许可协议进行许可。

posted @   SunStriKE  阅读(50)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· DeepSeek-R1本地部署如何选择适合你的版本?看这里
· 开源的 DeepSeek-R1「GitHub 热点速览」
· 传国玉玺易主,ai.com竟然跳转到国产AI
· 揭秘 Sdcb Chats 如何解析 DeepSeek-R1 思维链
· 自己如何在本地电脑从零搭建DeepSeek!手把手教学,快来看看! (建议收藏)
点击右上角即可分享
微信分享提示
more_horiz
keyboard_arrow_up light_mode palette
选择主题
menu