Transformers、CUDAの早期初期化による子プロセスの失敗を防ぐtorchao遅延インポート修正を提案

Transformersコミュニティは11月9日、依存パッケージのインポート時に発生する副作用を報告しました。torchaoがインストールされている環境では、transformers.modeling_utilsをインポートするだけでCUDAが早期に初期化され、その後forkで作成されCUDAを使用する子プロセスが失敗する可能性があります。報告者によると、この問題はvLLMの継続的インテグレーションテストで発生しました。共有イメージに含まれるtorchaoにより、その量子化機能を使わないジョブも影響を受けていました。問題報告

報告された環境は、Transformers 5.22.0、torchao 0.18.0、PyTorch 2.15、Linux GPUです。トリガーとなる呼び出しの流れは、量子化器の自動読み込みモジュールから始まります。このモジュールがtorchaoの量子化器をインポートすると、量子化器が実験的なsafetensorsサポートを即座に読み込み、それを介してGPUチェックが呼び出されます。報告には、インポート前後のCUDA初期化状態を比較できる最小再現コードが添えられています。現時点では、この影響は報告されたバージョンの組み合わせに限って捉えるべきです。再現手順と原因の説明

11月10日に提案された修正では、flatten_tensor_state_dictのインポートを、実際にシリアライズを行うget_state_dict_and_metadata()メソッド内へ移動します。また、独立したプロセスで回帰テストを追加し、モデルユーティリティのインポート時にtorchaoが読み込まれず、CUDAも初期化されないことを確認します。作者は関連するテスト30項目が通過し、INT8 safetensorsの保存・再読み込み後も出力が同一だったと報告しています。検証にはPyTorch 2.14.1が使われており、2.15はまだ対象に含まれていません。また、完全な型チェックでは、変更していないファイルにも診断が残っています。修正提案

技術的な要点は、GPU実行環境を初期化するタイミングです。PyTorchの公式ドキュメントでは、親プロセスがアクセラレータを初期化した後にforkすることで起きるエラーを「poison fork」と呼び、CUDAを使う子プロセスにはspawnまたはforkserverを使うよう求めています。遅延インポートにより、意図しない初期化経路を取り除けますが、デプロイ側では引き続きワーカープロセスの起動方式を確認する必要があります。共有推論イメージを使うチームにとって、この事例はオプションのパッケージがプロセスの挙動を変える可能性も示しています。今後は提案のレビュー、修正がマージされるバージョン、元の環境での検証結果を追う必要があります。PyTorchのマルチプロセスに関するドキュメント