Rapidformer を使用して PyTorch Transformer モデルトレーニングを最適化する方法 - Platform For AI

前提条件

開始する前に、以下が完了していることを確認してください。

Rapidformer のランタイムイメージがインストールされていること。詳細については、「PAI-Megatron-Patch ランタイムイメージのインストール」をご参照ください。
Rapidformer のパラメーターリファレンスを確認済みであること。詳細については、「パラメーター設定ガイド」をご参照ください。
Rapidformer の API リファレンスを確認済みであること。詳細については、「Rapidformer API」をご参照ください。

高速化手法の選択

次に進む前に、この表を参考にして適切な手法を選択してください。

	ブラックボックス	ホワイトボックス
仕組み	モデルとデータセットの名前を Rapidformer CLI に渡します。コードの変更は不要です。	Rapidformer テンプレート (`Finetuner` または `PreTrainer`) をサブクラス化し、必要なメソッドを実装してから、`--user-script` を使用して CLI 経由で呼び出します。
使用するケース	モデルとデータセットが Hugging Face または Megatron に登録されており、トレーニングループをカスタマイズする必要がない場合。	データ読み込み、モデルアーキテクチャ、またはフォワードパスロジックを制御する必要がある場合、または完全にカスタマイズされた Megatron モデルを使用している場合。
カスタマイズ	なし	データ、モデル、オプティマイザ、フォワードパスを完全に制御
高速化機能	混合精度トレーニング、ONNX Runtime グラフ最適化、ZeRO メモリ最適化、FSDP	ブラックボックスと同じ機能に加え、テンソル並列処理とパイプライン並列処理 (Megatron のみ)

既存のカスタム Trainer を持つプログラムに対して、Rapidformer は限定的な侵入型高速化 (Apex オプティマイザ、モデル状態の分割、計算グラフの最適化のみ) を提供します。混合精度トレーニングには大幅な変更が必要です。完全な高速化を適用するには、代わりにホワイトボックステンプレート手法を使用してください。

高速化フラグのリファレンス

このドキュメントのすべての CLI 例では、これらのフラグの組み合わせを使用します。各フラグの機能と有効にするタイミングを理解するために、この表をご参照ください。

フラグ	高速化技術	機能	組み合わせるフラグ
`--mixed-precision`	混合精度トレーニング	FP16 でトレーニングを行い、メモリ使用量を削減し、計算を高速化します	他の任意のフラグ
`--onnx-runtime-training`	計算グラフの最適化	ONNX Runtime を使用して計算グラフを最適化します	他の任意のフラグ
`--zero-1-memory-optimization`	ZeRO Stage 1：オプティマイザ状態の分割	オプティマイザの状態を GPU 間で分割します	`--mixed-precision`
`--zero-2-memory-optimization`	ZeRO Stage 2：オプティマイザと勾配の分割	オプティマイザの状態と勾配を GPU 間で分割します	`--mixed-precision`
`--zero-3-memory-optimization`	ZeRO Stage 3：完全なモデル状態の分割	オプティマイザの状態、勾配、モデルパラメーターを分割します	`--mixed-precision`
`--fsdp-memory-optimization`	FSDP (完全シャード化データ並列)	PyTorch FSDP を使用してモデルの状態をシャード化します	`--mixed-precision`
`--tensor-model-parallel-size <N>`	テンソル並列処理	演算子を N 個の GPU に分割します	Megatron の事前学習のみ
`--pipeline-model-parallel-size <N>`	パイプライン並列処理	モデルのレイヤーを N 個の GPU に分割します	Megatron の事前学習のみ
`--checkpoint-activations`	勾配チェックポイント	バックワードパス中に活性化を再計算して VRAM を節約します	他の任意のフラグ
`--data-impl mmap`	メモリマップドデータ	メモリマップドファイルを使用してディスクからデータセットを読み込み、データ読み込みのオーバーヘッドを削減します	事前学習のみ

完全なパラメーターリファレンスについては、「パラメーター設定ガイド」をご参照ください。

ブラックボックス：Hugging Face のファインチューニング

コードは不要です。データセットとモデルを Hugging Face に登録し、CLI を実行します。この例では、混合精度トレーニング、ONNX Runtime グラフ最適化、および ZeRO Stage 1 オプティマイザ状態分割を有効にします。

データセットを Hugging Face に登録するか、既存のものを使用します。--dataset-name を使用してデータセットを Rapidformer に渡します。
- Hugging Face データセットの登録
- 既存の Hugging Face データセットの参照
モデルを Hugging Face に登録するか、既存のものを使用します。--pretrained-model-name-or-path を使用してモデルを Rapidformer に渡します。
- Hugging Face モデルの登録
- 既存の Hugging Face モデルの参照

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 \
            --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 の事前学習

コードは不要です。mmap データセットを作成し、モデルを登録してから CLI を実行します。この例では、メモリマップドデータ読み込み、混合精度トレーニング、ONNX Runtime グラフ最適化、および FSDP モデル状態分割を有効にします。

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

モデルを Hugging Face に登録するか、既存のものを使用します。--pretrained-model-name-or-path を使用してモデルを Rapidformer に渡します。
- Hugging Face モデルの登録
- 既存の Hugging Face モデルの参照

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 のファインチューニング

Finetuner をサブクラス化し、4つのメソッドを実装します。Rapidformer は分散トレーニング、混合精度、ZeRO、および FSDP を自動的に処理します。コードでは、データ読み込み、モデル構築、フォワードロジック、および評価メトリクスのみを定義します。

テンプレートの構造

Finetuner 基底クラスには 4 つのメソッドが必要です。

メソッド	入力	出力	目的
`train_valid_test_datasets_provider`	なし	`train_dataset, valid_dataset, test_dataset, collate_fn`	トレーニング、検証、テスト用のデータセットを作成します
`model_optimizer_lr_scheduler_provider`	なし	`model, optimizer, lr_scheduler`	モデル、オプティマイザ、および学習率スケジューラを構築します
`run_forward_step`	`batch_or_iterator, model`	`loss`	フォワードパスのロジックを定義します
`run_compute_metrics`	`model, eval_dataloader`	`metric`	評価メトリクスを計算します (ファインチューニングのみ)

完全なメソッドシグネチャと戻り値の型の詳細については、「Rapidformer API」をご参照ください。

実装手順

「ブラックボックス：Hugging Face のファインチューニング」と同じアプローチでデータセットとモデルを準備し、以下の手順を実行します。

Rapidformer と Hugging Face のインターフェースをインポートします。

from transformers import AutoConfig, AutoTokenizer, 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

Finetuner サブクラスで 4 つのメソッドを実装します。

class MyFintuner(Finetuner):
    def __init__(self, engine):
        super().__init__(engine=engine)

    def train_valid_test_datasets_provider(self):
        args = get_args()
        tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path)

        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.pretrained_model_name_or_path)
        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]
        args = get_args()
        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

エンジンを初期化し、トレーナーを作成し、rapidformer_finetune_huggingface_bert_trainer.py として保存します。
```
engine = RapidformerEngine()
trainer = MyFintuner(engine=engine)
trainer.train()
```

--user-script でステップ 3 のファイルを指定し、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 --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 の事前学習

PreTrainer をサブクラス化し、3つのメソッドを実装します。Finetuner とは異なり、事前学習ではイテレータを介してデータを読み込むため、データ並列処理のブロードキャストのために mpu をインポートします。

テンプレートの構造

PreTrainer 基底クラスには 3 つのメソッドが必要です。

メソッド	入力	出力	目的
`train_valid_test_datasets_provider`	`train_val_test_num_samples`	`train_ds, valid_ds, test_ds`	トレーニング、検証、テスト用のデータセットを作成します
`model_optimizer_lr_scheduler_provider`	なし	`model, optimizer, lr_scheduler`	モデル、オプティマイザ、および学習率スケジューラを構築します
`run_forward_step`	`data_iterator, model`	`loss`	フォワードパスのロジックを定義します (イテレータベース)

メソッドの戻り値の型の詳細については、「ホワイトボックス：Finetuner テンプレートを使用した Hugging Face のファインチューニング」をご参照ください。

実装手順

「ブラックボックス：Hugging Face の事前学習」と同様にデータセットとモデルを準備し、以下の手順を実行します。

Rapidformer と Hugging Face のインターフェースをインポートします。

事前学習では、イテレータを介してデータを読み込みます。並列ワーカー間でデータをブロードキャストするために mpu をインポートします。
```
from megatron import mpu
from transformers import AutoModelForPreTraining, BertConfig, BertForPreTraining
from rapidformer import RapidformerEngine, get_args, PreTrainer
from rapidformer import build_train_valid_test_datasets_for_huggingface
```

PreTrainer サブクラスで 3 つのメソッドを実装します。

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

        data = next(data_iterator) if data_iterator is not None else 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']

エンジンを初期化し、トレーナーを作成し、rapidformer_pretrain_huggingface_bert_trainer.py として保存します。
```
engine = RapidformerEngine()
trainer = MyBertPreTrainer(engine=engine)
trainer.train()
```

--user-script でステップ 3 のファイルを指定し、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 を使用する既存の Hugging Face トレーニングコードには、Rapidformer API を段階的に追加します。このセクションでは、最も影響の大きい2つの変更点、データ並列処理と高速なオプティマイザについて説明します。

この方法は限定的な高速化 (Apex オプティマイザ、モデル状態の分割、計算グラフの最適化のみ) を提供します。混合精度トレーニングには、トレーニングループへの大幅な変更が必要です。完全な高速化のためには、代わりにFinetuner テンプレート手法を使用してください。

データ並列処理の追加

標準の DataLoader を finetuner.build_data_loader に置き換えます。このローダーはデータ並列処理をサポートし、バッチを自動的に GPU に移動させるため、トレーニングループから batch.to(device) を削除します。

+ from rapidformer import RapidformerEngine, Finetuner
+ 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"])

Apex Fused Adam オプティマイザの使用

AdamW を Rapidformer の Apex Fused Adam オプティマイザに置き換えます。engine.compose を呼び出して、モデル、オプティマイザ、学習率スケジューラを一緒にラップします。

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

カスタム Trainer ループで Apex オプティマイザと混合精度トレーニングを組み合わせるには、モデルを FP16 に切り替え、損失スケーリングを追加する必要があります。これは大幅な変更です。Finetuner テンプレートは、データ並列処理、Apex、PyTorch 混合精度トレーニング、Megatron オプティマイザ混合精度、および FairScale と DeepSpeed による VRAM 最適化を自動的に統合します。

ホワイトボックス：PreTrainer テンプレートを使用した Megatron の事前学習

Hugging Face や Data/Model Hub を使用しない、完全にカスタマイズされた Megatron モデルの場合、3 つすべての PreTrainer メソッドをカスタムロジックで実装します。データセットは torch.utils.data.Dataset から継承する必要があり、モデルは torch.nn.Module から継承する必要があります。

この方法では、Megatron モデルでのみ利用可能なテンソル並列処理とパイプライン並列処理も有効にできます。

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

カスタムデータセット作成ロジックで train_valid_test_datasets_provider を実装します。

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

カスタムモデルで model_optimizer_lr_scheduler_provider を実装します。

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

カスタムフォワードパスロジックで run_forward_step を実装します。

import torch
from megatron import get_tokenizer, mpu
from megatron.utils import get_ltor_masks_and_position_ids
from rapidformer import RapidformerEngine, get_args, PreTrainer

class MegatronGPTPreTrainer(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

        data = next(data_iterator) if data_iterator is not None else None
        data_b = mpu.broadcast_data(keys, data, datatype)

        tokens_ = data_b['text'].long()
        labels = tokens_[:, 1:].contiguous()
        tokens = tokens_[:, :-1].contiguous()

        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

エンジンを初期化し、トレーナーを作成し、rapidformer_pretrain_megatron_gpt_trainer.py として保存します。
```
engine = RapidformerEngine()
trainer = MegatronGPTPreTrainer(engine=engine)
trainer.train()
```

テンソル並列処理とパイプライン並列処理を有効にして、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_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

この例では、利用可能なすべての Megatron 固有の最適化を有効にします：テンソル並列処理 (--tensor-model-parallel-size 2)、パイプライン並列処理 (--pipeline-model-parallel-size 2)、グローバルバッチサイズによる勾配蓄積、メモリマップドデータ読み込み、ZeRO Stage 2 メモリ最適化、勾配チェックポイント、および混合精度トレーニング。

次のステップ

パラメーター設定ガイド — 完全な CLI パラメーターリファレンス
Rapidformer API — Finetuner、PreTrainer、および RapidformerEngine のメソッドシグネチャ