【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化

这篇具有很好参考价值的文章主要介绍了【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

相关博客
【深度学习】【分布式训练】一文捋顺千亿模型训练技术:流水线并行、张量并行和3D并行
【Megatron-DeepSpeed】张量并行工具代码mpu详解(四):张量并行版Embedding层及交叉熵的实现及测试
【Megatron-DeepSpeed】张量并行工具代码mpu详解(三):张量并行层的实现及测试
【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化
【Megatron-DeepSpeed】张量并行工具代码mpu详解(二):Collective通信操作的封装mappings
【深度学习】【分布式训练】DeepSpeed:AllReduce与ZeRO-DP
【深度学习】混合精度训练与显存分析
【深度学习】【分布式训练】Collective通信操作及Pytorch示例
【自然语言处理】【大模型】大语言模型BLOOM推理工具测试

Megatron-DeepSpeed:并行环境初始化

​ Megatron-DeepSpeed是DeepSpeed版本的NVIDIA Megatron-LM。像BLOOM、GLM-130B等主流大模型都是基于Megatron-DeepSpeed开发的。这里以BLOOM版本的Megetron-DeepSpeed为例,介绍其张量并行代码mpu的细节(位于megatron/mpu下)。

​ 理解该部分的代码需要对张量并行的原理以及集合通信有一定的理解,可以看文章:

  • 【深度学习】【分布式训练】Collective通信操作及Pytorch示例
  • 【深度学习】【分布式训练】一文捋顺千亿模型训练技术:流水线并行、张量并行和3D并行
  • 【深度学习】【分布式训练】DeepSpeed:AllReduce与ZeRO-DP

阅读建议:

  1. 本文仅会解析核心代码,并会不介绍所有代码;
  2. 本文会提供一些测试脚本来展现各部分代码的功能;
  3. 建议实际动手实操来加深理解;
  4. 建议对Collective通信以及分布式模型训练有一定理解,再阅读本文;

一、总览

​ mpu目录下核心文件有:

  • initialize.py:负责数据并行组、张量并行组和流水线并行组的初始化,以及获取与各类并行组相关的信息;
  • data.py:实现张量并行中的数据广播功能;
  • cross_entropy.py:张量并行版本的交叉熵;
  • layers.py:并行版本的Embedding层,以及列并行线性层和行并行线性层;
  • mappings.py:用于张量并行的通信操作;

二、初始化原理

​ Megatron-Deepspeed框架能够支持3D并行,而3D并行中显卡如何分配成不同的组等工作就是由mpu目录下的initialize.py完成的。

【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化,自然语言处理,自然语言处理,大模型,分布式训练,Megatron,张量并行

​ 假设有两个节点Node1Node2,每个节点有8个GPU,共计16个GPU。16个GPU的编号分别为Rank0、Rank1、…、Rank15。此外,假设用户设置流水线并行度为4,张量并行度为2

  • 流水线并行

    流水线并行会将整个模型划分为4份,这里称为sub_model_1至sub_model_4。每连续的4张GPU负责一个sub_model。即上图右上角中,相同颜色的GPU负责相同的sub_model

    每个流水线并行组对应一个完整的模型,最终的流水线并行组:

    [Rank0, Rank4, Rank8, Rank12],
    
    [Rank1, Rank5, Rank9, Rank13],
    
    [Rank2, Rank6, Rank10, Rank14],
    
    [Rank3, Rank7, Rank11, Rank15],
    
  • 张量并行

    张量并行会针对流水线并行中的sub_model来进行张量的拆分。即Rank0和Rank1负责一份sub_model_1,Rank2和Rank3负责另一份sub_model_1;Rank4和Rank5负责sub_model_2,Rank6和Rank7负责另一份sub_model_2;以此类推。上图右下角中,绿色线条表示单个张量并行组,每个张量并行组都共同负责某个具体的sub_model。

    最终的张量并行组:

    [Rank0, Rank1], [Rank2, Rank3],
    
    [Rank4, Rank5], [Rank6, Rank7],
    
    [Rank8, Rank9], [Rank10, Rank11],
    
    [Rank12, Rank13], [Rank14, Rank15],
    
  • 数据并行

    数据并行的目的是要保证并行中的相同模型参数读取相同的数据。经过流水线并行和张量并行后,Rank0和Rank2负责相同的模型参数,所以Rank0和Rank2是同一个数据并行组。上图左上角中的红色线条表示数据并行组。

    [Rank0, Rank2], [Rank1, Rank3],
    
    [Rank4, Rank6], [Rank5, Rank7],
    
    [Rank8, Rank10], [Rank9, Rank11],
    
    [Rank12, Rank14], [Rank13, Rank15],
    

三、初始化代码

​ 这里不对代码的所有细节进行解析。理解上面的原理后,细节部分只要花些时间即可弄明白。这里仅对整体的代码结构或者某些有代表性的函数进行说明

​ 总的来说,初始化的目标就是要赋予下面这些变量具体的值。

# Intra-layer model parallel group that the current rank belongs to.
_TENSOR_MODEL_PARALLEL_GROUP = None
# Inter-layer model parallel group that the current rank belongs to.
_PIPELINE_MODEL_PARALLEL_GROUP = None
# Model parallel group (both intra- and pipeline) that the current rank belongs to.
_MODEL_PARALLEL_GROUP = None
# Embedding group.
_EMBEDDING_GROUP = None
# Data parallel group that the current rank belongs to.
_DATA_PARALLEL_GROUP = None

_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None

# These values enable us to change the mpu sizes on the fly.
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_TENSOR_MODEL_PARALLEL_RANK = None
_MPU_PIPELINE_MODEL_PARALLEL_RANK = None

# A list of global ranks for each pipeline group to ease calculation of the source
# rank when broadcasting from the first or last pipeline stage
_PIPELINE_GLOBAL_RANKS = None

​ 初始化功能的核心是函数 initialize_model_parallel,其主要的参数就是用户指定的张量并行数tensor_model_parallel_size_以及流水线并行数pipeline_model_parallel_size_。基于这两个用户参数来计算出各种分组。

def initialize_model_parallel(tensor_model_parallel_size_=1,
                              pipeline_model_parallel_size_=1,
                              virtual_pipeline_model_parallel_size_=None):
    ...

​ 其余的函数主要功能是:获得初始化状态、获得当前rank所在的组、获得当前进行所在特定组里的rank等。这里功能会在后续的测试代码里展示具体的应用。

​ 此外,官方代码中的函数 destroy_model_parallel中应该是存在bug,下面是可正常运行的版本:

def destroy_model_parallel():
    """Set the groups to none."""
    global _TENSOR_MODEL_PARALLEL_GROUP
    _TENSOR_MODEL_PARALLEL_GROUP = None
    global _PIPELINE_MODEL_PARALLEL_GROUP
    _PIPELINE_MODEL_PARALLEL_GROUP = None
    global _DATA_PARALLEL_GROUP
    _DATA_PARALLEL_GROUP = None
    # 以下为新增
    global _MODEL_PARALLEL_GROUP
    _MODEL_PARALLEL_GROUP = None
    global _EMBEDDING_GROUP
    _EMBEDDING_GROUP = None

四、测试代码

1. 测试设置

  • 测试使用8张GPU;
  • 张量并行度为2,流水线并行度为2;
  • 依据上面介绍的原理,流水线并行组为:[Rank0, Rank4],[Rank1, Rank5],[Rank2, Rank6], [Rank3, Rank7]
  • 依据上面介绍的原理,张量并行组为:[Rank0, Rank1],[Rank2, Rank3],[Rank4,Rank5],[Rank6,Rank7]
  • 依据上面介绍的原理,数据并行组为: [Rank0, Rank2],[Rank1, Rank3],[Rank4,Rank6],[Rank5,Rank7]

2. 辅助代码

这里的辅助代码仍然使用原始单元测试的代码,即/megatron/mpu/tests/commons.py。 下面会对这些代码进行简单的注释。

# commons.py
import sys
sys.path.append("..")

import argparse
import os
import random
import numpy
import torch

import megatron.mpu as mpu # 由于本文测试代码位于项目的根目录下,因此修改了mpu的import方式


class IdentityLayer(torch.nn.Module):
    """
    一个单层网络,会在测试cross_entropy.py时使用。
    """
    def __init__(self, size, scale=1.0):
        super(IdentityLayer, self).__init__()
        self.weight = torch.nn.Parameter(scale * torch.randn(size))

    def forward(self):
        return self.weight


def set_random_seed(seed):
    """Set random seed for reproducability."""
    random.seed(seed)
    numpy.random.seed(seed)
    torch.manual_seed(seed)
    mpu.model_parallel_cuda_manual_seed(seed)


def initialize_distributed(backend='nccl'):
    """初始化分布式环境"""
    # Get local rank in case it is provided.
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=None,
                        help='local rank passed from distributed launcher')
    args = parser.parse_args()
    local_rank = args.local_rank

    # Get rank and world size.
    rank = int(os.getenv('RANK', '0')) # 当前进程所对应的rank(rank是显卡所对应的全局编号)
    world_size = int(os.getenv("WORLD_SIZE", '1')) # world_size是指所有可用显卡的数量

    print('> initializing torch.distributed with local rank: {}, '
          'rank: {}, world size: {}'.format(local_rank, rank, world_size))

    # Set the device id.
    device = rank % torch.cuda.device_count()
    if local_rank is not None:
        device = local_rank
    torch.cuda.set_device(device)

    # 初始化分布式环境所需要的相关代码
    init_method = 'tcp://'
    master_ip = os.getenv('MASTER_ADDR', 'localhost')
    master_port = os.getenv('MASTER_PORT', '6000')
    init_method += master_ip + ':' + master_port
    torch.distributed.init_process_group(
        backend=backend, # 使用gpu时,backend最好选择nccl
        world_size=world_size,
        rank=rank,
        init_method=init_method)
        
def print_separator(message):
    """
    输出辅助函数
    """
    torch.distributed.barrier() # 保证所有进程在此处保存同步,主要是为了防止多进行输出混乱
    filler_len = (78 - len(message)) // 2
    filler = '-' * filler_len
    string = '\n' + filler + ' {} '.format(message) + filler
    if torch.distributed.get_rank() == 0:
        print(string, flush=True)
    torch.distributed.barrier()

3. 测试代码

3.1 测试代码

​ 为了展示各个函数的功能,这里并不使用原始项目中的测试代码,而是单独撰写的代码。

# test_initialize.py
import sys
sys.path.append("..")

from commons import print_separator
from commons import initialize_distributed
import megatron.mpu as mpu
import torch

def run_test(
        tensor_model_parallel_size: int,
        pipeline_model_parallel_size:int):
    print_separator(f'> Test: TP={tensor_model_parallel_size}, PP={pipeline_model_parallel_size}')
    mpu.initialize_model_parallel(
            tensor_model_parallel_size,
            pipeline_model_parallel_size) # 并行初始化
    world_size = torch.distributed.get_world_size() # world_size, 总GPU数量
    global_rank = torch.distributed.get_rank() # 当前GPU的编号
    tp_world_size = mpu.get_tensor_model_parallel_world_size() # 每个张量并行组中GPU的数量
    pp_world_size = mpu.get_pipeline_model_parallel_world_size() # 每个流水线并行组中GPU的数量
    dp_world_size = mpu.get_data_parallel_world_size() # 每个数据并行组中的GPU数量
    tp_rank = mpu.get_tensor_model_parallel_rank() # 在张量并行组中的编号
    pp_rank = mpu.get_pipeline_model_parallel_rank() # 在流水线并行组中的编号
    dp_rank = mpu.get_data_parallel_rank() # 在数据并行组中的编号
    tp_group = mpu.get_tensor_model_parallel_group()
    tp_group = torch.distributed.distributed_c10d._pg_group_ranks[tp_group] # 当前GPU所在张量并行组的映射字典
    pp_group = mpu.get_pipeline_model_parallel_group()
    pp_group = torch.distributed.distributed_c10d._pg_group_ranks[pp_group] # 当前GPU所在流水线并行组的映射字典
    dp_group = mpu.get_data_parallel_group()
    dp_group = torch.distributed.distributed_c10d._pg_group_ranks[dp_group] # 当前GPU所在数据并行组的映射字典
    torch.distributed.barrier()
    info = f"="*20 + \
            f"\n> global_rank={global_rank}\n" + \
            f"> world_size={world_size}\n" + \
            f"> tp_world_size={tp_world_size}\n" + \
            f"> pp_world_size={pp_world_size}\n" + \
            f"> dp_world_size={dp_world_size}\n" + \
            f"> tp_rank={tp_rank}\n" + \
            f"> pp_rank={pp_rank}\n" + \
            f"> dp_rank={dp_rank}\n" + \
            f"> tp_group={tp_group}\n" + \
            f"> pp_group={pp_group}\n" + \
            f"> dp_group={dp_group}\n"
    print(info, flush=True)
    torch.distributed.barrier()


if __name__ == '__main__':
    initialize_distributed() # 初始化分布式环境
    tensor_model_parallel_size = 2
    pipeline_model_parallel_size = 2
    run_test(tensor_model_parallel_size, pipeline_model_parallel_size)
3.2 启动脚本
deepspeed test_initialize.py
3.3 输出结果

【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化,自然语言处理,自然语言处理,大模型,分布式训练,Megatron,张量并行

​ 上图是Rank5的相关信息,符合前面对于其张量并行组、流水线并行组和数据并行组的分析文章来源地址https://www.toymoban.com/news/detail-524226.html

到了这里,关于【Megatron-DeepSpeed】张量并行工具代码mpu详解(一):并行环境初始化的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • PTM:大模型加速方法或框架(预训练阶段/推理阶段)的简介、常用框架(Megatron-LM/Colossal-AI/DeepSpeed等,FasterTransformer/FastLLM/vLLM

    PTM:大模型加速方法或框架(预训练阶段/推理阶段)的简介、常用框架(Megatron-LM/Colossal-AI/DeepSpeed等,FasterTransformer/FastLLM/vLLM/TurboTransformers等)、案例应用之详细攻略 导读 :不同训练框架实现参数高效微调算法的效率会存在很大差异:比如使用Huggingface Transformers、DeepSpeed和Alpa训

    2024年02月07日
    浏览(42)
  • 分布式并行训练(DP、DDP、DeepSpeed)

    [pytorch distributed] 01 nn.DataParallel 数据并行初步 数据并行 vs. 模型并行 数据并行 :模型拷贝(per device),数据 split/chunk(对batch切分) 每个device上都拷贝一份完整模型,每个device分别处理1个batch的一部分(如batch_size=64, 2个device, 每device处理32个样本) 梯度反向传播时,每个设备上

    2024年02月07日
    浏览(45)
  • 全参数finetune Ziya-LLaMA-13B相关模型,目前支持数据并行+张量并行+ZeRO

    这个示例主要用于全参数finetune Ziya-LLaMA-13B 相关模型,目前支持 数据并行+张量并行+ZeRO Ziya-Finetune-Small,后续按照格式替换成自己的数据,目前代码直接用文件读取,非datasets读取,所以建议git clone下来然后在配置里引用对应的数据路径 以Ziya-LLaMA-13B-Pretrain-v1为例,因为开源的是

    2024年02月11日
    浏览(44)
  • 类ChatGPT逐行代码解读(2/2):从零实现ChatLLaMA、ColossalChat、DeepSpeed Chat

    本文为《类ChatGPT逐行代码解读》系列的第二篇,上一篇是:从零实现Transformer、ChatGLM-6B:从位置编码/缩放点积注意力/多头注意力开始 本文模型的特点是都加了RLHF,对于这4个模型而言:TRL、ChatLLaMA、ColossalChat、DeepSpeed Chat 如果只关注两个 则可以更多关注下ColossalChat、DeepS

    2024年02月03日
    浏览(49)
  • 探索 Paratest: 并行测试工具的利器

    Paratest 是一个 PHP 测试框架的并行执行工具,它可以帮助开发者在短时间内加速执行单元测试、功能测试等任务,从而提高开发效率和代码质量。 Paratest 是一款基于 PHP 编写的命令行工具,旨在通过并行化测试执行过程来减少测试所需的时间。它支持 Phing、Composer、PHPUnit 等多

    2024年04月23日
    浏览(42)
  • AI:DeepSpeed Chat(一款帮用户训练自己模型的工具且简单/低成本/快 RLHF 训练类ChatGPT高质量大模型)的简介、安装、使用方法之详细攻略

    AI:DeepSpeed Chat(一款帮用户训练自己模型的工具且简单/低成本/快 RLHF 训练类ChatGPT高质量大模型)的简介、安装、使用方法之详细攻略 目录 DeepSpeed Chat的简介 DeepSpeed-Chat的产生背景 DeepSpeed-Chat的简介 DeepSpeed-Chat的三大功能 DeepSpeed-RLHF 系统三大优势 DeepSpeed Chat的安装和使用方法

    2023年04月22日
    浏览(48)
  • verilog手撕代码4——串行数据和并行数据转换

    2023.4.26 把串行输入的数据存储起来,移位寄存器移位,得到并行数据 把并行数据缓存下来,再移位一个个数据输出

    2024年02月03日
    浏览(46)
  • 【论文阅读】Megatron-LM要点

    本文主要是对李沐老师的b站分享做一下自己的理解和总结。 李沐老师b站分享 模型结构无非就是那样,相比而言,想要训练更大的模型而又能平稳进行,是一项非常高超的技术! nvidia 跟gpipe类似,也是模型并行,但是在任务切割上面跟gpipe不一样。 gpipe: transformer也可以,

    2024年02月11日
    浏览(38)
  • [NLP]深入理解 Megatron-LM

    NVIDIA Megatron-LM 是一个基于 PyTorch 的分布式训练框架,用来训练基于Transformer的大型语言模型。Megatron-LM 综合应用了数据并行(Data Parallelism),张量并行(Tensor Parallelism)和流水线并行(Pipeline Parallelism)来复现 GPT-3. 在自然语言处理(NLP)领域,大型模型能够提供更精准和强

    2024年02月10日
    浏览(72)
  • 【配电网优化】基于串行和并行ADMM算法的配电网优化研究(Matlab代码实现)

    💥 💥 💞 💞 欢迎来到本博客 ❤️ ❤️ 💥 💥 🏆 博主优势: 🌞 🌞 🌞博客内容尽量做到思维缜密,逻辑清晰,为了方便读者。 ⛳ 座右铭: 行百里者,半于九十。 📋 📋 📋 本文目录如下: 🎁 🎁 🎁 目录 💥1 概述 📚2 运行结果 🎉3 参考文献 🌈4 Matlab代码实现 随

    2023年04月23日
    浏览(66)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包