Accelerate 0.24.0文档四：Megatron-LM

参考《Megatron-LM》

文章目录

- 一、Megatron-LM集成简介
- 二、环境配置
- - - 设置conda环境的步骤：
- 二、Accelerate Megatron-LM Plugin
- 三、自定义训练过程
- 四、检查点转换
- 五、文本生成
- 六、支持ROPE 、 ALiBi和Multi-Query Attention
- 七、注意事项

一、Megatron-LM集成简介

在这里插入图片描述

在大规模语言模型训练中，Accelerate通过整合Megatron-LM的特性实现了以下功能：

张量并行（TP）：降低内存占用，减少节点内的通信量。每个张量被分割成多个部分，每个部分位于不同的GPU上。在每个步骤中，相同的小批量数据由每个部分独立并行处理，然后在所有GPU间进行同步（all-reduce操作）。在简单的Transformer层中，这导致前向路径有2次all-reduce操作，后向路径也有2次。详情请参阅研究论文Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism和🤗博客文章中的The Technology Behind BLOOM Training部分。
管道并行（PP）：通过管道并行降低内存占用并实现大规模训练。采用PipeDream-Flush调度/1F1B调度和交错的1F1B调度来减少朴素管道并行的空闲时间。层被均匀分布到不同的管道并行阶段。例如，如果一个模型有24层，而我们有4个GPU进行管道并行，则每个GPU将有6层（24/4）。了解如何减少管道并行空闲时间的详细信息，请参阅研究论文Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM和🤗博客文章中的The Technology Behind BLOOM Training部分。
序列并行（SP）：在使用TP时降低内存占用，无需额外通信。只适用于TP。它通过在序列维度上对每个transformer层的输出进行分片来减少激活内存需求。例如，如果序列长度为1024，TP大小为4，每个GPU将为每个样本拥有256个标记（1024/4）。这提高了训练所支持的批量大小。了解更多详情，请参阅研究论文Reducing Activation Recomputation in Large Transformer Models。
**数据并行（DP）**通过分布式优化器：通过在DP排名之间分片优化器状态和梯度来降低内存占用。例如，使用Adam优化器进行混合精度训练时，每个参数占用12字节的内存。这会均等地分布到各个GPU，即如果有4个GPU，每个参数将占用3字节（12/4）。详情请参阅研究论文ZeRO: Memory Optimizations Toward Training Trillion Parameter Models和🤗博客文章中的The Technology Behind BLOOM Training部分。
选择性激活重新计算：通过智能激活检查点技术显著减少激活内存占用。它快速重新计算激活，不存储占用大量内存的激活，从而在内存和重新计算之间取得了很好的折衷。例如，对于GPT-3，这将导致激活所需内存减少70%，而重新计算激活仅增加2.7%的FLOPs开销。了解更多详情，请参阅研究论文Reducing Activation Recomputation in Large Transformer Models。
融合内核：包括融合Softmax、混合精度融合层归一化以及线性层权重梯度计算的融合梯度累积。此外，还有PyTorch JIT编译的融合GeLU和融合偏置+Dropout+残差添加。
支持索引数据集：为大规模训练提供了高效的二进制数据集格式。支持mmap、cached索引文件以及lazy加载格式。
检查点重塑和互操作性：为了更好地与多种工具（如🤗 Accelerate Big Model Inference、Megatron-DeepSpeed Inference等）兼容，提供了将Megatron-LM检查点的可变张量和管道并行大小重塑为🤗 Transformers分片检查点的实用工具。同时也支持将🤗 Transformers分片检查点转换为Megatron-LM检查点的功能，以便进行大规模训练。

二、环境配置

在开始之前，你需要安装最新版本的PyTorch、CUDA、NCCL以及NVIDIA APEX发行版，以及nltk库。你可以参考文档获取更多详情。另一种设置环境的方法是使用NGC提供的NVIDIA PyTorch容器，它包含所有必需的安装。

设置conda环境的步骤：

创建一个虚拟环境

假设机器上已安装CUDA 11.3，安装相应的PyTorch GPU版本

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

安装Nvidia APEX

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
cd ..

安装Megatron-LM

pip install git+https://github.com/huggingface/Megatron-LM.git

二、Accelerate Megatron-LM Plugin

在《Accelerate 0.24.0文档二：DeepSpeed集成》中介绍过，启用deepspeed功能时可通过运行accelerate config，以问答的形式完成deepspeed的配置。启用Megatron-LM时也可以直接通过accelerate config命令来完成相关配置，以下是执行此命令时的示例问题：

:~$ accelerate config --config_file "megatron_gpt_config.yaml"
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2
How many different machines will you use (use more than 1 for multi-node training)? [1]: 
Do you want to use DeepSpeed? [yes/NO]: 
Do you want to use FullyShardedDataParallel? [yes/NO]: 
Do you want to use Megatron-LM ? [yes/NO]: yes
What is the Tensor Parallelism degree/size? [1]:2
Do you want to enable Sequence Parallelism? [YES/no]: 
What is the Pipeline Parallelism degree/size? [1]:2
What is the number of micro-batches? [1]:2
Do you want to enable selective activation recomputation? [YES/no]: 
Do you want to use distributed optimizer which shards optimizer state and gradients across data pralellel ranks? [YES/no]: 
What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
How many GPU(s) should be used for distributed training? [1]:4
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16

之后会得到如下配置文件（yaml）：

~$ cat megatron_gpt_config.yaml 
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MEGATRON_LM
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config:
  megatron_lm_gradient_clipping: 1.0
  megatron_lm_num_micro_batches: 2
  megatron_lm_pp_degree: 2
  megatron_lm_recompute_activations: true
  megatron_lm_sequence_parallelism: true
  megatron_lm_tp_degree: 2
  megatron_lm_use_distributed_optimizer: true
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false

下面是使用官方的 megatron_lm_gpt_pretraining.py脚本进行最小化修改，以使用Megatron-LM进行GPT预训练的说明。修改的要点包括：

优化器和调度器的使用：Megatron-LM使用自己的优化器实现，需要使用与之兼容的调度器，所以用户需要创建 accelerate.utils.MegatronLMDummyScheduler，示例如下：

from accelerate.utils import MegatronLMDummyScheduler

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    lr_scheduler = MegatronLMDummyScheduler(
        optimizer=optimizer,
        total_num_steps=args.max_train_steps,
        warmup_num_steps=args.num_warmup_steps,
    )
else:
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )

总的批量大小：现在需要考虑张量和管道并行大小来获取有效的总批量大小。

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
else:
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

损失值的处理：使用Megatron-LM时，损失值已经在数据并行组内进行了平均。

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses.append(loss)
else:
    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses = torch.tensor(losses)
else:
    losses = torch.cat(losses)

模型保存：对于Megatron-LM，需要使用 accelerator.save_state 进行模型保存。

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    accelerator.save_state(args.output_dir)
else:
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(
        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
    )

开始训练：下面是使用 4 个 A100-80GB GPU 运行 gpt-large：

accelerate launch --config_file megatron_gpt_config.yaml \
examples/by_feature/megatron_lm_gpt_pretraining.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--block_size 1024 \
--learning_rate 5e-5 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 5 \
--with_tracking \
--report_to "wandb" \
--output_dir "awesome_model"

以下是输出日志中的一些重要摘录：

Loading extension module fused_dense_cuda...
>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
 > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
Building gpt model in the pre-training mode.
The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
Preparing dataloader
Preparing dataloader
Preparing model
 > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
 > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
Preparing optimizer
Preparing scheduler
> learning rate decay style: linear
10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]

使用 accelerate.utils.MegatronLMPlugin可以设置更多功能。

三、自定义训练过程

在使用Megatron-LM进行训练时，可以对训练步骤和数据加载器进行定制化调整。

定制训练步骤：在使用Megatron-LM时，如果需要对训练步骤进行特定定制，可以实现 accelerate.utils.AbstractTrainStep 或继承其相应的子类，如 accelerate.utils.GPTTrainStep、accelerate.utils.BertTrainStep 或 accelerate.utils.T5TrainStep，以适配Megatron-LM的特定要求和功能。

from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group


# Custom loss function for the Megatron model
class GPTTrainStepWithCustomLoss(GPTTrainStep):
    def __init__(self, megatron_args, **kwargs):
        super().__init__(megatron_args)
        self.kwargs = kwargs

    def get_loss_func(self):
        def loss_func(inputs, loss_mask, output_tensor):
            batch_size, seq_length = output_tensor.shape
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = losses.view(-1) * loss_mask

            # Resize and average loss per sample
            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
            loss_per_sample = loss_per_sample / loss_mask_per_sample

            # Calculate and scale weighting
            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
            weights = 1.0 + self.kwargs["alpha"] * weights
            # Calculate weighted average
            weighted_loss = (loss_per_sample * weights).mean()

            # Reduce loss across data parallel groups
            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])

            return weighted_loss, {"lm loss": averaged_loss[0]}

        return loss_func

    def get_forward_step_func(self):
        def forward_step(data_iterator, model):
            """Forward step."""
            # Get the batch.
            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)

            return output_tensor, partial(self.loss_func, tokens, loss_mask)

        return forward_step


def main():
    # Custom loss function for the Megatron model
    keytoken_ids = []
    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
    for keyword in keywords:
        ids = tokenizer([keyword]).input_ids[0]
        if len(ids) == 1:
            keytoken_ids.append(ids[0])
    accelerator.print(f"Keytoken ids: {keytoken_ids}")
    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
        "keytoken_ids": keytoken_ids,
        "alpha": 0.25,
    }

Megatron-LM Dataloader：在处理Megatron-LM数据集时，数据加载器仅在张量并行组的一个特定GPU（即rank 0）上可用，在其他GPU上，数据加载器不可用。针对这种情况，需要进行一些调整：

需要使用 MegatronLMDummyDataLoader 并传递必要的数据集参数（如 data_path、seq_length 等），更多详见参数列表。

from accelerate.utils import MegatronLMDummyDataLoader

megatron_dataloader_config = {
    "data_path": args.data_path,
    "splits_string": args.splits_string,
    "seq_length": args.block_size,
    "micro_batch_size": args.per_device_train_batch_size,
}
megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True

获取数据加载器： 数据加载器需要在张量并行组的rank 0上获取。为了在训练、验证和测试过程中使用不同的数据加载器，megatron_dataloader 被使用了3次，分别用于获取训练、验证和测试数据加载器。这里的数量和比例由 args.splits_string 参数决定。
```
model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
)
```

. 处理数据加载器可用性： 在训练和评估循环中，需要考虑数据加载器在其他GPU上不可用的情况。因此，在迭代过程中，必须确保仅在数据加载器不为空时进行迭代（即仅在rank 0上有数据加载器时才执行训练或评估），否则需要提供一个空的字典作为占位符。为了实现这一点，使用了一个while循环，在达到指定的训练步数 args.max_train_steps 时结束循环。

while completed_steps < args.max_train_steps:
    model.train()
    batch = next(train_dataloader) if train_dataloader is not None else {}
    outputs = model(**batch)
    loss = outputs.loss
    ...

    if completed_steps % eval_interval == 0:
        eval_completed_steps = 0
        losses = []
        while eval_completed_steps < eval_iters:
            model.eval()
            with torch.no_grad():
                batch = next(eval_dataloader) if eval_dataloader is not None else {}
                outputs = model(**batch)

这种调整确保了在处理Megatron-LM数据集时，训练和评估循环能够正确地运行，并在必要时适应数据加载器在部分GPU上不可用的情况。这也展示了 🤗 Accelerate 在适应不同情况下的灵活性和可扩展性。

四、检查点转换

Megatron-LM使用模型并行(model parallel)、流水线并行(pipeline parallel)和数据并行(data parallel)来支持分布式训练。而🤗Transformers使用基于参数分片(sharding)的方法。两者检查点格式并不兼容。

使用transformers库的checkpoint_reshaping_and_interoperability.py脚本，将Megatron-LM转换到通用Transformers 检查点：

python checkpoint_reshaping_and_interoperability.py \
--convert_checkpoint_from_megatron_to_transformers \
--load_path "gpt/iter_0005000" \
--save_path "gpt/trfs_checkpoint" \
--max_shard_size "200MB" \
--tokenizer_name "gpt2" \
--print-checkpoint-structure

使用 tp_size=2 、 pp_size=2 和 dp_size=2 将检查点从通用Transformers 转换到Megatron-LM：

python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
--load_path "gpt/trfs_checkpoint" \
--save_path "gpt/megatron_lm_checkpoint" \
--target_tensor_model_parallel_size 2 \
--target_pipeline_model_parallel_size 2 \
--target_data_parallel_size 2 \
--target_params_dtype "bf16" \
--make_vocab_size_divisible_by 128 \
--use_distributed_optimizer \
--print-checkpoint-structure

五、文本生成

使用Megatron-LM GPT模型的megatron_generate方法可进行高效的文本生成。首先设置require_logits=True，可以让Megatron-LM的GPT模型返回输出层的logits，以供后续生成所用。

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

Megatron-LM提供了megatron_generate方法专门用于GPT模型的生成，它利用了张量并行和管道并行来为一批输入完成生成。主要包括贪心采样(top-k、top-p)的批量生成,以及束搜索解码的单样本生成。

该方法的优点在于可以利用张量并行和管道并行来进行生成，已经默认使用了键值缓存和融合内核（fused kernels）。但是，使用这个方法需要将数据并行大小设置为1，禁用序列并行性和激活检查点，并且需要指定分词器的词汇文件和合并文件。以下是配置示例：

# 定义了分词器的词汇文件和合并文件的路径
vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
# 创建MegatronLMPlugin实例，同时传入了额外的Megatron参数，这些参数包括分词器的词汇文件和合并文件。
other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

# i通过 model.megatron_generate 方法进行了不同方式的文本生成
tokenizer.pad_token = tokenizer.eos_token
max_new_tokens = 64
batch_texts = [
    "Are you human?",
    "The purpose of life is",
    "The arsenal was constructed at the request of",
    "How are you doing these days?",
]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)

# top-p 采样
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_p=0.8,
    top_p_decay=0.5,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# top-k 采样
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_k=50,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# 在输入文本前加入bos标记
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# beam search => only takes single prompt
batch_texts = ["The purpose of life is"]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    num_beams=20,
    length_penalty=1.5,
)

# 生成的token被解码为文本，并通过 tokenizer.batch_decode 方法转换为可读的文本形式。
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

Top-p采样：
- 使用了 top_p 参数进行文本生成，这是一种基于概率的采样方法。在每个位置，模型将根据概率动态地选择概率大于给定阈值 top_p 的token作为下一个token的候选，直到生成指定数量的token（max_new_tokens）。top_p 的值为0.8，表示每次选择token时保留概率大于0.8的候选token。此外，还有一个 top_p_decay 参数，用于在生成过程中逐步降低 top_p 的值。
Top-k采样：
- 使用了 top_k 参数进行文本生成，这是另一种基于概率的采样方法。在每个位置，模型将在概率最高的 top_k 个token中进行选择作为下一个token的候选，直到生成指定数量的token（max_new_tokens）。这种方法不考虑token的概率大小，只根据排名选择概率最高的token。

通常来说,top-p sampling可以产生更多样而新颖的文本，而top-k sampling计算简单高效，但生成内容的多样性较差，需根据实际场景需求选择使用。

束搜索是遍历多种可能输出序列,同时保留topk个可能序列。它能更全面搜索输出空间,找到更佳解。但由于需要追踪多条序列,计算更复杂，所以一般只用于单样本。

最后，megatron_gpt2_generation.py 提供了使用egatron_generate 方法的端到端示例，其中包含配置文件 megatron_lm_gpt_generate_config.yaml，还有启动脚本megatron_lm_gpt_generate.sh 及其输出日志megatron_lm_gpt_generate.log。

六、支持ROPE 、 ALiBi和Multi-Query Attention

对于 ROPE/ALiBi 注意力，可设置 position_embedding_type为"absolute" 、“rotary” 或"alibi"：

other_megatron_args = {"position_embedding_type": "alibi"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

对于Multi-Query Attention，可设置attention_head_type 为"multihead" 或 “multiquery”：

other_megatron_args = {"attention_head_type": "multiquery"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

七、注意事项

这段文字主要介绍了使用Megatron-LM集成时🤗Transformers的一些使用限制和注意事项:

目前只支持GPT2、Megatron-BERT和T5这三类模型结构，涵盖了decoder-only、encoder-only和encoder-decoder三种情况。
模型forward时只返回loss，不直接返回logit，这是因为管道、张量和数据并行之间存在相当复杂的相互作用。这些损失是在数据并行排名中进行了平均。对于大多数情况，特别是在使用Megatron-LM特性运行预训练任务时，可以使用损失来轻松计算困惑度。对于GPT模型，除了损失外，还支持返回logits。这些logits并未在数据并行排名中被收集，可以使用 accelerator.utils.gather_across_data_parallel_groups 来收集跨数据并行排名的logits，然后与标签一起用于计算各种性能指标。
主进程为最后一个流水线阶段rank,is_main_process等接口按此判断。
在 accelerator.prepare 调用中，会创建一个与给定Transformers模型对应的Megatron-LM模型，具有随机权重。需要使用 accelerator.load_state 来加载具有匹配TP、PP和DP分区的Megatron-LM检查点。
gradient_accumulation_steps需要为1。Megatron中的micro batch类似累积梯度。
保存和加载检查点用save_state和load_state接口。
支持Megatron和Transformers三种模型的对应关系，而其他 🤗 transformers 的模型架构则不受支持
- Megatron-LM BertModel 对应于 🤗 transformers 中具有 megatron-bert 模型类型的模型，如 MegatronBERT。
- Megatron-LM GPTModel 对应于 🤗 transformers 中具有 gpt2 模型类型的模型，如 OpenAI GPT2。
- Megatron-LM T5Model 对应于 🤗 transformers 中具有 t5 模型类型的模型，如 T5 和 MT5。