Megatron-LM

Megatron-LM 支持大规模训练大型 Transformer 语言模型。它为预训练基于 Transformer 的语言模型（例如 GPT（仅解码器）、BERT（仅编码器）和 T5（编码器-解码器））提供高效的张量、流水线和序列模型并行性。有关详细信息以及幕后工作原理，请参阅 github repo。

集成了什么？

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

a. 张量并行 (TP)：减少内存占用，而无需在节点内 ranks 上进行太多额外的通信。每个张量被分成多个块，每个分片驻留在单独的 GPU 上。在每个步骤中，每个分片独立并行地处理相同的 mini-batch 数据，然后跨所有 GPU 进行同步（all-reduce 操作）。在一个简单的 Transformer 层中，这会在前向路径中导致 2 个 all-reduces，在后向路径中导致 2 个。有关更多详细信息，请参阅研究论文 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 以及博客文章的此部分 BLOOM 训练背后的技术。

b. 流水线并行 (PP)：减少内存占用，并通过节点间并行化实现大规模训练。通过 PipeDream-Flush 调度/1F1B 调度和交错 1F1B 调度减少了朴素 PP 的气泡。层均匀分布在 PP 阶段。例如，如果一个模型有 24 层，并且我们有 4 个 GPU 用于流水线并行，则每个 GPU 将有 6 层 (24/4)。有关减少 PP 空闲时间的调度的更多详细信息，请参阅研究论文 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM 以及博客文章的此部分 BLOOM 训练背后的技术。

c. 序列并行 (SP)：减少内存占用，而无需任何额外的通信。仅在使用 TP 时适用。它减少了所需的激活内存，因为它通过用 reduce-scatter 替换后，防止相同的副本位于张量并行 ranks 的 all-reduce 上，并且 no-op 操作将被 all-gather 替换。由于 all-reduce = reduce-scatter + all-gather，这节省了大量的激活内存，而没有增加通信成本。简而言之，它沿着序列维度对每个 Transformer 层的输出进行分片，例如，如果序列长度为 1024，TP 大小为 4，则每个 GPU 对于每个样本将有 256 个 token (1024/4)。这增加了训练可以支持的 batch size。有关更多详细信息，请参阅研究论文 Reducing Activation Recomputation in Large Transformer Models。

d. 通过分布式优化器的 数据并行 (DP)：通过在 DP ranks 之间分片优化器状态和梯度（而不是在数据并行 ranks 之间复制优化器状态的传统方法）来减少内存占用。例如，当使用 Adam 优化器进行混合精度训练时，每个参数占用 12 字节的内存。这在 GPU 之间平均分配，即，如果我们有 4 个 GPU，则每个参数将占用 3 字节 (12/4)。有关更多详细信息，请参阅研究论文 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 以及博客的以下部分 BLOOM 训练背后的技术。

e. 选择性激活重计算：通过智能激活检查点显著减少激活的内存占用。它不存储占用大量内存的激活，同时快速重新计算，从而在内存和重计算之间实现良好的权衡。例如，对于 GPT-3，这导致激活所需的内存减少 70%，而激活重计算的 FLOPs 开销仅为 2.7%。有关更多详细信息，请参阅研究论文 Reducing Activation Recomputation in Large Transformer Models。

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 环境的逐步方法

创建一个虚拟环境

conda create --name ml

假设机器已安装 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

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

Accelerate 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 预训练为例。使用 Megatron-LM 对官方 run_clm_no_trainer.py 进行的最小更改如下

由于 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,
    )

现在获取总 batch size 的详细信息需要认知张量和流水线并行大小。下面显示了获取有效总 batch size 的示例

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
    )

就这样！我们准备好了 🚀。请在路径 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 索引数据集的高级功能

为了利用更多功能，请阅读以下详细信息。

以下是自定义 Train Step 时使用 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,
    }

为了使用 Megatron-LM 数据集，还需要进行一些更改。这些数据集的数据加载器仅在每个张量并行组的 rank 0 上可用。因此，有些 rank 数据加载器将不可用，这需要调整训练循环。能够做到这一切表明 Accelerate 是多么灵活和可扩展。所需的更改如下。

a. 对于 Megatron-LM 索引数据集，我们需要使用 MegatronLMDummyDataLoader 并将所需的数据集 args 传递给它，例如 data_path、seq_length 等。有关可用 args 的列表，请参阅此处。

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. 更改训练和评估循环，因为数据加载器仅在张量并行 ranks 0 上可用。因此，我们只需要在数据加载器不是 None 的情况下迭代，否则提供空字典。因此，我们使用 while 循环进行循环，并在 completed_steps 等于 args.max_train_steps 时中断。这类似于 Megatron-LM 设置，其中用户在使用 Megaton-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)

检查点重塑和互操作性的实用程序

这些脚本存在于 Transformers 库中各自的模型下。目前，它可用于 GPT 模型 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。

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 函数以进行文本生成

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

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

Megatron-LM GPT 模型的 megatron_generate 方法：当使用贪婪算法（带/不带 top_k/top_p 采样）时，这将使用张量和流水线并行性来完成一批输入的生成，当使用 beam search 解码时，这将使用张量和流水线并行性来完成单个 prompt 输入的生成。仅支持 transformers generate 功能的子集。这将有助于通过张量和流水线并行性使用大型模型进行生成（默认情况下已进行键值缓存并使用融合内核）。这需要数据并行大小为 1，禁用序列并行和激活检查点。它还需要指定 tokenizer 的词汇表文件和 merges 文件的路径。以下示例显示了如何为 Megatron-LM GPT 模型配置和使用 megatron_generate 方法。

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

使用 Megatron-LM GPT 模型的 megatron_generate 方法的端到端示例可在 megatron_gpt2_generation.py 中找到，配置文件为 megatron_lm_gpt_generate_config.yaml。带有 accelerate launch 命令的 bash 脚本可在 megatron_lm_gpt_generate.sh 中找到。脚本的输出日志可在 megatron_lm_gpt_generate.log 中找到。

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

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

对于多查询注意力机制，请将 attention_head_type 以及 ("multihead" | "multiquery") 传递给 MegatronLMPlugin，如下所示。

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

注意事项

支持 Transformers GPT2、Megatron-BERT 和 T5 模型。这涵盖了解码器专用、编码器专用和编码器-解码器模型类别。
由于幕后流水线、张量和数据并行之间存在相当复杂的相互作用，因此仅从模型前向传递返回损失。 model(**batch_data) 调用返回跨数据并行 ranks 平均的损失。这对于大多数使用 Megatron-LM 功能运行预训练作业的情况来说是可以接受的，并且您可以轻松地使用损失计算 perplexity。对于 GPT 模型，支持返回 logits 以及损失。这些 logits 不会在数据并行 ranks 之间收集。使用 accelerator.utils.gather_across_data_parallel_groups 在数据并行 ranks 之间收集 logits。这些 logits 以及标签可用于计算各种性能指标。
主进程是最后一个 rank，因为损失/logits 在流水线的最后阶段可用。当使用 Megatron-LM 集成时，accelerator.is_main_process 和 accelerator.is_local_main_process 对于最后一个 rank 返回 True。
在 accelerator.prepare 调用中，将创建一个与给定 Transformers 模型对应的 Megatron-LM 模型，并带有随机权重。请使用 accelerator.load_state 加载具有匹配 TP、PP 和 DP 分区的 Megatron-LM 检查点。
目前，检查点重塑和互操作性支持仅适用于 GPT。很快将扩展到 BERT 和 T5。
gradient_accumulation_steps 需要为 1。当使用 Megatron-LM 时，流水线并行设置中的微批次与梯度累积同义。
当使用 Megatron-LM 时，请使用 accelerator.save_state 和 accelerator.load_state 来保存和加载检查点。
以下是从 Megatron-LM 模型架构到等效 transformers 模型架构的映射。仅支持这些 transformers 模型架构。

a. Megatron-LM BertModel : transformers 模型，其配置的模型类型中包含 megatron-bert，例如，MegatronBERT

b. Megatron-LM GPTModel : transformers 模型，其配置的模型类型中包含 gpt2，例如，OpenAI GPT2

c. Megatron-LM T5Model : transformers 模型，其配置的模型类型中包含 t5，例如，T5 和 MT5

< > 在 GitHub 上更新