# Megatron-LM

[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)は、大規模なTransformer言語モデルの学習を可能にします。[GPT](https://huggingface.co/papers/2005.14165)（Decoder Only）、[BERT](https://huggingface.co/papers/1810.04805)（Encoder Only）、[T5](https://huggingface.co/papers/1910.10683)（Encoder-Decoder）などのTransformerベースの言語モデルの事前学習において、効率的なテンソル、パイプライン、シーケンスベースのモデル並列性を提供します。詳細な情報や内部の仕組みについては、GitHubの[リポジトリ](https://github.com/NVIDIA/Megatron-LM)を参照してください。

## 何が統合されていますか？

Accelerateは、BERT（エンコーダー）、GPT（デコーダー）、またはT5モデル（エンコーダーとデコーダー）の大規模な事前学習/ファインチューニングを可能にするため、Megatron-LMの以下の機能を統合しています：

a. **テンソル並列性 (TP)**: ノード内のランク間で追加の通信をほとんど増やさずにメモリフットプリントを削減します。各テンソルは複数のチャンクに分割され、各シャードは別の GPU に配置されます。各ステップで、同じミニバッチのデータが各シャードによって独立かつ並列に処理され、その後、すべての GPU 間で同期 (`all-reduce` 操作) が行われます。単純な Transformer レイヤーでは、これはフォワードパスで 2 回の `all-reduces`、バックワードパスで 2 回の `all-reduces` を生じさせます。詳細については、研究論文 [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://huggingface.co/papers/1909.08053) およびブログ記事のこのセクション [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism) を参照してください。

b. **Pipeline Parallelism (PP)**: メモリフットプリントを削減し、ノード間並列化による大規模な学習を可能にします。PipeDream-Flush スケジュール/1F1B スケジュールおよび Interleaved 1F1B スケジュールにより、素朴な PP のバブルを削減します。レイヤーは PP ステージ間で均等に分散されます。例えば、モデルが`24`層を持ち、パイプライン並列化に`4`個の GPU を使用する場合、各 GPU には`6`層（24/4）が割り当てられます。PP のアイドル時間を削減するためのスケジュールの詳細については、研究論文[Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://huggingface.co/papers/2104.04473)およびブログ記事のこのセクション[The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism)を参照してください。

c. **シーケンス並列性 (SP)**: 追加の通信なしでメモリフットプリントを削減します。TP を使用する場合にのみ適用できます。`all-reduce` 後にテンソル並列ランク上で同じコピーが保持されるのを防ぎ、代わりに `reduce-scatter` と `no-op` 操作を `all-gather` に置き換えることで、必要なアクティベーションメモリを削減します。`all-reduce = reduce-scatter + all-gather` であるため、追加の通信コストなしで大量のアクティベーションメモリを節約できます。簡単に言うと、各 Transformer レイヤーの出力をシーケンス次元に沿ってシャードします。例えば、シーケンス長が `1024` で TP サイズが `4` の場合、各 GPU は各サンプルに対して `256` トークン (1024/4) を保持します。これにより、トレーニングでサポートできるバッチサイズが増加します。詳細については、研究論文 [Reducing Activation Recomputation in Large Transformer Models](https://huggingface.co/papers/2205.05198) を参照してください。

d. **Distributed Optimizer による Data Parallelism (DP)**: DP ランク間でオプティマイザの状態と勾配をシャーディングすることで、メモリ使用量を削減します（従来のデータ並列ランク間でオプティマイザの状態を複製する方法とは対照的です）。例えば、混合精度学習で Adam オプティマイザを使用する場合、各パラメータは 12 バイトのメモリを占有します。これは GPU 間で均等に分散されるため、4 つの GPU を使用している場合、各パラメータは 3 バイト（12/4）を占有することになります。詳細については、研究論文 [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://huggingface.co/papers/1910.02054) およびブログ [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#zero-data-parallelism) の該当セクションを参照してください。

e. **Expert Parallelism (EP)** Megatron-LM における Expert parallelism は、Mixture-of-Experts (MoE) レイヤーに使用されます。MoE レイヤーでは多数の「experts」（小さなフィードフォワードネットワーク）が存在しますが、各トークンに対してはごく少数の expert だけが活性化されます。すべての expert をすべての GPU に配置するのではなく、Megatron は異なる expert を異なる GPU に分散配置します。これが expert parallelism です。学習中は、トークンが選択された expert をホストする GPU にルーティングされ、そこで計算された後、元の場所に戻されます。これによりメモリコストが削減されます。大規模モデルでは、tensor/pipeline parallelism と組み合わせて使用されることがよくあります。 f. **Full Activation Recomputation**: スマートな activation checkpointing により、activation のメモリフットプリントを大幅に削減します。メモリを大量に消費するものの再計算が高速な activation を保存しないことで、メモリと再計算の間に優れたトレードオフを実現します。例えば、GPT-3 の場合、activation の再計算に伴う FLOPs のオーバーヘッドはわずか 2.7% であるにもかかわらず、activation に必要なメモリを 70% 削減できます。詳細については、研究論文 [Reducing Activation Recomputation in Large Transformer Models](https://huggingface.co/papers/2205.05198) を参照してください。

g. **Fused Kernels**: Fused Softmax、混合精度 Fused Layer Norm、および線形層の重み勾配計算への Fused gradient accumulation。PyTorch JIT によるコンパイル済み Fused GeLU、および Fused Bias+Dropout+Residual addition。

h. **インデックス付きデータセットのサポート**: 大規模トレーニング向けの効率的なデータセットのバイナリ形式。`mmap`、`cached`インデックスファイル、および`lazy`ローダー形式のサポート。

i. **チェックポイントのリシェイプと相互運用性**: 可変なテンソル並列サイズおよびパイプライン並列サイズを持つ Megatron-LM のチェックポイントを、Accelerate Big Model Inference や Megatron-DeepSpeed Inference など多数のツールで優れたサポートを受ける Transformers のシャーディング済みチェックポイントへ変形するためのユーティリティ。大規模トレーニング向けに、Transformers のシャーディング済みチェックポイントを可変なテンソル並列サイズおよびパイプライン並列サイズを持つ Megatron-LM のチェックポイントへ変換するサポートも利用可能です。

## 前提条件 

最新のpytorch、cuda、nccl、NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start)リリースおよびnltkライブラリのインストールが必要です。詳細は[ドキュメント](https://github.com/NVIDIA/Megatron-LM#setup)を参照してください。環境をセットアップする別の方法として、NGCから必要なインストールがすべて含まれたNVIDIA PyTorch Containerを取得する方法があります。

以下は、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 9a1c0d05c992c8a241da384ab27dce2021bb56dd
you need to manually move gpt_builders.py to megatron/training and update
include = [
    "megatron.core", 
    "megatron.core.*",
    "megatron.training",
    "megatron.training.*",
    "megatron.legacy",
    "megatron.legacy.*",
]
in pyproject.toml file to unblock yourself from using Megatron
pip install --no-use-pep517 -e .
```

## Megaton-LMのチェックポイントを準備する
モデルのファインチューニングを行う場合は、torch dist形式のチェックポイントが用意されていることを確認してください。Hugging Faceのモデルにのみアクセスできる場合は、Megatronが受け入れるtorch dist形式のチェックポイントへの変換を検討してください。例として、slimeのスクリプトを使用する方法を挙げます。GLMモデルを例に示します:
```
source /your/path/to/slime/scripts/models/glm4.5-355B-A32B.sh
srun torchrun --nproc-per-node 8 \
   /your/path/to/slime/tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /your/path/to/huggingface/models/GLM4.5-355B-A32B \
    --save /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist

```
変換後、以下の点を確認してください。1. `/your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist` 配下で、`latest_checkpointed_iteration.txt` の中身を `release` から `0` に変更し、ディレクトリ `release` を `iter_0000000` にリネームします。2. 設定ファイルで、オプティマイザの状態が不要になるよう `megatron_lm_no_load_optim` が true になっていることを確認します。

## Accelerate Megatron-LM プラグイン

重要な機能は、`accelerate config` コマンドによって直接サポートされています。Megatron-LM の機能を使用する場合の、対応する質問の例を以下に示します。

```bash
:~$ 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`を作成する必要があります。例を以下に示します。

```python
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. 総バッチサイズの詳細を取得するには、テンソル並列サイズとパイプライン並列サイズを考慮する必要があります。
有効な総バッチサイズを取得する例を以下に示します。

```python
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を使用する場合、損失はすでにデータ並列グループ内で平均化されています

```python
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` を使用してモデルを保存する必要があります。

```python
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` モデルアーキテクチャで実行してみましょう。

```bash
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"
```

以下は、出力ログからの重要な抜粋です：

```bash
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 Indexed Datasetsを活用するための高度な機能

より多くの機能を活用するには、以下の詳細をご確認ください。

1. 以下は、Megatron-LM を使用しながら Train Step をカスタマイズするために必要な変更の例です。
`accelerate.utils.AbstractTrainStep`を実装するか、対応する子クラスである`accelerate.utils.GPTTrainStep`、`accelerate.utils.BertTrainStep`、`accelerate.utils.T5TrainStep`を継承します。

```python
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のデータセットを使用するには、さらにいくつかの変更が必要です。これらのデータセット用のDataloader
は、各テンソル並列グループのランク0でのみ利用可能です。そのため、dataloaderが利用できないランクが存在し、トレーニングループの調整が必要です。これらを実現できることは、Accelerateがどれだけ柔軟で拡張性が高いかを示しています。必要な変更は以下の通りです。

a. Megatron-LMのインデックス付きデータセットを使用する場合は、`MegatronLMDummyDataLoader`を使用し、`data_path`や`seq_length`などの必要なデータセット引数を渡す必要があります。利用可能な引数の一覧は[こちら](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py)を参照してください。
    
```python
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` の比率に従って、学習用、検証用、テスト用のデータローダーを取得します。
    
```python
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の構成に似ており、Megaton-LMのインデックス付きデータセットを使用する際は、ユーザーが`max_train_steps`を指定する必要があります。これにより、Accelerateがどれだけ柔軟で拡張性が高いかがわかります。

```python
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. これらのスクリプトは、Transformersライブラリ内の各モデルの下に配置されています。
Currently, it is available for GPT model [checkpoint_reshaping_and_interoperability.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py)

2. 以下は、Megatron-LMのチェックポイントを汎用的なTransformersシャード済みチェックポイントに変換する例です。
```bash
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. `tp_size=2`、`pp_size=2`、`dp_size=2` で、transformers から megatron へのチェックポイントの変換。
```bash
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`を設定する必要があります。
これらはパイプラインの最終段階で利用可能になります。
```python
megatron_lm_plugin = MegatronLMPlugin(return_logits=True)
```

2. Megatron-LM GPTモデルの`megatron_generate`メソッド: Tensor ParallelismとPipeline Parallelismを使用して処理を完了します
`megatron_generate`メソッドは、greedy（top_k/top_pサンプリングあり/なし）を使用する際のバッチ入力に対する生成、およびビームサーチデコーディングを使用する際の個別プロンプト入力に対する生成をサポートします。transformersのgenerateの機能の一部のみがサポートされています。これにより、テンソル並列およびパイプライン並列を用いた大規模モデルの生成が容易になります（キーバリューキャッシュはすでに実装されており、デフォルトで融合カーネルを使用します）。この機能を使用するには、データ並列サイズが1であること、シーケンス並列およびアクティベーションチェックポイントが無効化されていることが必要です。また、トークナイザのvocabファイルおよびmergesファイルへのパスを指定する必要があります。以下の例では、Megatron-LM GPTモデルに対して`megatron_generate`メソッドを設定して使用する方法を示しています。
```python
# 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-LM GPTモデルの`megatron_generate`メソッドを使用するエンドツーエンドの例は、
設定ファイル [megatron_lm_gpt_generate_config.yaml](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/Configs/megatron_lm_gpt_generate_config.yaml) を使用した [megatron_gpt2_generation.py](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/inference/megatron_gpt2_generation.py)。accelerate launch コマンドを含む bash スクリプトは [megatron_lm_gpt_generate.sh](https://github.com/pacman100/accelerate-megatron-test/blob/main/megatron_lm_gpt_generate.sh) で利用できます。スクリプトの出力ログは [megatron_lm_gpt_generate.log](https://github.com/pacman100/accelerate-megatron-test/blob/main/output_logs/megatron_lm_gpt_generate.log) で利用できます。

## ROPEおよびALiBi位置埋め込み、マルチクエリアテンションのサポート

1. ROPE/ALiBi アテンションを使用する場合は、`MegatronLMPlugin` に `position_embedding_type` を `("absolute" | "rotary" | "alibi")` として渡します。以下に例を示します。
```python
other_megatron_args = {"position_embedding_type": "alibi"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
```

2. Multi-Query Attention の場合は、`MegatronLMPlugin` に `attention_head_type` を `("multihead" | "multiquery")` として渡します。以下に例を示します。
```python
other_megatron_args = {"attention_head_type": "multiquery"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
```

## 注意事項

1. TransformersのGPT2、Megatron-BERT、T5モデルをサポートします。
これは、Decoder only、Encode only、および Encoder-Decoder のモデルクラスを扱います。

2. モデルのフォワードパスから返されるのは、lossのみです。
裏側では、パイプライン並列性、テンソル並列性、データ並列性の複雑な相互作用が存在します。`model(**batch_data)` の呼び出しは、データ並列ランク全体で平均化されたロス（複数可）を返します。これは、Megatron-LM の機能を使用して事前学習ジョブを実行するほとんどのケースで問題ありません。この場合、ロスを使用して`perplexity`を簡単に計算できます。GPT モデルでは、ロス（複数可）に加えて logits を返すこともサポートされています。これらの logits は、データ並列ランク間で集約されていません。データ並列ランク間で logits を集約するには、`accelerator.utils.gather_across_data_parallel_groups`を使用してください。これらの logits とラベルを使用して、さまざまなパフォーマンス指標を計算できます。

3. 損失/logitsがパイプラインの最終ステージで利用可能であるため、メインプロセスは最後のランクです。
Megatron-LM統合を使用する場合、`accelerator.is_main_process`および`accelerator.is_local_main_process`は最後のランクに対して`True`を返します。

4. `accelerator.prepare` の呼び出しにおいて、指定された Transformers モデルに対応する Megatron-LM モデルが作成されます
ランダムな重み付きです。TP、PP、DPのパーティションが一致するMegatron-LMのチェックポイントをロードするには、`accelerator.load_state`を使用してください。

5. 現在、チェックポイントのリシェイプおよび相互運用性のサポートは、GPT のみで利用可能です。
まもなく、BERTおよびT5にも拡張される予定です。

6. `gradient_accumulation_steps` は 1 である必要があります。Megatron-LM を使用する場合、パイプライン並列性におけるマイクロバッチ
setting は勾配蓄積と同義です。

7. Megatron-LMを使用する場合は、チェックポイントの保存と読み込みに`accelerator.save_state`と`accelerator.load_state`を使用してください。

8. 以下は、Megatron-LMのモデルアーキテクチャから、同等のTransformersモデルアーキテクチャへのマッピングです。
サポートされているのは、これらのtransformersモデルアーキテクチャのみです。

a. Megatron-LM [BertModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/bert/bert_model.py) : config の model type に `megatron-bert` を含む transformers モデル。例: [MegatronBERT](https://huggingface.co/docs/transformers/model_doc/megatron-bert)
    
b. Megatron-LM [GPTModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/gpt/gpt_model.py): config の model type に `gpt2` を含む transformers モデル。例: [OpenAI GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
   
c. Megatron-LM [T5Model](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/T5/t5_model.py): config の model type に `t5` を含む transformers モデル。例: [T5](https://huggingface.co/docs/transformers/model_doc/t5)、[MT5](https://huggingface.co/docs/transformers/model_doc/mt5)