TL;DR
NVIDIA H100 80GB GPUを24基(3ノード)利用した70B LLMの分散学習で発生した主要トラブルと、その解決策をまとめました。
-
SlurmのGPU割り当て: --gpus-per-task=1 が原因で invalid device ordinal エラーが発生。
--gpus-per-node=8
で解決。 - メモリ枯渇: DeepSpeed Zero-3が有効にならず各GPUが78GBを専有。応急策としてQLoRA(4bit量子化)を導入し、約41GBに抑制。
- 語彙サイズ不一致: Tokenizerとモデルの語彙数が異なり CUDA assert エラーが発生。model.resize_token_embeddings() で解決。
- 再現手順: こちら に sbatchスクリプトと設定の雛形を掲載。
環境と前提
本記事の作業は、以下の環境で実施しました。
項目 | 値 | 備考 |
---|---|---|
ハードウェア | ||
ノード数 | 3 | |
GPU | NVIDIA H100 80GB × 24 | 1ノードあたり8基 |
ソフトウェア | ||
ジョブ管理 | Slurm | |
CUDA / NCCL | 12.x / 2.x | |
Python | 3.10 | |
PyTorch | 2.x | |
Transformers | 4.x | |
Accelerate | 0.32.0 | |
PEFT | 0.x | |
bitsandbytes | 0.x |
結論:分散学習 最終設定の要点
数々の試行錯誤の末にたどり着いた、安定稼働のための設定の要点を先にまとめます。
- Slurm ジョブスクリプト: マルチノード分散学習では、--gpus-per-task ではなく --gpus-per-node を指定します。これにより、各ノード内の全プロセスが全てのGPUを認識でき、local_rank に応じた適切なデバイス割り当てが可能になります。
- メモリ効率化: DeepSpeed Zero-3の完全な適用が難しい場合、QLoRA (4bit量子化) と Gradient Checkpointing の併用が有効です。これにより、70Bモデルでも1GPUあたりのメモリ使用量をH100(80GB)の半分程度に抑えることができました。
- 語彙数(vocab_size)の整合性: モデルとTokenizerの語彙数が異なる場合は、学習開始前に model.resize_token_embeddings(len(tokenizer)) を実行し、サイズを一致させることが不可欠です。
- vLLMの安定起動: 推論サーバーとしてvLLMを使用する場合、起動コマンド投入後、APIエンドポイントが有効になるまで数十秒〜数分かかります。sleep と ヘルスチェックのループ を組み合わせ、サーバーが完全に準備完了するのを待ってからリクエストを送信する仕組みを導入しました。
ハマりどころと原因究明の記録
ここからは、実際に遭遇した問題とその解決に至るまでの経緯を時系列で紹介します。
a) Slurmの罠:--gpus-per-task が引き起こした invalid device ordinal エラー
最初に直面した最大の壁は、分散通信テストで発生したCUDAエラーでした。
CUDA failure 101 'invalid device ordinal'
原因は、Slurmのジョブスクリプトで指定した --gpus-per-task=1
にありました。このオプションは、各タスク(プロセス)に単一のGPUを排他的に割り当てるため、どのプロセスからも CUDA_VISIBLE_DEVICES が 0 と認識されてしまいます。結果、全プロセスが同じGPU 0にアクセスしようと競合し、エラーが発生していました。
【解決策】
--gpus-per-node=8
に変更し、各ノード内の全プロセスが8つのGPU全てを可視化できるようにしました。これにより、各プロセスは自身の local_rank(0〜7)を使って、担当すべきGPUを正しく選択できるようになりました。
b) Zero-3不発:各GPUが約78GBメモリに張り付くOOM地獄
分散通信は確立できたものの、70Bモデルの学習を開始すると即座に CUDA out of memory
エラーが発生しました。モニタリングすると、各GPUが約78GBものメモリを確保しており、DeepSpeed Zero-3によるメモリ削減が全く機能していないことが判明しました。
これは、Transformers.Trainer とDeepSpeedの連携設定が不十分で、各プロセスがモデル全体(FP16で約140GB)をロードしようとしていたためです。
【解決策】
根本解決には時間が必要と判断し、応急策として QLoRA(4bit量子化) を導入しました。bitsandbytesライブラリを用いることで、モデルの重みを4bitに量子化し、メモリ使用量を劇的に削減。最終的に1GPUあたり約41GBでの安定学習に成功しました。
c) 語彙サイズ不一致:pad_token に起因する CUDA assert エラー
強化学習(GRPO)の実装を試みた際、Tokenizer(語彙数128,000)とモデル(語彙数128,256)の語彙数が異なっていることが原因で、CUDA device-side assert triggered
エラーに遭遇しました。具体的には、Tokenizerの pad_token のID(128001)が、モデルの埋め込み層の範囲外となっていました。
【解決策】
学習スクリプトの冒頭で model.resize_token_embeddings(len(tokenizer))
を呼び出し、モデル側の埋め込み層のサイズをTokenizerの語彙数に強制的に合わせることで解決しました。
d) vLLMの気まぐれ:サーバー起動後の接続エラー
評価パイプラインでvLLMを導入した際、スクリプトでサーバーを起動した直後にクライアントが接続できず、Connection error
が多発しました。これは、vLLMが起動プロセスを完了し、APIリクエストを受け付け可能になるまでに時間がかかるためでした。
【解決策】
単純な sleep 300
で待機する方法に加え、より堅牢なヘルスチェック処理を実装しました。
# vLLM起動後、ヘルスチェックで準備完了を待つ
for i in {1..10}; do
if curl -s http://127.0.0.1:8000/health >/dev/null 2>&1; then
echo "✅ vLLM server is ready!"
success=true
break
fi
sleep 30
done
再現手順(コピペ用雛形)
Slurmジョブスクリプト (sbatch)
3ノード×8GPU(合計24GPU)を使用する場合の基本的なスクリプトです。
#!/bin/bash
#SBATCH --job-name=llm70b-train
#SBATCH --nodes=3
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=8
#SBATCH --time=12:00:00
# MASTERノードのIPアドレスとポートを環境変数に設定
export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n 1)
export MASTER_PORT=29500
# srunで分散学習スクリプトを実行
srun --label python train.py \
--model_name_or_path "deepseek-ai/DeepSeek-R1-Distill-Llama-70B" \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 2 \
--gradient_checkpointing True \
--fp16 True
QLoRA (4bit量子化) 設定
bitsandbytes を用いた4bit量子化のPythonコードスニペットです。
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch
# 4bit量子化設定
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.float16
)
# モデルのロード時に適用
model = AutoModelForCausalLM.from_pretrained(
"deepseek-ai/DeepSeek-R1-Distill-Llama-70B",
quantization_config=bnb_config,
device_map="auto" # 自動でデバイスに配置
)
評価と限界
学習済みモデルの性能評価には、HLEベンチマークを使用しました。
【重要】本評価は限定的な条件下での参考値です。
- 内部サブセット評価 (n=698): DeepSeek-R1-Distill-Llama-70B の最大コンテキスト長(当時2048トークンと解釈)を超える長文問題がHLEベンチマークに含まれていたため、これらを除外した698問のみで評価を行いました。
- 非互換スコア: 上記の理由により、算出されたスコア(Accuracy: 4.58%)は、全データセットで評価された公式スコアとは互換性がありません。
教訓(チェックリスト)
今回の経験から得られた教訓を、今後のためのチェックリストとしてまとめます。
- Slurmオプション: 分散学習では --gpus-per-node を使うか?
- 量子化: 大規模モデルを扱う際、QLoRAなどのメモリ削減策を検討したか?
- CPUオフロード: モデルのマージなど、GPUメモリを大量に消費するタスクはCPUでの実行を検討したか?
- 語彙サイズ: モデルとTokenizerの語彙数が一致しているか?
- 非同期処理の待機: vLLMなど、サーバー・クライアント型の処理で適切な待機/ヘルスチェック処理を入れたか?
- コンテキスト長: 評価タスクの最大トークン長が、モデルの対応範囲内であるか?
本プロジェクトは、国立研究開発法人新エネルギー・産業技術総合開発機構(NEDO)の「日本語版医療特化型LLMの社会実装に向けた安全性検証・実証」における基盤モデルの開発プロジェクトの一環として行われます。