すべてのプロダクト
Search
ドキュメントセンター

Platform For AI:ユースケース:トランスフォーマーモデルのトレーニング高速化

最終更新日:Mar 11, 2026

PAI-Rapidformer のブラックボックス方式(CLI ベース)またはホワイトボックス方式(コードテンプレート)を用いて、PyTorch トランスフォーマーモデルのトレーニングを高速化します。

前提条件

高速化手法

Rapidformer では、以下の 2 種類の高速化手法をサポートしています:

ブラックボックス:Hugging Face モデルのファインチューニング

  1. Hugging Face にデータセットを登録するか、既存のデータセットを使用します。--dataset-name オプションで Rapidformer に渡します。

    詳細については、「Hugging Face データセットの登録」および「既存の Hugging Face データセット一覧の照会」をご参照ください。

  2. Hugging Face にモデルを登録するか、既存のモデルを使用します。--pretrained-model-name-or-path オプションで Rapidformer に渡します。

    詳細については、「Hugging Face モデルの登録」および「既存の Hugging Face モデル一覧の照会」をご参照ください。

  3. Rapidformer CLI の設定を行い、トレーニングを開始します。

    #!/bin/bash
    export CUDA_VISIBLE_DEVICES=4,5,6,7
    export MASTER_ADDR=localhost
    export MASTER_PORT=6010
    export NNODES=1
    export NODE_RANK=0
    
    rapidformer --task sequence_classification \ # タスク名
                --pretrained-model-name-or-path 'bert-base-cased' \  # 登録済みモデル名
                --data-path glue \                      # 登録済みデータパス名
                --data-name mrpc \                      # 登録済みデータファイル名
                --epochs 3 \                               # トレーニング反復回数
                --micro-batch-size 16 \                    # 各 GPU 上のバッチサイズ
                --global-batch-size 64 \                   # 分散トレーニング全体のバッチサイズ
                --lr 2e-5 \                                # 学習率
                --lr-decay-style linear \                  # 学習率減衰ポリシー
                --lr-warmup-iters 100 \                    # 学習率ウォームアップステップ数
                --weight-decay 1e-2 \                      # 学習率係数
                --clip-grad 1.0 \                          # 勾配クリップ係数
                --seed 42 \                                # 乱数シード
                --mixed-precision \                        # 混合精度トレーニングを有効化
                --onnx-runtime-training \                  # 計算グラフ最適化を有効化
                --zero-1-memory-optimization \             # オプティマイザ状態のパーティショニングを有効化

    各パラメータの詳細については、「パラメータ設定ガイド」をご参照ください。

ブラックボックス:Hugging Face モデルの事前学習

  1. 事前学習用の mmap 形式データセットを作成します。

    詳細については、「Megatron データ処理スクリプト」をご参照ください。mmap データセット作成コマンドの例:

    python preprocess_data.py \
      --input book_wiki_owtv2_small.json  \
      --output-prefix gpt_small \
      --vocab gpt2-vocab.json \
      --dataset-impl mmap \
      --tokenizer-type GPT2BPETokenizer \
      --merge-file gpt2-merges.txt \
      --append-eod
  2. Hugging Face にモデルを登録するか、既存のモデルを使用します。--pretrained-model-name-or-path オプションで Rapidformer に渡します。

    詳細については、「Hugging Face モデルの登録」および「既存の Hugging Face モデル一覧の照会」をご参照ください。

  3. Rapidformer CLI の設定を行い、トレーニングを開始します。

    #!/bin/bash
    export CUDA_VISIBLE_DEVICES=4,5,6,7
    export MASTER_ADDR=localhost
    export MASTER_PORT=6010
    export NNODES=1
    export NODE_RANK=0
    
    rapidformer --task pretraining \
           --pretrained-model-name-or-path 'bert-base-uncased' \
           --num-layers 12 \
           --hidden-size 768 \
           --num-attention-heads 12 \
           --micro-batch-size 16 \
           --global-batch-size 128 \               # 勾配蓄積を有効化
           --seq-length 512 \
           --tokenizer-type BertWordPieceLowerCase \
           --max-position-embeddings 512 \
           --train-iters 100 \
           --data-path book_wiki_owtv2_small_text_sentence \
           --vocab-file bert-en-uncased-vocab.txt  \
           --data-impl mmap \
           --split 980,20 \
           --lr 1e-3 \
           --lr-decay-style linear \
           --min-lr 0.0 \
           --lr-decay-iters 2000 \
           --weight-decay 1e-2 \
           --clip-grad 1.0 \
           --lr-warmup-fraction .01 \
           --mixed-precision \                    # 混合精度トレーニングを有効化
           --onnx-runtime-training \              # 計算グラフ最適化を有効化
           --fsdp-memory-optimization \           # モデル状態のパーティショニングを有効化

    各パラメータの詳細については、「パラメータ設定ガイド」をご参照ください。

ホワイトボックス:Finetuner テンプレートを用いた Hugging Face モデルのファインチューニング

Rapidformer の Finetuner コードテンプレートを用いると、Hugging Face モデルのファインチューニングジョブを迅速に作成できます。このテンプレートには、以下の 4 つの関数が含まれています:

  • train_valid_test_datasets_provider:データセットの作成

  • model_optimizer_lr_scheduler_provider:モデル、オプティマイザー、および学習率スケジューラの構築

  • run_forward_step:順伝播処理ロジックの定義

  • run_compute_metrics:トレーニングおよび評価精度の計算

これらの関数の詳細については、「Rapidformer API」をご参照ください。入力および出力の概要:

class MyFintuner(Finetuner):

    def __init__(self, engine):
        super().__init__(engine=engine)

    # トレーニング/検証/テストデータセットの取得
    # 入力:なし
    # 出力:3 つのオブジェクトと 1 つのオブジェクト関数
    def train_valid_test_datasets_provider(self):

        return train_dataset, valid_dataset, test_dataset, collate_f

    # モデル/オプティマイザー/学習率スケジューラの作成
    # 入力:なし
    # 出力:3 つのオブジェクト
    def model_optimizer_lr_scheduler_provider(self):

        return model, optimizer, lr_scheduler

    # 順伝播ロジックの記述
    # 入力:バッチまたはイテレータ、モデル
    # 出力:損失
    def run_forward_step(self, batch_or_iterator, model):
        return loss

    # ファインチューニング専用の検証データセット評価ロジックの記述
    # 入力:モデル、検証データセットのデータローダー
    # 出力:メトリックオブジェクト
    def run_compute_metrics(self, model, eval_dataloader):
        return metric
                

カスタムコードテンプレートの内容を理解した後、ブラックボックス:Hugging Face モデルのファインチューニング に従ってデータセットおよびモデルを準備します。その後、以下の手順を実行します:

  1. Rapidformer および Hugging Face のインターフェイスをインポートします。

    from transformers.easytexmier import AutoConfig, BertForSequenceClassification
    from datasets import load_dataset, load_metric
    from rapidformer import RapidformerEngine
    from rapidformer import get_args
    from rapidformer import get_logger
    from rapidformer import get_timers
    from rapidformer import Finetuner
    from rapidformer import Pretrainer
    from rapidformer import build_train_valid_test_datasets_for_huggingface
  2. コードテンプレート内の 4 つの関数を完成させます。

    class MyFintuner(Finetuner):
        def __init__(self,engine):
            super().__init__(engine=engine)
    
        def train_valid_test_datasets_provider(self):
            tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
    
            def tokenize_function(examples):
                # max_length=None => モデルの最大長を使用(実際にはデフォルト)
                outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
                return outputs
    
            datasets = load_dataset(args.dataset_path, args.dataset_name)
            # データセットのすべての分割に対して、先ほど定義したメソッドを適用
            tokenized_datasets = datasets.map(
                tokenize_function,
                batched=True,
                remove_columns=["idx", "sentence1", "sentence2"],
            )
            tokenized_datasets.rename_column_("label", "labels")
    
            train_dataset = tokenized_datasets["train"]
            valid_dataset = tokenized_datasets['validation']
            test_dataset = tokenized_datasets['test']
    
            def collate_fn(examples):
                return tokenizer.pad(examples, padding="longest", return_tensors="pt")
    
            return train_dataset, valid_dataset, test_dataset, collate_fn
    
        def model_optimizer_lr_scheduler_provider(self):
            args = get_args()
            model = BertForSequenceClassification.from_pretrained(args.load)
            return model, None, None
    
        def run_forward_step(self, batch, model):
            output_tensor = model(**batch)
            return output_tensor.loss
    
        # 各エポック終了後に検証データセットでメトリックを実行
        def run_compute_metrics(self, model, eval_dataloader):
            model = model[0]
            metric = load_metric(args.dataset_path, args.dataset_name)
            for step, batch in enumerate(eval_dataloader):
                with torch.no_grad():
                    outputs = model(**batch)
                predictions = outputs.logits.argmax(dim=-1)
    
                metric.add_batch(
                    predictions=self.gather(predictions),
                    references=self.gather(batch["labels"]),
                )
    
            eval_metric = metric.compute()
            return eval_metric
                            
  3. Rapidformer エンジンを初期化し、トレーナーオブジェクトを作成して finetune() を呼び出し、rapidformer_finetune_hugging_face_bert_trainer.py として保存します。

    engine = RapidformerEngine()
    trainer = MyFintuner(engine=engine)
    trainer.train()
  4. CLI をベースに起動スクリプトを準備します。--user-scriptrapidformer_finetune_hugging_face_bert_trainer.py に設定し、高速化スイッチを有効化します。

    #!/bin/bash
    export CUDA_VISIBLE_DEVICES=4,5,6,7
    export MASTER_ADDR=localhost
    export MASTER_PORT=6010
    export NNODES=1
    export NODE_RANK=0
    
    rapidformer --user-script rapidformer_finetune_huggingface_bert_trainer.py
                --task sequence_classification \
                --pretrained-model-name-or-path 'bert-base-cased' \
                --data-path glue \
                --data-name mrpc \
                --epochs 3 \
                --micro-batch-size 16 \
                --global-batch-size 16 \
                --lr 2e-5 \
                --lr-decay-style linear \
                --lr-warmup-iters 100 \
                --weight-decay 1e-2 \
                --clip-grad 1.0 \
                --mixed-precision                                 # 混合精度トレーニングを有効化
                --zero-3-memory-optimization \                    # モデル状態のパーティショニングを有効化
                --onnx-runtime-training \                         # 計算グラフ最適化を有効化

ホワイトボックス:Pretrainer テンプレートを用いた Hugging Face モデルの事前学習

Rapidformer の Pretrainer コードテンプレートを用いると、Hugging Face モデルの事前学習タスクを迅速に作成できます。このテンプレートには、以下の関数が含まれています:

  • train_valid_test_datasets_provider:データセットの作成

  • model_optimizer_lr_scheduler_provider:モデル、オプティマイザー、および学習率スケジューラの構築

  • run_forward_step:順伝播処理ロジックの定義

これらの関数の詳細については、「Rapidformer API」をご参照ください。入力および出力の詳細については、「ホワイトボックス:Finetuner テンプレートを用いた Hugging Face モデルのファインチューニング」をご参照ください。

カスタムコードテンプレートの内容を理解した後、ブラックボックス:Hugging Face モデルのファインチューニング に従ってデータセットおよびモデルを準備します。その後、以下の手順を実行します:

  1. Rapidformer および Hugging Face のインターフェイスをインポートします。

    説明

    事前学習ではイテレータを用いてデータを読み込むため、データ並列処理のために mpu をインポートします。

    from megatron import mpu
    from transformers import BertConfig, BertForPreTraining
    from rapidformer import RapidformerEngine, get_args, PreTrainer
    from rapidformer import build_train_valid_test_datasets_for_huggingface
  2. Pretrainer を継承し、事前学習コードを完成させます。

    class MyBertPreTrainer(PreTrainer):
    
        def __init__(self,engine):
            super().__init__(engine=engine)
    
        def train_valid_test_datasets_provider(self, train_val_test_num_samples):
            args = get_args()
    
            train_ds, valid_ds, test_ds = build_train_valid_test_datasets_for_huggingface(
                data_prefix=args.data_path,
                data_impl=args.data_impl,
                splits_string=args.split,
                train_valid_test_num_samples=train_val_test_num_samples,
                max_seq_length=args.seq_length,
                masked_lm_prob=args.mask_prob,
                short_seq_prob=args.short_seq_prob,
                seed=args.seed,
                skip_warmup=(not args.mmap_warmup),
                binary_head=True)
    
            return train_ds, valid_ds, test_ds
    
        def model_optimizer_lr_scheduler_provider(self):
            args = get_args()
            model = AutoModelForPreTraining.from_pretrained(args.pretrained_model_name_or_path)
            return model, None, None
    
        def run_forward_step(self, data_iterator, model):
            # アイテムとその型。
            keys = ['input_ids', 'attention_mask', 'token_type_ids', 'labels', 'next_sentence_label']
            datatype = torch.int64
    
            # データのブロードキャスト。
            if data_iterator is not None:
                data = next(data_iterator)
            else:
                data = None
            data_b = mpu.broadcast_data(keys, data, datatype)
            input_ids = data_b['input_ids'].long()
            attention_mask = data_b['attention_mask'].long()
            token_type_ids = data_b['token_type_ids'].long()
            labels = data_b['labels'].long()
            next_sentence_label = data_b['next_sentence_label'].long()
            output_tensor = model(input_ids=input_ids, attention_mask=attention_mask,
                                  token_type_ids=token_type_ids, labels=labels, next_sentence_label=next_sentence_label)
    
            return output_tensor['loss']
  3. Rapidformer エンジンを初期化し、トレーナーオブジェクトを作成して pretrain() を呼び出し、rapidformer_pretrain_hugging_face_bert_trainer.py として保存します。

    engine = RapidformerEngine()
    trainer = MyBertPreTrainer(engine=engine)
    trainer.train()
  4. CLI をベースに起動スクリプトを準備し、高速化スイッチを有効化します。

    #!/bin/bash
    export CUDA_VISIBLE_DEVICES=4,5,6,7
    export MASTER_ADDR=localhost
    export MASTER_PORT=6010
    export NNODES=1
    export NODE_RANK=0
    
    DATA_PATH=book_wiki_owtv2_small_text_sentence
    
    rapidformer --user-script rapidformer_pretrain_huggingface_bert_trainer.py \
           --pretrained-model-name-or-path 'bert-base-uncased' \
           --num-layers 12 \
           --hidden-size 768 \
           --num-attention-heads 12 \
           --micro-batch-size 16 \
           --global-batch-size 64 \
           --seq-length 512 \
           --tokenizer-type BertWordPieceLowerCase \
           --max-position-embeddings 512 \
           --train-iters 100 \
           --data-path $DATA_PATH \
           --vocab-file bert-en-uncased-vocab.txt  \
           --data-impl mmap \                               # データ高速化を有効化
           --split 980,20 \
           --lr 1e-3 \
           --lr-decay-style linear \
           --weight-decay 1e-2 \
           --clip-grad 1.0 \
           --lr-warmup-fraction .01 \
           --zero-3-memory-optimization \                    # モデル状態のパーティショニングを有効化
           --onnx-runtime-training \                         # 計算グラフ最適化を有効化
           --mixed-precision                                 # 混合精度トレーニング

ホワイトボックス:カスタム Trainer を用いた Hugging Face モデルのファインチューニング

カスタム Trainer を使用するプログラムの場合、Rapidformer による高速化は限定的です(Apex オプティマイザー、モデル状態のパーティショニング、計算グラフ最適化のみ)。混合精度トレーニングには多数の修正が必要です。より良い結果を得るには、前述のテンプレートベースの手法をご利用ください。本セクションでは、典型的な Hugging Face ファインチューニングコードに対する侵襲的な高速化を示します。

Hugging Face ファインチューニングコードの例:

import torch
from datasets import load_dataset, load_metric
from torch.utils.data import DataLoader
from transformers import (
    AdamW,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    BertForSequenceClassification,

)

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
datasets = load_dataset("glue", "mrpc")
metric = load_metric("glue", "mrpc")

def tokenize_function(examples):
    # max_length=None => use the model max length (it's actually the default)
    outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
    return outputs

tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    remove_columns=["idx", "sentence1", "sentence2"],
)

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)

optimizer = AdamW(params=model.parameters(), lr=args.lr, correct_bias=True)

lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=args.lr_warmup_iters,
    num_training_steps=args.train_iters
)

device = torch.device("cuda", args.local_rank)

for epoch in range(args.epochs):
    model.train()
    for step, batch in enumerate(train_dataloader):
        batch.to(device)
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    model.eval()
    for step, batch in enumerate(eval_dataloader):
        batch.to(device)
        with torch.no_grad():
            outputs = model(**batch)
            predictions = outputs.logits.argmax(dim=-1)
            metric.add_batch(
                    predictions=engine.gather(predictions),
                    references=engine.gather(batch["labels"]))

     eval_metric = metric.compute()
     print("epoch {}: {}".format(epoch, eval_metric))

このコードには、データ並列処理のサポートがない、オプティマイザーが遅い、混合精度トレーニングが未実装などの課題があります。以下では、Rapidformer の API を用いてこのコードを修正します。

  1. データ並列処理のサポートを追加します。

    finetuner オブジェクトを作成し、finetuner.build_data_loader を呼び出してデータローダーを作成します。このローダーはデータ並列処理をサポートし、自動的にデータを GPU に送信します。元のコードから batch.to(device) を削除します。

    + from rapidformer import RapidformerEngine
    + engine = RapidformerEngine()
    + finetuner = Finetuner(engine=engine)
    
    - train_dataloader = DataLoader(tokenized_datasets["train"])
    - eval_dataloader = DataLoader(tokenized_datasets["train"])
    
    + train_dataloader = finetuner.build_data_loader(tokenized_datasets["train"])
    + eval_dataloader = finetuner.build_data_loader(tokenized_datasets["validation"])
  2. データ並列処理の上位に Apex オプティマイザーを適用します。

    オプティマイザーを高速な Apex Fused Adam に置き換えます。元のオプティマイザーを削除し、Rapidformer の Fused Adam を使用します。engine.compose を呼び出して、モデル、オプティマイザー、および学習率スケジューラをカプセル化します。

    + from rapidformer import RapidformerEngine
    + engine = RapidformerEngine()
    + finetuner = Finetuner(engine=engine)
    
    - optimizer = AdamW(params=model.parameters(), lr=args.lr, correct_bias=True)
    - lr_scheduler = get_linear_schedule_with_warmup(optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_iters,
        num_training_steps=args.train_iters
    )
    
    
    + lr_scheduler = partial(
            get_linear_schedule_with_warmup,
            num_warmup_steps=args.lr_warmup_iters,
            num_training_steps=args.train_iters
        )
    
    + model, optimizer, lr_scheduler = engine.compose(model_obj=model,
          lr_scheduler_fn=lr_scheduler)
    説明

    Apex オプティマイザーおよび混合精度をデータ並列処理と併用するのは複雑です。混合精度トレーニングでは、モデルを fp16 に切り替え、損失スケーリングを行う必要があります。トレーナーなしでフロントエンドプログラムを変更するのは非常に複雑です。代わりに、トレーナーに基づくソリューションをご利用ください。Rapidformer の Finetuner は、データ並列処理、Apex、PyTorch 混合精度トレーニング、Megatron オプティマイザーの混合精度トレーニング、FairScale および DeepSpeed による VRAM 最適化を統合しています。

ホワイトボックス:Pretrainer テンプレートを用いた Megatron モデルの事前学習

ホワイトボックス:カスタム Trainer を用いた Hugging Face モデルのファインチューニング の手法を理解した後、Data および Model Hub をバイパスすることで、さらに柔軟性を高めます。train_valid_test_datasets_provider でカスタムデータ作成ロジックを記述し、model_optimizer_lr_scheduler_provider でカスタムモデルを記述し、run_forward_step でカスタム順伝播ロジックを記述します。

  1. 事前学習用の mmap 形式データセットを作成します。

    詳細については、「Megatron データ処理スクリプト」をご参照ください。mmap データセット作成コマンドの例:

    python preprocess_data.py \
      --input /apsarapangu/disk2/jerry.lp/pretrain_datasets/en/book_wiki_owtv2_small.json  \
      --output-prefix /apsarapangu/disk2/jerry.lp/pretrain_datasets/en/gpt_small \
      --vocab gpt2-vocab.json \
      --dataset-impl mmap \
      --tokenizer-type GPT2BPETokenizer \
      --merge-file gpt2-merges.txt \
      --append-eod
  2. Pretrainer を継承し、事前学習コード内で train_valid_test_datasets_provider のカスタム関数を完成させます。

    サードパーティライブラリに依存せずに、トレーニング、検証、テストデータセットを作成するカスタムロジックを記述します。データセットは torch.utils.data.Dataset を継承する必要があります。

    from rapidformer import RapidformerEngine, get_args, PreTrainer
    
    class MegatronGPTPreTrainer(PreTrainer):
        def __init__(self,
                     engine,
                     ):
            super().__init__(engine=engine)
    
        def train_valid_test_datasets_provider(self, train_val_test_num_samples):
            args = get_args()
    
            train_ds, valid_ds, test_ds = build_train_valid_test_datasets(
                data_prefix=args.data_path,
                data_impl=args.data_impl,
                splits_string=args.split,
                train_valid_test_num_samples=train_val_test_num_samples,
                seq_length=args.seq_length,
                seed=args.seed,
                skip_warmup=(not args.mmap_warmup))
    
            return train_ds, valid_ds, test_ds
  3. Pretrainer を継承し、事前学習コード内で model_optimizer_lr_scheduler_provider のカスタム関数を完成させます。

    サードパーティライブラリに依存せずに、カスタムモデルオブジェクトを作成するカスタムロジックを記述します。モデルは torch.nn.Module を継承する必要があります。

    from rapidformer import RapidformerEngine, get_args, PreTrainer
    from yourmodel import GPTModel
    
    class MegatronGPTPreTrainer(PreTrainer):
        def __init__(self,
                     engine,
                     ):
            super().__init__(engine=engine)
    
        def model_optimizer_lr_scheduler_provider(self):
            model = GPTModel()
            return model, None, None
  4. Pretrainer を継承し、事前学習コード内で run_forward_step のカスタム関数を完成させます。

    from rapidformer import RapidformerEngine, get_args, PreTrainer
    
    class MyGPTPreTrainer(PreTrainer):
        def __init__(self,
                     engine,
                     ):
            super().__init__(engine=engine)
    
    
        def run_forward_step(self, data_iterator, model):
            """順伝播ステップ。"""
            args = get_args()
    
            tokenizer = get_tokenizer()
    
            # アイテムとその型。
            keys = ['text']
            datatype = torch.int64
    
            # データのブロードキャスト。
            if data_iterator is not None:
                data = next(data_iterator)
            else:
                data = None
            data_b = mpu.broadcast_data(keys, data, datatype)
    
            # アンパック。
            tokens_ = data_b['text'].long()
            labels = tokens_[:, 1:].contiguous()
            tokens = tokens_[:, :-1].contiguous()
    
            # マスクおよび位置 ID の取得。
            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
                tokens,
                tokenizer.eod,
                args.reset_position_ids,
                args.reset_attention_mask,
                args.eod_mask_loss)
    
            output_tensor = model(tokens, position_ids, attention_mask,
                                  labels=labels)
    
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
    
            return loss
    
    
                            
  5. Rapidformer エンジンを初期化し、トレーナーオブジェクトを作成して pretrain() を呼び出し、rapidformer_pretrain_megatron_gpt_trainer.py として保存します。

    engine = RapidformerEngine()
    trainer = MyGPTPreTrainer(engine=engine)
    trainer.train()
  6. 起動スクリプトを準備し、高速化スイッチを有効化します。

    #!/bin/bash
    export CUDA_VISIBLE_DEVICES=4,5,6,7
    export MASTER_ADDR=localhost
    export MASTER_PORT=6010
    export NNODES=1
    export NODE_RANK=0
    
    DATA_PATH=book_wiki_owtv2_small_text_sentence
    PRETRAINED_CHECKPOINT=
    
    rapidformer --user-script rapidformer_pretrain_megatron_gpt_trainer.py \
           --tensor-model-parallel-size 2 \          # 演算子分割最適化を有効化
           --pipeline-model-parallel-size 2 \        # パイプライン並列化最適化を有効化
           --num-layers 12 \
           --hidden-size 768 \
           --num-attention-heads 12 \
           --micro-batch-size 16 \
           --global-batch-size 128 \                  # 勾配蓄積最適化を有効化
           --seq-length 512 \
           --tokenizer-type GPT2BPETokenizer \
           --max-position-embeddings 512 \
           --train-iters 100 \
           --data-path $DATA_PATH \
           --vocab-file gpt2-vocab.json \
           --merge-file gpt2-merges.txt \
           --data-impl mmap \                         # データ高速化を有効化
           --split 980,20 \
           --lr 1e-3 \
           --lr-decay-style linear \
           --weight-decay 1e-2 \
           --clip-grad 1.0 \
           --lr-warmup-fraction .01 \
           --log-interval 1 \
           --zero-2-memory-optimization \              # モデル状態のパーティショニングを有効化
           --checkpoint-activations \                  # グラデーションチェックポイントを有効化
           --mixed-precision                           # 混合精度トレーニングを有効化