エラー解決

【PyTorch】Gradient Checkpointingで大規模モデルを学習する際のメモリ不足エラー解決法

· 13分で読めます · 2 閲覧

問題の概要：大規模モデル学習時のメモリ不足エラー

PyTorchを使用して大規模なニューラルネットワーク（例：数十億パラメータを持つTransformerモデル）を学習する際、以下のようなメモリ不足エラーに遭遇することがあります。

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 10.00 GiB total capacity; 8.50 GiB already allocated; 0 bytes free; 8.70 GiB reserved in total by PyTorch)

あるいは、バッチサイズを極端に小さくしてもエラーが解消されず、モデルの順伝播（フォワードパス）中に必要な中間活性化の保存に大量のGPUメモリが消費されてしまう状況です。この問題は、モデルの層数が深くなったり、パラメータ数が増えたりするほど顕著になります。バックプロパゲーション（誤差逆伝播）を実行するためには、フォワードパスで計算されたすべての中間テンソル（活性化）をメモリ上に保持しておく必要があるためです。

原因の解説：計算グラフとメモリ消費の関係

このメモリ不足の根本的な原因は、PyTorchの自動微分エンジン「Autograd」の仕組みにあります。デフォルトでは、フォワードパスで計算されたすべての中間テンソルは、バックプロパゲーションで勾配を計算するためにメモリ上に保持されます。モデルが大規模になるほど、この「中間活性化の保存」に必要なメモリ量は線形的に増加します。

具体的には、N層のモデルを学習する場合、O(N)のメモリが中間活性化の保存に必要となります。これが、限られたGPUメモリで非常に深いモデルを学習することを困難にしています。Gradient Checkpointing（勾配チェックポイント）は、この計算コスト（メモリ使用量）と計算時間のトレードオフを管理する技術です。

解決方法：Gradient Checkpointingの実装手順

Gradient Checkpointingの核心は、「中間活性化をすべて保存する代わりに、一部のみを保存し、必要になった時点で再計算する」というアイデアです。これにより、メモリ使用量をO(N)からO(√N)程度に削減できますが、その分、フォワードパスの再計算が発生するため計算時間は増加します（通常は20-30%程度）。

ステップ1: torch.utils.checkpointのインポート

まず、必要なモジュールをインポートします。PyTorch 1.0以降では、標準ライブラリにこの機能が含まれています。

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint, checkpoint_sequential

ステップ2: カスタムモデルでのcheckpoint関数の適用

モデルの一部のセクション（例えば、Transformerのブロック）をチェックポイントとして指定します。ここでは、シンプルなResNet風のブロックを例にします。

class MyBlock(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.linear1 = nn.Linear(in_dim, out_dim)
        self.linear2 = nn.Linear(out_dim, out_dim)
        self.relu = nn.ReLU()

    def forward(self, x):
        # このブロック全体をチェックポイントとする
        # 引数として「この関数オブジェクト」と「必要な入力」を渡す
        return checkpoint(self._forward, x)

    def _forward(self, x):
        # 実際のフォワードパスの計算
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.relu(x)
        return x

checkpoint関数の第一引数は関数オブジェクト、第二引数以降はその関数への入力です。チェックポイント化されたセクションでは、フォワードパス時に中間活性化が保存されず、バックワードパスで必要になった時点で_forward関数が再度実行されます。

ステップ3: Sequentialモデルでのcheckpoint_sequentialの適用

モデルがnn.Sequentialで構成されている場合は、checkpoint_sequential関数を使用するとより簡単です。

# 大きなSequentialモデルを作成
layers = [nn.Linear(1024, 1024) for _ in range(20)]
model = nn.Sequential(*layers, nn.ReLU())

# 学習ループ内で、フォワードパス時にチェックポイントを適用
def forward_with_checkpoint(input):
    # 20層を4つのセグメントに分割してチェックポイント化
    segments = 4
    return checkpoint_sequential(model, segments, input)

ステップ4: Transformerモデルへの適用例（実践的）

Hugging Face Transformersライブラリなど、既存の大規模モデルを学習する場合の実践的な例です。

from transformers import AutoModelForCausalLM, AutoConfig
import torch

# モデルの設定と読み込み
config = AutoConfig.from_pretrained("gpt2")
config.gradient_checkpointing = True  # 設定で有効化
model = AutoModelForCausalLM.from_pretrained("gpt2", config=config)

# または、読み込み後に有効化する方法
model.gradient_checkpointing_enable()

print(f"Gradient checkpointing is enabled: {model.is_gradient_checkpointing}")

ステップ5: 学習ループでの注意点と検証

Gradient Checkpointingを有効にしたら、メモリ使用量が削減されていることを確認し、学習が正常に進むか検証します。

# メモリ使用量のモニタリング
print(f"Initial memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

# ダミーデータで試す
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
input_ids = torch.randint(0, 50257, (2, 1024)).cuda()  # 小さなバッチ
labels = input_ids.clone()

# 学習ステップ
model.train()
optimizer.zero_grad()
outputs = model(input_ids, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()

print(f"Peak memory allocated: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
print(f"Loss: {loss.item():.4f}")

コード例・コマンド例：エラー発生時の対処法

Gradient Checkpointingを適用してもメモリ不足が解消されない場合、以下の点を確認してください。

# よくある問題1: チェックポイント化されていない巨大なテンソル
# 回避策: 大きなテンソルの生成をチェックポイント内に含める
def forward(self, x):
    # 悪い例: 大きな中間テンソルがチェックポイント外で生成される
    big_tensor = torch.randn(1000, 1000, device=x.device)  # これがメモリを圧迫
    return checkpoint(self._forward, x, big_tensor)  # big_tensorも保存対象になる

    # 良い例: 大きなテンソルの生成もチェックポイント関数内で行う
    return checkpoint(self._forward_with_creation, x)

def _forward_with_creation(self, x):
    big_tensor = torch.randn(1000, 1000, device=x.device)  # 必要時のみ生成
    # ... 続く処理

# よくある問題2: inplace操作との非互換性
# PyTorchのチェックポイントは、inplace操作（例: relu_, add_）と相性が悪い場合があります。
# エラーメッセージ: "RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation"
# 解決策: inplace操作を避ける、またはモデル設計を見直す

# 悪い例
x.relu_()  # inplace操作

# 良い例
x = torch.relu(x)  # 関数型の操作

まとめ・補足情報

Gradient Checkpointingは、限られたGPUメモリで大規模モデルを学習するための強力な技術です。本記事で解説したように、torch.utils.checkpointモジュールを利用することで、比較的簡単に既存のモデルに適用できます。

主な利点:

  • GPUメモリ使用量を大幅に削減できる（O(N) → O(√N)）
  • バッチサイズを増やしたり、より大きなモデルを学習したりできる
  • PyTorch標準機能のため、追加インストールが不要

注意点とベストプラクティス:

  • 計算時間の増加: メモリ使用量の削減と引き換えに、計算時間は20-30%程度増加します。これはトレードオフとして理解する必要があります。
  • チェックポイントの粒度: あまりに細かい単位でチェックポイントを設定すると再計算のオーバーヘッドが大きくなります。逆に、大きなブロック単位ではメモリ削減効果が小さくなります。モデルの構造に応じて適切な単位（例：Transformerの1ブロックごと）を選択しましょう。
  • 検証時の無効化: モデルの検証（評価）時には、勾配計算が不要なため、チェックポイントを無効にすることで推論速度を向上させられます。torch.no_grad()コンテキスト内では、チェックポイントのオーバーヘッドは事実上無視できますが、明示的に無効化する方法も検討してください。
  • 他のメモリ最適化技術との組み合わせ: 混合精度訓練（AMP）、モデル並列化、オフローディングなど、他のメモリ最適化技術と組み合わせることで、さらに大規模なモデルの学習が可能になります。

最終的に、Gradient Checkpointingは「メモリ」と「計算時間」の貴重なトレードオフを管理するツールです。実際のプロジェクトでは、利用可能なハードウェアリソースと学習時間の制約を考慮し、この技術を効果的に活用してください。

この記事は役に立ちましたか？