AI用機械学習プラットフォーム (PAI) チームによって開発されたPAI-Rapidformerツールは、トランスフォーマーのトレーニングを加速するためのさまざまな方法を提供します。 ブラックボックスまたはホワイトボックスの方法でトレーニングを加速するには、PAI-Rapidformerイメージのみをインストールする必要があります。 このトピックでは、PAI-Rapidformerを使用してPyTorch変圧器のトレーニングを加速する方法について説明します。
前提条件
PAI-Rapidformerイメージがインストールされています。 詳細については、「Pai-Megatron-Patchイメージのインストール」をご参照ください。
あなたはPAI-Rapidformerのパラメータ設定に精通しています。 詳細については、「パラメーター設定」をご参照ください。
あなたはPAI-Rapidformerの方法に精通しています。 詳細は、「API」をご参照ください。
背景情報
PAI Rapidformerを使用することにより、ブラックボックスまたはホワイトボックスの方法で変圧器のトレーニングを加速することができます。
ブラックボックス方式
PAI-Rapidformerはコードを記述する必要性なしで変圧器の訓練を加速するのに簡単な構成を使用できるCLIを提供します。 トランスフォーマーのトレーニングをブラックボックスで高速化する場合は、トランスフォーマーとデータセットを事前に登録してください。 次のセクションでは、トレーニングアクセラレーションをブラックボックスで実装する方法の例を示します。
ホワイトボックス方式
PAI Rapidformerはコードを書く必要性なしで変圧器の訓練を加速できるCLIを提供します。 PAI-Rapidformerは、必要に応じてデータセットとトランスフォーマーをカスタマイズできるコードテンプレートも提供します。 次に、
-- user-scriptパラメーターを使用して、トレーニングの高速化のためにカスタマイズされたデータセットとトランスフォーマーをCLIに渡すことができます。 次のセクションでは、トレーニングアクセラレーションをホワイトボックスで実装する方法の例を示します。データハブまたはモデルハブを使用した高速化
データハブまたはモデルハブを使用せずに完全にカスタマイズされたアクセラレーション (より柔軟)
ブラックボックス方式でハギングフェイストランスフォーマーの微調整を加速
データセットをHugging Faceに登録するか、登録されている目的のデータセットを検索します。 この方法では、
-- dataset-nameパラメーターを使用して、手順3のPAI-Rapidformerにデータセットを渡すことができます。詳細については、「データセットをハギングフェイスに登録する」と「データセット」ページをご参照ください。
トランスフォーマーをハギングフェイスに登録するか、登録されている目的のトランスフォーマーを検索します。 この方法では、
-- pretrained-model-name-or-pathパラメーターを使用して、手順3で変換器をPAI-Rapidformerに渡すことができます。詳細については、「トランスフォーマーにモデルを追加する方法」と「モデル」ページをご参照ください。
PAI-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 \ # The task name. --pretrained-model-name-or-path 'bert-base-cased' \ # The name of the registered transformer. --data-path glue \ # The path of the registered dataset. --data-name mrpc \ # The file name of the registered dataset. --epochs 3 \ # The number of epochs. --micro-batch-size 16 \ # The size of the mini-batch on a single GPU. --global-batch-size 64 \ # The size of the global batch on all GPUs during distributed training. --lr 2e-5 \ # The learning rate. --lr-decay-style linear \ # The scheme of learning rate decay. --lr-warmup-iters 100 \ # The learning rate warmup. --weight-decay 1e-2 \ # The weight decay value. --clip-grad 1.0 \ # The gradient clipping value. --seed 42 \ # The random seed. --mixed-precision \ # Enables FP16 data transfers. --onnx-runtime-training \ # Enables graph optimization provided by ONNX Runtime. --zero-1-memory-optimization \ # Uses Zero Redundancy Optimizer (ZeRO) to partition optimizer states.上記のスクリプトのパラメーターの詳細については、「パラメーター設定」をご参照ください。
ブラックボックス方式でハギングフェイストランスフォーマーの事前トレーニングを加速
作成します。事前トレーニングデータセットとしてのmmapファイル。
詳細については、「Megatron-LM」をご参照ください。サンプルコード:
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トランスフォーマーをハギングフェイスに登録するか、登録されている目的のトランスフォーマーを検索します。 この方法では、
-- pretrained-model-name-or-pathパラメーターを使用して、手順3で変換器をPAI-Rapidformerに渡すことができます。詳細については、「トランスフォーマーにモデルを追加する方法」と「モデル」ページをご参照ください。
PAI-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 \ # The size of the global batch. --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 \ # Enables FP16 data transfers. --onnx-runtime-training \ # Enables graph optimization provided by ONNX Runtime. --fsdp-memory-optimization \ # Uses Fully Sharded Data Parallel (FSDP) to partition optimizer states, gradients, and parameters.上記のスクリプトのパラメーターの詳細については、「パラメーター設定」をご参照ください。
Finetunerコードテンプレートを使用して、ホワイトボックス方式でHugging Face transformerの微調整を高速化
このセクションでは、PAI-Rapidformerが提供するFinetunerコードテンプレートを使用して、Hugging Faceトランスフォーマーの微調整を加速する方法について説明します。 コードテンプレートの次の4つのメソッドに注意してください。
train_valid_test_datasets_provider: このメソッドを呼び出してデータを生成します。model_optimizer_lr_scheduler_provider: このメソッドを呼び出して、トランスフォーマー、オプティマイザ、およびラーニングレートスケジューラを構築できます。run_forward_step: このメソッドを呼び出して、前方伝播のロジックを定義できます。run_compute_metrics: このメソッドを呼び出して、トランスフォーマーをトレーニングし、同時に精度を計算できます。
これらのメソッドの詳細については、「API」をご参照ください。 次の情報は、これらのメソッドの入力と出力を示しています。
class MyFintuner(Finetuner):
def __init__(self, engine):
super().__init__(engine=engine)
# Obtains a training, verification, or test dataset.
# Input: none.
# Output: three objects and an object method.
def train_valid_test_datasets_provider(self):
return train_dataset, valid_dataset, test_dataset, collate_f
# Creates a transformer, an optimizer, or a learning rate scheduler.
# Input: none.
# Output: three objects.
def model_optimizer_lr_scheduler_provider(self):
return model, optimizer, lr_scheduer
# Defines the logic of forward propagation.
# Input: a batch or an iterator, and a transformer.
# Output: loss.
def run_forward_step(self, batch_or_iterator, model):
return loss
# Defines the logic of using the verification dataset. This logic is tailored to the fine-tuning of transformers.
# Input: a transformer and a data loader for the verification dataset.
# Output: a metric object.
def run_compute_metrics(self, model, eval_dataloader):
return metric
コードテンプレートに慣れたら、「ブラックボックス方式でハグ面トランスフォーマーの微調整を加速する」セクションの指示に従って、データセットとトランスフォーマーを準備します。 次に、次の手順を実行します。
PAI-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コードテンプレートの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 => use the model max length (it's actually the default) outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) return outputs datasets = load_dataset(args.dataset_path, args.dataset_name) # Apply the method we just defined to all the examples in all the splits of the dataset 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 # after each epoch run metric on eval dataset 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_metricPAI-Rapidformerを初期化し、トレーナーオブジェクトを作成します。
finetune()メソッドを呼び出し、rapidformer_finetune_huggingface_bert_trainer.pyという名前のファイルとして出力を保存します。engine = RapidformerEngine() trainer = MyFintuner(engine=engine) trainer.train()PAI-RapidformerのCLIで起動スクリプトを設定します。 スタートアップスクリプトで、
-- user-scriptパラメーターをrapidformer_finetune_huggingface_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 # Enables FP16 data transfers. --zero-3-memory-optimization \ # Uses ZeRO to partition optimizer states, gradients, and parameters. --onnx-runtime-training \ # Enables graph optimization provided by ONNX Runtime.
Pretrainerコードテンプレートを使用して、ホワイトボックス方式でハグ顔トランスフォーマーの事前トレーニングを高速化
PAI-Rapidformerが提供するPretrainerコードテンプレートを使用して、ホワイトボックス方式でHugging Face transformerの事前トレーニングを加速する場合は、コードテンプレートの次の方法に注意してください。
train_valid_test_datasets_provider: このメソッドを呼び出してデータを生成します。model_optimizer_lr_scheduler_provider: このメソッドを呼び出して、トランスフォーマー、オプティマイザ、およびラーニングレートスケジューラを構築できます。run_forward_step: このメソッドを呼び出して、前方伝播のロジックを定義できます。
これらのメソッドの詳細については、「API」をご参照ください。 これらのメソッドの入力と出力の詳細については、このトピックの「Finetunerコードテンプレートを使用して、ホワイトボックス方式でハグ面トランスフォーマーの微調整を加速する」を参照してください。
コードテンプレートに慣れたら、「ブラックボックス方式でハグ面トランスフォーマーの微調整を加速する」セクションの指示に従って、データセットとトランスフォーマーを準備します。 次に、次の手順を実行します。
PAI-RapidformerとHugging Faceトランスの方法をインポートします。
説明モデルの事前トレーニングはイテレータを使用してデータを読み取るため、並列データ処理を実現するにはMPU (Micro Processor Unit) を使用する必要があります。
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_huggingfacePreTrainerから派生したクラスを使用して、事前トレーニングアクセラレーションのメソッドの設定を完了します。 サンプルコード:
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): # Items and their type. keys = ['input_ids', 'attention_mask', 'token_type_ids', 'labels', 'next_sentence_label'] datatype = torch.int64 # Broadcast data. 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']PAI-Rapidformerを初期化し、トレーナーオブジェクトを作成します。
pretrain()メソッドを呼び出し、出力をrapidformer_pretrain_huggingface_bert_trainer.pyという名前のファイルとして保存します。engine = RapidformerEngine() trainer = MyBertPreTrainer(engine=engine) trainer.train()PAI-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 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 \ # Enables acceleration. --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 \ # Uses ZeRO to partition optimizer states, gradients, and parameters. --onnx-runtime-training \ # Enables graph optimization provided by ONNX Runtime. --mixed-precision # Enables FP16 data transfers.
カスタムトレーナーを使用して、ホワイトボックス方式でハギングフェイストランスフォーマーの微調整を高速化
カスタムトレーナーに基づくアクセラレーションの場合、PAI-Rapidformerは、Apex FusedAdam、モデル状態分割、グラフ最適化など、いくつかのアクセラレーション機能のみを提供します。 混合精度トレーニングではトレーニングプロセスに複雑な変更が必要になるため、PAI-Rapidformerが提供するコードテンプレートを使用してトレーニングを高速化することをお勧めします。 このセクションでは、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))先行コードは、並列データトレーニング、オプティマイザの高速実行、または混合精度トレーニングをサポートしていません。 したがって、コードを最適化する必要があります。 PAI-Rapidformerのメソッドを使用して、コードを最適化できます。 そうするには、次の手順を実行します。
並列データ処理を有効にします。
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"])Apex FusedAdamを有効にします。
元のオプティマイザを、PAI-Rappiformerが提供するApex FusedAdamに置き換えます。
engine.comポーズを呼び出して、トランスフォーマー、オプティマイザー、およびラーニングレートオプティマイザーをカプセル化することで、オプティマイザーを置き換えることができます。+ 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 FusedAdam、および混合精度を有効にする場合、混合精度トレーニングでは、トレーニングプロセスの変更、FP16データ転送の有効化、損失スケーリングの有効化などの特定の操作を実行する必要があります。 これは、複雑なフロントエンドの変更を伴う。 このような変更を避けるために、トレーナーを使用してアクセラレーションを実装できます。 PAI-Rapidformerが提供するFinetunerコードテンプレートは、並列データ処理、Apex FusedAdam、Pytorch混合精度トレーニング、Megatronオプティマイザが提供する混合精度トレーニング、DeepSeedとFairScaleが提供するメモリ最適化の加速など、さまざまなアクセラレーション機能をサポートします。。
Pretrainerコードテンプレートを使用して、ホワイトボックス方式でメガトロン変圧器の事前トレーニングを高速化
「Accelerate the Accelerate the fine-tuning of a Hugging Face transformer in a white-box manner by using a custom trainer」セクションで説明されている加速方法と比較して、この加速方法はより柔軟です。 Data HubまたはModel Hubを使用する必要はありません。 代わりに、train_valid_test_datasets_provider、model_optimizer_lr_scheduler_provider、およびrun_forward_stepメソッドを呼び出して、カスタムデータ生成、カスタムトランスフォーマー構築、およびフォワードプロパゲーションのロジックを定義できます。
事前トレーニングデータセットとしてのmmapファイルを作成します。
詳細については、「Megatron-LM」をご参照ください。サンプルコード:
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-eodPreTrainerから派生したクラスを使用して、カスタムデータの生成に使用する
train_valid_test_datasets_providerメソッドを構成します。カスタムデータ生成のロジックを定義して、サードパーティのライブラリを使用することなく、トレーニング、検証、およびテストデータセットを作成できます。 データセットは、
torch.utils.data.Datasetクラスを継承する必要があります。 サンプルコード：rapidformer importからの
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_dsPreTrainerから派生したクラスを使用して、カスタムトランスフォーマーの構築に使用する
model_optimizer_lr_scheduler_providerメソッドを構成します。カスタム変圧器構築のロジックを定義して、サードパーティのライブラリを使用することなく、カスタム変圧器を構築できます。 モデルは
torch.nn.Moduleクラスを継承する必要があります。 サンプルコード：rapidformer importからの
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, NonePreTrainerから派生したクラスを使用して、前方伝播を実現するために使用される
run_forward_stepメソッドを構成します。rapidformer importからの
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): """Forward step.""" args = get_args() tokenizer = get_tokenizer() # Items and their type. keys = ['text'] datatype = torch.int64 # Broadcast data. if data_iterator is not None: data = next(data_iterator) else: data = None data_b = mpu.broadcast_data(keys, data, datatype) # Unpack. tokens_ = data_b['text'].long() labels = tokens_[:, 1:].contiguous() tokens = tokens_[:, :-1].contiguous() # Get the masks and postition ids. 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 lossPAI-Rapidformerを初期化し、トレーナーオブジェクトを作成します。 次に、
pretrain()メソッドを呼び出し、出力をrapidformer_pretrain_megatron_gpt_trainer.pyという名前のファイルとして保存します。engine = RapidformerEngine() trainer = MyGPTPreTrainer(engine=engine) trainer.train()PAI-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 DATA_PATH=book_wiki_owtv2_small_text_sentence PRETRAINED_CHECKPOINT= rapidformer --user-script rapidformer_pretrain_megatron_gpt_trainer.py \ --tensor-model-parallel-size 2 \ # The size of tensor parallelism. --pipeline-model-parallel-size 2 \ # The size of pipeline parallelism. --num-layers 12 \ --hidden-size 768 \ --num-attention-heads 12 \ --micro-batch-size 16 \ --global-batch-size 128 \ # The size of the global batch. --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 \ # Enables acceleration. --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-3-memory-optimization \ # Uses ZeRO to partition optimizer states and gradients. --checkpoint-activations \ # Enables checkpoints. --mixed-precision # Enables FP16 data transfers.