Accelerate 0.24.0文档 四：Megatron-LM

参考《Megatron-LM》

一、Megatron-LM集成简介

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

1. 张量并行（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部分。

2. 管道并行（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部分。

3. 序列并行（SP）：在使用TP时降低内存占用，无需额外通信。只适用于TP。它通过在序列维度上对每个transformer层的输出进行分片来减少激活内存需求。例如，如果序列长度为1024，TP大小为4，每个GPU将为每个样本拥有256个标记（1024/4）。这提高了训练所支持的批量大小。了解更多详情，请参阅研究论文Reducing Activation Recomputation in Large Transformer Models。

4. **数据并行（DP）**通过分布式优化器：通过在DP排名之间分片优化器状态和梯度来降低内存占用。例如，使用Adam优化器进行混合精度训练时，每个参数占用12字节的内存。这会均等地分布到各个GPU，即如果有4个GPU，每个参数将占用3字节（12/4）。详情请参阅研究论文ZeRO: Memory Optimizations Toward Training Trillion Parameter Models和🤗博客文章中的The Technology Behind BLOOM Training部分。

5. 选择性激活重新计算：通过智能激活检查点技术显著减少激活内存占用。它快速重新计算激活，不存储占用大量内存的激活，从而在内存和重新计算之间取得了很好的折衷。例如，对于GPT-3，这将导致激活所需内存减少70%，而重新计算激活仅增加2.7%的FLOPs开销。了解更多详情，请参阅研究论文Reducing Activation Recomputation in Large Transformer Models。

6. 融合内核：包括融合Softmax、混合精度融合层归一化以及线性层权重梯度计算的融合梯度累积。此外，还有PyTorch JIT编译的融合GeLU和融合偏置+Dropout+残差添加。

7. 支持索引数据集：为大规模训练提供了高效的二进制数据集格式。支持mmap、cached索引文件以及lazy加载格式。

8. 检查点重塑和互操作性：为了更好地与多种工具（如🤗 Accelerate Big Model Inference、Megatron-DeepSpeed Inference等）兼容，提供了将Megatron-LM检查点的可变张量和管道并行大小重塑为🤗 Transformers分片检查点的实用工具。同时也支持将🤗 Transformers分片检查点转换为Megatron-LM检查点的功能，以便进行大规模训练。

二、环境配置

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

设置conda环境的步骤：

1. 创建一个虚拟环境

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

   conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
   1

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 ..
   1
   2
   3
   4

4. 安装Megatron-LM

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

二、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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

之后会得到如下配置文件（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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

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

1. 优化器和调度器的使用：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,
       )
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13
   14
   15

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
   1
   2
   3
   4

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)
   1
   2
   3
   4
   5
   6
   7
   8
   9

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
       )
   1
   2
   3
   4
   5
   6
   7

5. 开始训练：下面是使用 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"
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13
   14

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

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]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

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

三、自定义训练过程

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

1. 定制训练步骤：在使用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,
       }
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20
   21
   22
   23
   24
   25
   26
   27
   28
   29
   30
   31
   32
   33
   34
   35
   36
   37
   38
   39
   40
   41
   42
   43
   44
   45
   46
   47
   48
   49
   50
   51
   52
   53
   54
   55
   56
   57
   58
   59
   60

2. 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
     1
     2
     3
     4
     5
     6
     7
     8
     9
     10

   • 获取数据加载器： 数据加载器需要在张量并行组的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
     )
     1
     2
     3

   • . 处理数据加载器可用性： 在训练和评估循环中，需要考虑数据加载器在其他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)
     1
     2
     3
     4
     5
     6
     7
     8
     9
     10
     11
     12
     13
     14
     15

  这种调整确保了在处理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
1
2
3
4
5
6
7

  使用 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
1
2
3
4
5
6
7
8
9
10

五、文本生成

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

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)
1

  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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

1. 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 的值。

2. 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)
1
2

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

七、注意事项

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

1. 目前只支持GPT2、Megatron-BERT和T5这三类模型结构，涵盖了decoder-only、encoder-only和encoder-decoder三种情况。

2. 模型forward时只返回loss，不直接返回logit，这是因为管道、张量和数据并行之间存在相当复杂的相互作用。这些损失是在数据并行排名中进行了平均。对于大多数情况，特别是在使用Megatron-LM特性运行预训练任务时，可以使用损失来轻松计算困惑度。对于GPT模型，除了损失外，还支持返回logits。这些logits并未在数据并行排名中被收集，可以使用 accelerator.utils.gather_across_data_parallel_groups 来收集跨数据并行排名的logits，然后与标签一起用于计算各种性能指标。

3. 主进程为最后一个流水线阶段rank,is_main_process等接口按此判断。

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

5. gradient_accumulation_steps需要为1。Megatron中的micro batch类似累积梯度。

6. 保存和加载检查点用save_state和load_state接口。

7. 支持Megatron和Transformers三种模型的对应关系，而其他 🤗 transformers 的模型架构则不受支持

   • Megatron-LM BertModel 对应于 🤗 transformers 中具有 megatron-bert 模型类型的模型，如 MegatronBERT。
   • Megatron-LM GPTModel 对应于 🤗 transformers 中具有 gpt2 模型类型的模型，如 OpenAI GPT2。
   • Megatron-LM T5Model 对应于 🤗 transformers 中具有 t5 模型类型的模型，如 T5 和 MT5。