12GBのGPUで動画モデルを微調整する方法：ZeRO-3を用いたマルチGPU実験の記録

この記事では、以前サーバー上で動画モデルの微調整を試みた際の過程を記録しておく。

Video-ORAを使用しており、その役割はモデルに動画を視聴させ、保持すべき時間範囲を出力させ、それを後続のプログラムに引き継いで編集を行うことです。私はSFTを担当しており、具体的には動画と対応する正解を用いて教師あり微調整を行い、LoRAを用いて少量の追加パラメータを学習させています。

サーバーにはRTX 3060が数枚あり、各カード12GBです。当初、1枚のGPUで8フレームしか表示していなかったにもかかわらず、ビデオメモリの使用量は約11.24 GiBに達し、これ以上フレーム数を増やすとすぐにOOM（ビデオメモリ不足）に陥ってしまいました。動画は動作の前後一連の流れを見る必要があるため、ずっとフレーム数を減らし続けるわけにはいきません。

（実のところ、最初から単カードではうまくいかないことは分かっていたんだけど、それでも試してみたかったんだ）

最後の2枚のGPUカードで、64フレームのサンプルを含む多段階トレーニングを完了し、4枚のGPUカードでの更新、保存、および再読み込みも成功した。その中で最も時間がかかったのは重みの読み込みであり、さらに当初見落としていた問題として、動画の解像度制限がまったく期待通りに機能していなかったことが判明した。

1. グラフィックメモリは一体どこに使われているのか

その時点では、すでにBF16、LoRA、勾配チェックポイントを使用しており、視覚エンコーダーとマルチモーダル接続層も凍結済みで、1回のカーブ処理で1つのサンプルしか処理していなかった。

フレーム数が少ない場合、2枚のDDPカードで動作するが、VRAMのピーク使用量は依然として11.2 GiB近くになる。DDPは各カードにモデルのコピーを保持するが、今回は1枚のカードでは収まらないという問題を解決する必要があるため、引き続きDeepSpeed ZeRO-3を試行し、パラメータなどのトレーニング状態を分割して異なるカードに配置する。

ただし、調査を行う際には、まずOOMがどの段階で発生したのかを特定する必要があります。

ある時、シングルカードチェックが caching_allocator_warmup() 8.35 GiB の申請に失敗した。一見するとモデルが大きすぎるように見えるが、実際には検証スクリプトが同じプロセス内で 2 回呼び出されていたためである。 get_model_processor()、2つ目のモデルの読み込みを開始した。その後、リロードチェックを独立したプロセスに分離し、トレーニングプロセスは更新と保存のみを担当するようにした。

また、ZeRO-3で重みを読み込む際のgather（つまり、シャードパラメータを一時的に集約する処理）で、数回エラーが発生しました。動画はまだ入力されていないため、フレームレートを下げるだけでは当然解決できません。

その後、エラーが発生した箇所ごとに分けて調査しました。読み込み段階では重みと仮割り当てを確認し、バッチを受け取った後は、動画の入力と学習プロセスを確認しました。こうすることで、少なくともOOMエラーが発生したからといって、すぐにすべてのパラメータを小さくしてしまうようなことは避けられます。

2. 動画の予算が有効になっていない

この部分はガカよりも先に確認する価値がある。

もともと設定されていた max_pixels、動画の解像度はすでに制限されていると思っていた。解像度を下げてから実際のバッチを確認したところ、テンソルはそれに応じて縮小されていなかった。当時使用していたQwen3.5の動画処理ブランチでは、このパラメータに基づいて入力が制限されていなかった。

元の 8 フレームのサンプル、video_grid_thw はい [4,72,40]、テキストと合わせると合計 3058 トークンになります。その後、動画のトークン予算に変更し、関連する処理コードをインポートする前に以下のように設定しました：

export VIDEO_MAX_TOKEN_NUM=196
export VIDEO_MIN_TOKEN_NUM=128

同じサンプルを再処理したところ、その違いは明らかでした：

最初の2行のビジュアルトークンが約2880から720に減少した。この分のスペースを節約した結果、従来のシングルGPU環境でも、16秒・32フレームの連続3回の更新が正常に実行され、VRAMのピーク使用量は約11.28 GiBとなった。

つまり、当初の「8フレームで上限に達した」という結果は、その入力設定がVRAMを大量に消費していたことを示すに過ぎず、このグラフィックカードの固定フレームレート上限と見なすことはできない。今回は、1フレームあたりの空間的なディテールを若干抑えることで、より多くのフレーム数を確保した。ディテールの損失がタスクの成果にどのような影響を与えるかについては、別途評価する必要がある。

このグリッドも誤解されやすい。video_grid_thw 記録されているのは、時間、高さ、幅の3方向のグリッドです。当時の時間パッチサイズは2であったため、 [16,36,20] これは 16 フレームではなく、32 フレームの入力に対応しています。

VIDEO_MAX_TOKEN_NUM=196 また、動画全体が最大で196トークンしかないという意味でもありません。この処理パスでは、各フレームのスペース予算に制約が設けられていますが、上記の32フレームのサンプルを統合した後の視覚トークンは、依然として以下の通りです：

16 × (36 / 2) × (20 / 2) = 2880

その後のログには、実際のサンプリングフレーム数、グリッド、および入力トークンをすべて記録しました。設定ファイルに何を書いたかに関わらず、最終的にはモデルが実際に受け取った量が重要になります。これら2つの環境変数の使い方は当時のプロセッサに合わせて設定されたものであり、モデルや処理ライブラリを変更した場合は、改めて確認する必要があります。

3. ZeRO-3 の重みの読み込みに関する問題

2つのカードは当初、次のようなログの後に停止していました：

Video-ORA ZeRO-3 loaded 250 parameters

その後、欠員を報告しながら、 _initialize_missing_keys() 付近でOOMが発生。モデル構造のみを構築し、実際の重みをロードしないメタデバイスプローブで検証したところ、チェックポイントのファイルとインデックスは正しくアラインメントされていたため、問題は古いカスタムロードロジックにあることが判明した。

その違いの一つは、重み付けキーの名前です：

checkpoint:  model.language_model.layers.0....
model state:       language_model.layers.0....

プレフィックマッピングを追加するだけでは不十分で、再帰的な読み込み時の検索、コピー、および初期化完了の判定もすべて対応させる必要があります。そうしないと、重みファイルが明らかに存在しているにもかかわらず、フレームワークはパラメータが読み込まれていないものとみなして、パラメータが欠落している場合の初期化プロセスに入ってしまうことになります。

もう一つの落とし穴は、TensorとParameterを混同して使用してしまうことです。ローダーは通常の state_dict() 値を取得するが、取得する値は必ず nn.Parameter。最初にエラーが発生したのは、共有重みに関するものでした。 lm_head、特例処理を追加したところ、通常の視覚層のバイアスで再び同じエラーが発生した。

この問題は CPU を見れば一目瞭然です：

layer = torch.nn.Linear(2, 2)

type(layer.state_dict()["weight"])
# torch.Tensor

layer.state_dict(keep_vars=True)["weight"] is layer.weight
# True

調査の過程で、「関数が変更されたにもかかわらず、実行されているのは依然として古い関数である」という現象にも遭遇した。モジュール内のローダーが置き換えられたにもかかわらず、別のモジュールがインポート時に保持していた関数への参照が変更されていなかったためである。互換性パッチが積み重なるほど確認が難しくなったため、その後、メインブランチのトレーニングを隔離環境に移行し、上流のローダーを検証した上で、古いパッチを無効化した。

当時の環境設定は以下の通りです。これは今後のログ確認に役立てるためのものであり、必ずしもこの一般的なバージョン構成通りにインストールする必要があるわけではありません：

隔離環境では元の環境が維持され、その他の依存関係は検索パスを通じて再利用される。その後、再実行した際に以前のエラーが再び発生し、ログを確認したところ、起動コマンドに古いインタプリタと古いローダーが指定されていたことが判明した。それ以来、実際のインタプリタ、依存関係のバージョン、スクリプトのハッシュ値も実行記録に併記するようにしている。

重み付けの検証で一度誤検知がありました：ZeRO-3 の未集約パラメータが次のように表示される場合があります [0]、このローカル形状をそのまま完全なチェックポイントと比較しても、当然一致しません。検証時には、代わりに論理形状を使用します：

def logical_parameter_shape(parameter):
    return list(getattr(parameter, "ds_shape", parameter.shape))

数値比較では、各ランクを同じ順序でギャザーに送り込み、同じdtypeに変換されたソース重みと要素ごとに比較を行います。最終的に724個のチェックポイント項目がすべて一致し、この読み込みステップが完了して初めて、照合が完了したものとみなされます。

4. トレーニングがステップ0で停止している

2枚のカード間の通信については、まずNCCLのall-reduceでテストしました：各カードに1つずつ配置しました 1.0、双方とも受信しました 2.0。モデルの読み込みも完了しましたが、トレーニングログには次のように表示されています：

global_step/max_steps: 0/1
train_loss: 0
train_steps_per_second: 0

その原因は、最小テストでサンプルが1つしか用意されていなかったためです。サンプリングアルゴリズムはデータをランク順に分類するため、2枚のカードがそれぞれ利用可能なバッチを取得できませんでした。シャッフルを無効にしても意味がなく、順序を変えてもデータが突然1つ増えることはありません。

単一サンプルの検査を、以下のように変更する world_size 複製し、各トレーニングプロセスに1つずつ割り当てる。このスイッチはスモークテストでのみ使用される。トレーニング前にはDataLoaderによるチェックも追加されており、どのrankがbatchを取得できなかったか、事前にエラーとして報告される。

別のウィンドウに切り替えて複数回のテストを行ったところ、3つのサンプルのうち実際には2つしか使用されていなかったことが判明した。当時のサンプラーのスライシングパスは以下の通りであった。 total_samples // world_size 処理の過程で、2枚のカードのうち、末尾の1枚が失われてしまった。そこで、テストを4つの異なるウィンドウに変更し、各カードから実際に読み取られたサンプルを順次記録するとともに、サンプルの行数がカード枚数で割り切れるかどうかを確認した。

この2回の実行ではどちらもプログラムは起動し、「データがすべて使用されていない」というメッセージも直接表示されませんでした。進捗バーと損失値だけを見ていたため、見落としやすかったです。

パラメータ更新のチェックでもちょっとしたハプニングがありました。最初のステップで、LoRAテンソルの半分について勾配がすべてゼロになっていましたが、それらはすべて lora_A。このとき、A はランダムに初期化され、B は 0 に初期化され、増分は B × A、A への最初のバックプロパゲーションにおける勾配は、すべてゼロの B と乗算される。B は依然として更新されるため、「各パラメータの最初のステップで非ゼロの勾配がある」という基準だけで、学習が正常に行われているかどうかを判断することはできない。

5. 2カードおよび4カードのテスト結果

2枚のカードを用いた多段階テストでは、4つの異なるウィンドウが使用され、その中で最も長いものは32秒・64フレームで、実際にはrank 0でトレーニングが行われました：

video_grid_thw = [[32, 28, 28]]
input_tokens = 6770
supervised_tokens = 182

3つのステップ（順方向、逆方向、オプティマイザーの更新、アダプターの保存）を連続して完了し、4つのウィンドウすべてがバッチログに記録されました。2つのプロセスのそれぞれのピーク値は以下の通りです：

allocated は、PyTorch がテンソルに割り当てたビデオメモリであり、reserved にはそのキャッシュプールで確保された領域も含まれます。ここに記録されているのは PyTorch の指標であり、直接的には nvidia-smi カード全体の占有率。

これは長さが混在する一連のテストであり、2枚のグラフィックカードが同時にそれぞれ64フレームずつ処理したものと解釈してはなりません。ここでいう2枚のグラフィックカードによる64フレームの空間グリッドは、前述の単一グラフィックカードによるテストとは異なるため、2組の数値を単純に比較して、マルチカード構成でどの程度の割合が節約されたかを算出することはできません。

ビデオメモリ以外にも、解答が実際にトレーニングに組み込まれているかどうかを確認しました：assistant に対応する監督トークンが空ではなく、解答が切り捨てられておらず、残りの入力にはマスクが保持されていること。ウィンドウ分割後の時間も揃える必要があります。ソース動画の時間からウィンドウの開始点を差し引き、ウィンドウをまたぐイベントはまず現在のウィンドウの範囲内に切り詰めます。そうしないと、動画は局所的に0秒から始まるのに、解答には元の動画の時間が記載されたままになり、トレーニングは実行できるものの、入力されるのは位置がずれたサンプルになってしまいます。

使用した ZeRO の設定を以下に示します。通信バケットとプリフェッチバケットは小さめに設定し、通信オーバーラップは無効化しています。また、CPU や NVMe オフロードは使用していません。このバージョンでは、まずは容量の問題を解決することを優先しており、スループットの最適化は行っていません。

{
  "bf16": {"enabled": "auto"},
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": false,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5000000,
    "stage3_prefetch_bucket_size": 5000000,
    "stage3_param_persistence_threshold": 10000,
    "stage3_max_live_parameters": 100000000,
    "stage3_max_reuse_distance": 100000000,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

その後の4枚のGPUカードによるテストでは、同じ読み込みおよびトレーニングのエントリポイントを使用し、各ランクで異なるウィンドウを読み込み、連続的な更新、保存、および再読み込みもすべて成功した。その際は32秒のウィンドウ、32フレームの入力を使用し、各GPUカードのピーク時の予約メモリは約5.76～5.82 GiBであった。フレーム数は2枚のGPUでのテストとは異なるため、ここでは線形スケーリングの比較対象とはしない。以前の3枚のGPUを用いた単ステップテストでも速度向上は見られなかったが、今回のマルチGPU連携テストでは、まず「収まるかどうか」という問題を解決した。

6. アダプタのバックロードテスト

保存されたアダプターのサイズはおよそ10 MBです。トレーニングプロセスを終了した後、ベースモデルとアダプターを再読み込みすると、256個のテンソルが保存ファイルと要素ごとに一致しました。

その後、ショートテスト用のアダプターを使って、正解が示されていない生成テストを行ったところ、次のような結果が出た。 <argmax>11</argmax> と重複する <1.0>、出力はまだ不安定です。トレーニングと推論のパッケージ、テンプレート、動画の予算を確認しましたが、その時点では具体的な原因は特定できませんでした。これはこれまでの数ステップにおける簡易テストの結果であり、今後の正式な微調整については別途効果評価を行う予定です。

7. THE END

その後の正式な微調整においても、このマルチカード対応の入口が引き継がれた。 今回の環境構築における最大の功労者：ds-4.1f（本当に速い）＋ 6astra（的を射たアドバイス）