Megatron-LM

Megatron-LM 使大规模训练变压器语言模型成为可能。 它提供了高效的张量并行、管道并行和基于序列的模型并行，用于预训练基于变压器的语言模型，如 GPT（仅解码器）、BERT（仅编码器）和 T5（编码器-解码器）。 有关详细信息和工作原理，请参阅 GitHub 仓库。

集成了哪些功能？

Accelerate 集成了 Megatron-LM 的以下功能，以支持大规模预训练/微调 BERT（编码器）、GPT（解码器）或 T5 模型（编码器和解码器）：

a. 张量并行（TP）：在节点内秩之间减少内存占用，而不会增加太多额外的通信。 每个张量被分成多个块，每个块位于不同的 GPU 上。在每一步，相同的 mini-batch 数据由每个块独立并行处理，然后在所有 GPU 之间同步（all-reduce 操作）。 在简单的变压器层中，这会导致前向路径中有 2 次 all-reduce，后向路径中有 2 次 all-reduce。 更多详细信息，请参阅研究论文 Megatron-LM: 使用模型并行训练数十亿参数的语言模型 和博客文章的这一部分 BLOOM 训练背后的技术。

b. 管道并行（PP）：通过节点间并行化减少内存占用，支持大规模训练。 通过 PipeDream-Flush 调度/1F1B 调度和交错 1F1B 调度减少朴素 PP 的气泡。 层在 PP 阶段之间均匀分布。例如，如果模型有 24 层，我们有 4 个 GPU 用于管道并行，每个 GPU 将有 6 层（24/4）。 有关减少 PP 空闲时间的调度的更多详细信息，请参阅研究论文 使用 Megatron-LM 在 GPU 集群上高效训练大规模语言模型 和博客文章的这一部分 BLOOM 训练背后的技术。

c. 序列并行（SP）：在不增加额外通信的情况下减少内存占用。仅在使用 TP 时适用。 它通过用 reduce-scatter 替换 all-reduce 后的 all-gather 操作来减少激活内存需求，防止相同的副本出现在张量并行秩上。 由于 all-reduce = reduce-scatter + all-gather，这在不增加通信成本的情况下节省了大量的激活内存。 简单来说，它将每个变压器层的输出沿序列维度切分，例如，如果序列长度为 1024，TP 大小为 4，每个 GPU 将有 256 个标记（1024/4）用于每个样本。 这增加了支持的训练批次大小。更多详细信息，请参阅研究论文 减少大规模变压器模型中的激活重计算。

d. 数据并行（DP） 通过分布式优化器：通过在 DP 秩之间切分优化器状态和梯度来减少内存占用 （与传统方法在数据并行秩之间复制优化器状态相比）。 例如，当使用混合精度训练的 Adam 优化器时，每个参数占用 12 字节的内存。 这在 GPU 之间均匀分布，即，如果有 4 个 GPU，每个参数将占用 3 字节（12/4）。 更多详细信息，请参阅研究论文 ZeRO: 朝向训练万亿参数模型的内存优化 和博客文章的这一部分 BLOOM 训练背后的技术。

e. 选择性激活重计算：通过智能激活检查点显著减少激活内存占用。 它不存储占用大量内存的激活，而是在重计算时速度较快，从而在内存和重计算之间取得很好的平衡。 例如，对于 GPT-3，这可以将激活所需的内存减少 70%，而重计算激活的 FLOPs 开销仅为 2.7%。 更多详细信息，请参阅研究论文 减少大规模变压器模型中的激活重计算。

f. 融合内核：融合 Softmax、混合精度融合层归一化和线性层权重梯度计算的融合梯度累积。 PyTorch JIT 编译的融合 GeLU 和融合 Bias+Dropout+Residual 加法。

g. 支持索引数据集：用于大规模训练的数据集的高效二进制格式。支持 mmap、cached 索引文件和 lazy 加载器格式。

h. 检查点重塑和互操作性：用于将不同张量和管道并行大小的 Megatron-LM 检查点重塑为受广泛工具支持的 Transformers 分片检查点 （如 Accelerate Big Model Inference、Megatron-DeepSpeed Inference 等）。 还支持将 Transformers 分片检查点转换为不同张量和管道并行大小的 Megatron-LM 检查点，以支持大规模训练。

先决条件

你需要安装最新版本的 PyTorch、CUDA、NCCL 和 NVIDIA APEX 以及 nltk 库。 更多详细信息，请参阅 文档。 另一种设置环境的方法是从 NGC 拉取一个包含所有所需安装的 NVIDIA PyTorch 容器。

以下是设置 conda 环境的步骤：

1. 创建一个虚拟环境

conda create --name ml

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

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

3. 安装 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 ..

4. 安装 Megatron-LM

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.5.0
pip install --no-use-pep517 -e .

加速 Megatron-LM 插件

重要功能直接通过 accelerate config 命令支持。 使用 Megatron-LM 功能的相应问题示例如下：

:~$ 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 parallel 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

生成的配置如下所示：

~$ 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

我们将以 GPT 预训练为例。对官方的 run_clm_no_trainer.py 进行最小修改以使用 Megatron-LM 的步骤如下：

1. 由于 Megatron-LM 使用其自己的 Optimizer 实现，因此需要使用与之兼容的调度器。因此，仅支持 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,
    )

2. 获取总批次大小的详细信息现在需要了解张量并行和管道并行的大小。 获取有效总批次大小的示例如下所示：

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

3. 在使用 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)

4. 对于 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
    )

好了！我们准备就绪 🚀。请在 accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py 路径下的 examples 文件夹中找到示例脚本。 让我们使用 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 索引数据集的高级功能

要利用更多功能，请参阅以下详细信息。

1. 以下是在使用 Megatron-LM 时自定义训练步骤所需更改的示例。 你将实现 accelerate.utils.AbstractTrainStep 或继承其相应的子类 accelerate.utils.GPTTrainStep、accelerate.utils.BertTrainStep 或 accelerate.utils.T5TrainStep。

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,
    }

2. 使用 Megatron-LM 数据集时，需要进行一些额外的更改。这些数据集的数据加载器仅在每个张量并行组的 rank 0 上可用。因此，有些 rank 上将没有数据加载器，这需要对训练循环进行调整。能够完成这些操作展示了 Accelerate 的灵活性和可扩展性。所需的更改如下。

a. 对于 Megatron-LM 索引数据集，我们需要使用 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

b. 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
)

c. 训练和评估循环的更改，因为数据加载器仅在张量并行等级 0 上可用 因此，我们只有在数据加载器不为 None 时才进行迭代，否则提供一个空字典 因此，我们使用 while 循环，并在 completed_steps 等于 args.max_train_steps 时退出循环 这类似于 Megatron-LM 的设置，用户在使用 Megatron-LM 索引数据集时需要提供 max_train_steps。 这展示了 Accelerate 的灵活性和可扩展性。

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)

Checkpoint 重塑和互操作性工具

1. 这些脚本位于 Transformers 库中相应模型的目录下。 目前，GPT 模型的脚本为 checkpoint_reshaping_and_interoperability.py

2. 以下是一个将 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

3. 将 checkpoint 从 transformers 转换为 megatron，其中 tp_size=2，pp_size=2 和 dp_size=2。

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 模型支持返回 logits 和 megatron_generate 函数用于文本生成

1. 返回 logits 需要在 MegatronLMPlugin 中设置 require_logits=True，如下所示。 这些 logits 将在管道的最后一个阶段可用。

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

2. megatron_generate 方法用于 Megatron-LM GPT 模型：这将使用张量并行和管道并行来完成一批输入的生成，无论是在使用贪心搜索时（带或不带 top_k/top_p 采样）还是在使用束搜索解码时对单个提示输入进行生成。仅支持 transformers 生成方法的一部分功能。这将有助于通过张量并行和管道并行使用大型模型进行生成（已经实现了键值缓存，并默认使用融合内核）。 这要求数据并行大小为 1，序列并行和激活检查点功能必须禁用。 还需要指定分词器的词汇表文件和合并文件的路径。 以下示例展示了如何配置和使用 megatron_generate 方法来生成 Megatron-LM GPT 模型。

# specifying tokenizer's vocab and merges file
vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

# inference using `megatron_generate` functionality
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 sampling
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 sampling
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)

# adding `bos` token at the start
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,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

3. 一个使用 megatron_generate 方法的端到端示例，适用于 Megatron-LM GPT 模型，可以在 megatron_gpt2_generation.py 中找到，配置文件为 megatron_lm_gpt_generate_config.yaml。包含加速启动命令的 bash 脚本可以在 megatron_lm_gpt_generate.sh 中找到。脚本的输出日志可以在 megatron_lm_gpt_generate.log 中找到。

支持 ROPE 和 ALiBi 位置嵌入及多查询注意力

1. 对于 ROPE/ALiBi 注意力，将 position_embedding_type 设置为 ("absolute" | "rotary" | "alibi") 并传递给 MegatronLMPlugin，如下所示。

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

2. 对于多查询注意力机制，将 attention_head_type 设置为 ("multihead" | "multiquery") 并传递给 MegatronLMPlugin，如下所示。

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

注意事项

1. 支持 Transformers GPT2、Megatron-BERT 和 T5 模型。 这涵盖了仅解码器、仅编码器和编码器-解码器模型类。

2. 模型前向传递仅返回损失，因为背后有相当复杂的管道、张量和数据并行性交互。 model(**batch_data) 调用返回在数据并行等级上平均的损失。 对于大多数情况，使用 Megatron-LM 特性运行预训练任务时，这已经足够，你可以轻松地使用损失计算 困惑度。 对于 GPT 模型，支持返回损失之外的 logits。 这些 logits 不会在数据并行等级上收集。使用 accelerator.utils.gather_across_data_parallel_groups 在数据并行等级上收集 logits。这些 logits 与标签可以用于计算各种性能指标。

3. 主进程是管道的最后一个等级，因为损失/logits 在管道的最后一个阶段可用。 使用 Megatron-LM 集成时，accelerator.is_main_process 和 accelerator.is_local_main_process 会在最后一个等级返回 True。

4. 在 accelerator.prepare 调用中，会创建一个与给定 Transformers 模型对应的 Megatron-LM 模型，权重随机初始化。 请使用 accelerator.load_state 加载与 TP、PP 和 DP 分区匹配的 Megatron-LM 检查点。

5. 目前，检查点重塑和互操作性支持仅适用于 GPT。 很快将扩展到 BERT 和 T5。

6. gradient_accumulation_steps 需要设置为 1。使用 Megatron-LM 时，管道并行设置中的微批次与梯度累积同义。

7. 使用 Megatron-LM 时，使用 accelerator.save_state 和 accelerator.load_state 保存和加载检查点。

8. 以下是 Megatron-LM 模型架构与等效的 Transformers 模型架构的映射。 仅支持这些 Transformers 模型架构。

a. Megatron-LM BertModel： config 中模型类型为 megatron-bert 的 Transformers 模型，例如， MegatronBERT

b. Megatron-LM GPTModel： config 中模型类型为 gpt2 的 Transformers 模型，例如， OpenAI GPT2

c. Megatron-LM T5Model： config 中模型类型为 t5 的 Transformers 模型，例如， T5 和 MT5