大模型训练常见工具概念

备忘用

NCCL

优化GPU内通信的原语

NCCL是一个独立的库，用作GPU的标准通信路由，实现了 all-reduce/all0gather/reduce/broadcast/reduce-scatter、基于发送、接受的通信模式（scatter, gather, all-to-all）。为PCIe、NVLink、NVswitch等平台、InfiniBand Verbs和TCP、IP socket等网络优化以取得高带宽。

NCCL支持单、多节点任意数量的GPU，可以用作单、多进程任务。

使用:

安装NCCL库
修改应用，链接到库
include头文件nccl.h到应用
创建一个communicator
使用NCCL收集通信原语，实现数据交流。

TVM

Open Deep Learning Compiler Stack

开源机器学习编译框架，用作CPU、GPU、ML加速卡。

目的是在任何硬件后台上高效优化、运行计算任务。

TVM同深度学习框架运作，来提供对不同后端的端到端编译。

工作流程：

从PyTorch等框架导入模型
翻译为Relay，为TVM的高级模型语言。被导入进TVM的模型会用Relay进行标示。Relay是函数型语言并且是神经网络的IR (intermediate representation)。Relay支持：
- 传统的数据流式表示
- 函数式作用域，let-binding令其成为全功能、有区分性的语言
- 允许用户混合两种编程风格
提供图级优化，来优化模型
Lowering到张量表达式（TE）表示，将高级表示转化为低级表示称为lower。高级优化后，Relay运行FuseOps pass，将模型分区到多个小子图，并将子图lower到TE表示。TE是一种特定领域的语言，描述张量计算。TE也提供的schedule，进行特定低级循环优化，例如tiling、向量化、并行化等，为了将Relay转化成TE，TVM引用了一个Tensor Operator Inventory（TOPI），其中包含了一些常见张量操作的预定义模板。
使用自动微调米快AutoTVM或AutoScheduler搜索最优schedule。Schedule指定TE中的一个操作或子图的低级循环优化。自动微调模块寻找最佳schedule，并且根据花销等对齐进行比较。TVM中有两种自动微调模块：
- AutoTVM：基于模板的自动微调模块。运行搜索算法，在用户定义模板中，寻找可调节的最佳值。大多数的操作来说，TOPI已经提供了他们的模板。
- AutoScheduler (a.k.a. Ansor)，免模板自动微调模块，分析计算定义自动生成搜索空间，根据生成的搜索空间寻找最佳schedule。
选取模型编译的最优配置。微调之后，自动微调模块生成json格式的微调记录，为每个子图选取最优schedule。
lower到张量钟健表示(TIR),TVM的低级中间表示。在基于微调步骤选取最优配置后，每个TE子图被lower到TIR，可用低级有优化方法进行优化。然后，优化的TIR被lower到目标硬件平台的TIR。这是最终代码生成阶段，来提供一个可以部署生产的最优模型。TVM支持一些不同编译器后台，包括LLVM、特定编译器（NVCC等）、嵌入和特定目标。
编译得到机器码。这个处理结束后，特定编译器生成的代码可以lower到机器码。

XLA

Accelerated Linear AlgebrA

机器学习编译器，用作GPU、CPU、ML加速卡。

从一些ML框架中提取模型，框架包括PyTorch，TensorFlow、JAX等，对这些进行优化，满足在GPU、CPU、ML加速卡上高性能运行。

寻找计算图中适合编译的子图进行加速。

pytorch简单示例

import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = MNIST().train().to(device)
loss_fn = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

for data, target in train_loader:
    optimizer.zero_grad()
    data = data.to(device)
    target = target.to(device)
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()

optimizer.step()
xm.mark_step()

DeepSpeed

DeepSpeed是对transformer结构的大型模型的计算框架，可以用作训练、推理、模型压缩等。

DeepSpeed对PyTorch进行轻量级wrap，包含主流SOTA训练技巧，例如分布式训练、混合精度、梯度累积、存档点等。

利用的一些优化算保证了计算性能、内存性能、可拓展性、通讯效率、数据效率。

特性包括：

混合精度分布式训练
- 16位混合精度
- 单GPU、多GPU、多节点
模型并行
- 自定义模型并行
- 与Megatron-LM整合
流水线并行
- 3D 并行(DP+TP+PP)
ZeRO
- 优化器状态和梯度分区
- activation分区
- 固定buffer优化
- 连续内存优化
ZeRO-Offload
- 利用CPU+GPU的内存训练
- 支持单GPU的10B模型训练
Ultra-fast dense transformer kernels
稀疏注意力
- 高效内存和计算的稀疏内核
- 比dense长10倍的seq
- 灵活支持不同稀疏结构
1bit Adam, 0/1 Adam, 1-bit LAMB
- 自定义通信的收集函数
- 节省26倍通信量
其他内存和带宽优化
- 只能梯度累积
- 通信、计算重叠
训练特性
- 简化训练API
- 梯度裁切
- 混合精度自动损失放缩
训练优化器
- Fused Adam和任意torch.optim.Optimizer
- 优化内存带宽的FP16优化器
- 使用LAMB优化器的大批次训练
- 使用ZeRO优化器的高效内存训练
- CPU-Adam
训练checkpoint
高级参数搜索
- 学习率范围测试
- 1Cycle Learning Rate Schedule
简化数据loader
高效数据
- 通过curriculum learning 进行高效数据采样，通过随机分层令牌丢弃进行有效的数据路由
- 在GPT-3/BERT预训练和GPT/ViT微调中，节省至多2倍的数据和时间
- 在相同数据、时间下，进一步提升模型质量
Curriculum Learning
- 基于curriculum learning的数据流水线，提供更简单的训练样本
- 稳定并提升速度、批次大小、学习率，并保证token级别的收敛速度
- 请注意，上面的数据效率库提供了更通用的curriculum learning支持。这个curriculum learning仍然受到支持，但我们建议使用Data Efficiency library。
Progressive Layer Dropping
- 高效、鲁棒训练
- 预训练收敛速度提升
性能分析、debug
Mixture of Experts (Moe)

FSDP

Fully Sharded Data Parallel

是一种数据并行训练策略，与传统DP不同，传统DP在每个GPU里保存一份模型参数、梯度、优化器状态。

FSDP将这些状态在所有节点上进行分区，这样也可以进行offload。

下图展示了两个数据并行节点的fsdp。

工作流程：

constructor：

将模型参数分片，分给每一个rank。
正向传播：
- 对于此fsdp单元，运行all-gather，从每个rank收集所有shard，得到所有参数
- 进行正向传播计算
- 丢弃刚刚收集的所有参数分片
反向传播：
- 对于此fsdp单元，运行all-gather，从每个rank收集所有shard，得到所有参数
- 进行反向传播计算
- 运行reduce-scatter来同步梯度
- 丢弃参数

在pytorch中使用时对模型进行wrap即可，有两种方法。

自动wrap：对模型整体进行wrap
手动wrap：可对模型中的指定层级进行wrap，之后设置wrap模式。

Megatron

transformer结构的大语言模型训练框架。

采用混合精度，高效的模型并行（张量、序列、管道）和数据并行，多节点预训练。

其中，其采用的模型并行可为对transformer进行module级进行划分。

官方提供的简单语言模型训练demo如下，删除部分代码

import torch
from functools import partial
from megatron import get_args
from megatron.arguments import core_transformer_config_from_args
from megatron import print_rank_0
from megatron import get_timers
from megatron import get_tokenizer
from megatron.core import tensor_parallel
from megatron.core.enums import ModelType
from megatron.data.gpt_dataset import build_train_valid_test_datasets
from megatron.core.models.gpt import GPTModel
from megatron.training import pretrain
from megatron.utils import get_ltor_masks_and_position_ids
from megatron.utils import average_losses_across_data_parallel_group

def model_provider(pre_process=True, post_process=True):
    构建模型
    return model


def get_batch(data_iterator):
    """Generate a batch"""
    args = get_args()
    tokenizer = get_tokenizer()

    # Items and their type.
    keys = ['text']
    datatype = torch.int64

    # Broadcast data.
    if data_iterator is not None:
        data = next(data_iterator)
    else:
        data = None
    data_b = tensor_parallel.broadcast_data(keys, data, datatype)

    # Unpack.
    tokens_ = data_b['text'].long()
    labels = tokens_[:, 1:].contiguous()
    tokens = tokens_[:, :-1].contiguous()

    # Get the masks and postition ids.
    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
        tokens,
        tokenizer.eod,
        args.reset_position_ids,
        args.reset_attention_mask,
        args.eod_mask_loss)

    return tokens, labels, loss_mask, attention_mask, position_ids

def loss_func(loss_mask, output_tensor):
    losses = output_tensor.float()
    loss_mask = loss_mask.view(-1).float()
    loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()

    # Reduce loss for logging.
    averaged_loss = average_losses_across_data_parallel_group([loss])

    return loss, {'lm loss': averaged_loss[0]}


def forward_step(data_iterator, model):
    """Forward step."""
    args = get_args()
    timers = get_timers()

    # Get the batch.
    timers('batch-generator', log_level=2).start()
    tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
        data_iterator)
    timers('batch-generator').stop()

    output_tensor = model(tokens, position_ids, attention_mask,
                          labels=labels)

    return output_tensor, partial(loss_func, loss_mask)


def train_valid_test_datasets_provider(train_val_test_num_samples):
	构建三个dataset
    return train_ds, valid_ds, test_ds


if __name__ == "__main__":

    pretrain(train_valid_test_datasets_provider, model_provider,
             ModelType.encoder_or_decoder,
             forward_step,
             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}

LLVM

LLVM是一个集合，包含了模组化、可复用的编译器和工具链。

标签：loss,训练,优化,模型,常见,mask,import,data
From： https://www.cnblogs.com/antelx/p/17523886.html