勾配累積

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

Dennis Lee

Data Engineer, Amazon

分散学習

コーストピックのフローチャート：データ準備、分散学習、効率的な学習、オプティマイザー。

効率的な学習

コーストピックのフローチャート：データ準備、分散学習、効率的な学習、オプティマイザー。

学習効率の向上

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

学習効率の向上

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

学習効率の向上

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

勾配累積によるメモリ効率の向上

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

大きいバッチサイズの問題

大きいバッチサイズ：安定した勾配推定により学習が高速化
GPUメモリがバッチサイズを制限する

大きいバッチサイズではメモリ不足エラーが発生し、小さいバッチサイズでは学習が正常に完了することを示す図。

勾配累積の仕組み

大きいバッチを小さいバッチに分割する図。

勾配累積：小さいバッチで勾配を合算
大きいバッチでモデルを実質的に学習
勾配合算後にモデルパラメータを更新

複数バッチの勾配の合算として勾配累積を示す図。

PyTorch、Accelerator、Trainer

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

PyTorch、Accelerator、Trainer

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

PyTorch、Accelerator、Trainer

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if ((index + 1) 
        % gradient_accumulation_steps == 0):

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchによる勾配累積

for index, batch in enumerate(dataloader):
    inputs, targets = (batch["input_ids"], 
                       batch["labels"])
    inputs, targets = (inputs.to(device), 
                       targets.to(device))
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if ((index + 1) 
        % gradient_accumulation_steps == 0):
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

複数バッチの勾配の合算として勾配累積を示す図。

PyTorchからAcceleratorへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

PyTorchからAcceleratorへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):

        inputs, targets = (batch["input_ids"],
                           batch["labels"])

複数バッチの勾配の合算として勾配累積を示す図。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):

        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss

複数バッチの勾配の合算として勾配累積を示す図。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss

複数バッチの勾配の合算として勾配累積を示す図。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)

複数バッチの勾配の合算として勾配累積を示す図。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)

複数バッチの勾配の合算として勾配累積を示す図。

Acceleratorによる勾配累積

accelerator = \
    Accelerator(gradient_accumulation_steps=2)

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = (batch["input_ids"],
                           batch["labels"])
        outputs = model(inputs, 
                        labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

複数バッチの勾配の合算として勾配累積を示す図。

AcceleratorからTrainerへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

AcceleratorからTrainerへ

PyTorch、Accelerator、Trainerの使いやすさとカスタマイズ性を比較するグラフ。

Trainerによる勾配累積

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=2)

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

trainer.train()

{'epoch': 1.0, 'eval_loss': 0.73, 'eval_accuracy': 0.03, 'eval_f1': 0.05}
{'epoch': 2.0, 'eval_loss': 0.68, 'eval_accuracy': 0.19, 'eval_f1': 0.25}

練習しましょう！

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