混合精度トレーニング

PyTorch による効率的な AI モデルトレーニング

Dennis Lee

Data Engineer, Amazon

混合精度トレーニングによる計算の高速化

メモリ効率・通信効率・計算効率を表すアイコン。

精度を下げて計算を高速化

FP32とFP16での浮動小数点数の表現を示す図。

精度を下げて計算を高速化

FP32とFP16での浮動小数点数の表現を示す図。

精度を下げて計算を高速化

FP32とFP16での浮動小数点数の表現を示す図。

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象
アンダーフロー防止のため損失をスケーリング

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象
アンダーフロー防止のため損失をスケーリング

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象
アンダーフロー防止のため損失をスケーリング

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象
アンダーフロー防止のため損失をスケーリング

混合精度トレーニングとは？

混合精度トレーニングの手順を示す図。

混合精度トレーニング：FP16とFP32の計算を組み合わせてトレーニングを高速化
アンダーフロー：数値が精度の下限を下回り0になる現象
アンダーフロー防止のため損失をスケーリング

PyTorchによる実装

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較したチャート。

PyTorchによる混合精度トレーニング

scaler = torch.amp.GradScaler()


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]

    with torch.autocast(device_type="cpu", dtype=torch.float16):

         outputs = model(inputs, labels=targets)
         loss = outputs.loss

    scaler.scale(loss).backward()

    scaler.step(optimizer)

    scaler.update()

    optimizer.zero_grad()

PyTorchからAcceleratorへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較したチャート。

PyTorchからAcceleratorへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較したチャート。

Acceleratorによる混合精度トレーニング

accelerator = Accelerator(mixed_precision="fp16")


model, optimizer, train_dataloader, lr_scheduler = \
    accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

AcceleratorからTrainerへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較したチャート。

AcceleratorからTrainerへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較したチャート。

Trainerによる混合精度トレーニング

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    fp16=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    compute_metrics=compute_metrics,
)

trainer.train()

練習しましょう！

PyTorch による効率的な AI モデルトレーニング

Preparing Video For Download...