はじめに
Turing CTO室に所属している東京科学大学(Institute of Science Tokyo)の藤井です。
本記事は、LLM, VLM開発の裏で行われるリアルなデバッグ作業の様子を紹介します。
LLM, VLMの開発の裏には本記事で紹介するような地道なデバッグ作業が多数あるのですが、なかなかその実態が伝わっていないように思います。できるだけ詳細にデバッグ作業の様子を記しましたので、実際の現場で行われている作業を追体験いただけますと幸いです。
Background
まず、デバッグ作業を行う前の背景と経緯について説明します。
学習に使用している依存関係の継続的なupdateのために、学習ライブラリが新しいversionのPyTorch、CUDA Toolkitで動作するかどうか常に検証しています。
本記事では、その作業を行う中で問題に直面し、何が問題なのか突き止めるまでの過程を紹介しています。
PyTorch
PyTorch公式からリリースされているstable versionの最新版は、執筆時点(2025/10/03)ではPyTorch 2.8.0となっています。こちらは、 CUDA Toolkit 12.6, 12.8, 12.9向けにbuildされており、先日リリースされたCUDA Toolkit 13.0向けにbuildを行ったPyTorchは配布されていません。
stable version(2.8.0)の対応状況 (pytorch.orgより)
また、PyTorch 2.8.0をCUDA Toolkit 13.0でself buildするのにも問題が発生します。
CUDA Toolkit 13.0では、12.xとはlibパスが異なるものが複数存在し、build scriptを変更する必要があり、そのままではbuildすることができません。
CUDA Toolkit 13.0に対応したPyTorchを利用するには大きく分けて以下の3つの方法があります。
- PyTorch 2.8.0 or 2.9.x からself buildによるsource installを行う方法
- PyTorch Nightlyを利用してinstallを行う
- NGC PyTorchを利用する
1の方法は、PyTorch 2.8.0系を利用する場合、build設定の更新などをすべて自前で行う必要があり、作業コストが多く、ミスが発生しやすいため、できれば行いたくありません。加えて、stable versionではない、2.9.xを利用してsource buildすることは、distributed backendにmpiを利用したい等の理由や、Nightlyでサポートされていない組み合わせのLibraryを利用したい場合を除き、buildにかかる時間と手間を考えますと避けたいのが実情です。そのため、選択肢2のNightlyを利用する方が1よりは望ましいといえます。
では、選択肢2が、最良でしょうか？
検証を行う作業者がsudo等の権限を有している環境の場合はNVIDIA Driverをupdateできるため、選択肢2でも良いかも知れませんが、学習環境がスパコンなどの場合、NVIDIA Driverを気軽にupdateすることが出来ないため選択肢2も避けた方が良いでしょう。
CUDA Toolkit 13.xでbuildされているPyTorchを12.xをサポートしているDriver環境で動作させるとwarningが出るように、場合によっては望ましくない挙動をする可能性があります。
そこで、本検証では、クラスターのDriver versionをイジる必要がない、選択肢3のNGC Containerを利用することにしました。
pytorch.orgにおける対応状況 (2025/10/04)
Singularity
ABCIやTSUBAME等のスーパーコンピュータ上で、コンテナを利用するには、SingularityやApptainerを利用します。今回のブログでは、singularityを利用します。
NVIDIA NGC PyTorchより適切なContainerを選択します。
今回は、PyTorch 2.9.0a0+50eac811a6かつ、CUDA Toolkit 13.0.1に対応している25.09を利用することにしました。
singularityのsifファイルを作成するために、以下のようにdefファイルを作成します。
Bootstrap: docker
From: nvcr.io/nvidia/pytorch:25.09-py3
%post
pip uninstall -y <uninstalled packages>
pip install --no-cache-dir <installed packages>
次に、以下のコマンドでbuildを行います。
cd /local
mkdir -p tmp
export export SINGULARITY_TMPDIR=/local/tmp
singularity build ngc-pytorch-25.09.sif ngc-pytorch-25.09.def
build時は、NFS、Lustreなどの共有ストレージ上ではなく、GPUノードのローカルストレージにて処理を行っています。また、singularity build時のtmpディレクトリもローカルストレージになるように環境変数を設定しています。
Megatron-LM
上記で用意したNGC PyTorch 25.09を利用して、大規模LLM開発で利用されるライブラリMegatron-LMにて検証のための学習を実施し、checkpoint save/loadが正しく行えるか(= 以前のversionと同様のloss, grad-normの推移をするか、FLOP/s/GPUが低下していないか、checkpoint saveした状態からloadを行い復帰できるか)を検証しました。
その際、以下のようなエラーに直面しました。
megatron/core/dist_checkpointing/strategies/torch.py", line 845, in get_reformulation_metadata
[rank0]: ckpt_global_shape = ckpt_metadata.mcore_data[sh_ten.key][
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: AttributeError: 'Metadata' object has no attribute 'mcore_data'. Did you mean:
エラーの詳細については割愛しますが、概要としては、PyTorch DCP(Distributed Checkpoint)において保存される.metadataを読み込んだ際、本来あるべき.mcore_dataというフィールドが存在しないという問題です。
これにより、学習が停止した際、その後、学習を再開することができなくなっています。
この問題を受けて、(1) Megatron-LMの実装ミス (2)checkpoint形式のミス (3)NGC PyTorchのPyTorch実装の仕様変更のいずれかを原因と考え、調査を開始することにしました。
以下では、デバッグ作業の様子を時系列順に説明します。
Debug
調査方針
すでに述べたように問題は、本来あるべきmetada情報が存在しないことです。
原因として考えられる3つの要因について順に検討してみましょう。
1. Megatron-LMの実装ミス
以下に示すように、PyTorchの元々の実装には、Distributed CheckpointのMetadataに.mcore_dataというフィールドはありません。そのため、単純に実装ミスにより存在しない、もしくは別の場所にあるべき情報に誤ってアクセスしてしまっている可能性が考えられます。
もちろん、Megatron-LMは入念にCIにより検査されていますが、時折、バグが存在することがあり、このようなミスである可能性は捨てきれません。そのため、.mcore_dataというフィールドが本当に、metadataに直接存在するべきであるのか,checkpoint save時にも同じ処理が実施されているのかを調査することにしました。
2. checkpoint形式のミス
PyTorchの仕様変更により、metadataの情報の持ち方が変更になった可能性を検討しました。
こちらのcommitのようにuse_collectivesという新しい引数が追加され、global metadaではなくlocalなrankごとのmetadaを作成する機能が実装されつつあったため、その変更の影響を受けた可能性を考えました。関係しそうなファイルの変更履歴のうち、PyTorch 2.7.1以降に相当するものを確認し、影響を及ぼしそうな変更がないかチェックすることにしました。
3. PyTorchの仕様変更
2と似ていますが、checkpoint save/loadの際にやりとりされる形式それ自体ではなく、他のPyTorchの仕様変更に巻き込まれた形で影響が発生した可能性を考慮して、以下のリンクに示すようにPyTorch 2.7.1と2.9.0-rc6の間のめぼしい変更を調査することにしました。
調査ログ
コンテナ内のpackage
動作確認ができるている環境では、PyTorchをソースビルドしていたため、NGC PyTorchにあらかじめ入っているnvidia-modeloptなどが影響を及ぼしている可能性を考え、こちらの可能性を潰すためにsingularityのsif fileのbuildし直しを実施しました。
これには、過去にHuggingFace形式からMegatron-LM形式へのcheckpoint convertスクリプトにて問題が発生した際の原因に、huggingface transformersと特定のversionのnvidia-modeloptの間の相性が悪く影響を及ぼしたことがあったためです。
以下のように、uninstallするライブラリを設定することでmodeloptの影響を排除した検証が可能です。
Bootstrap: docker
From: nvcr.io/nvidia/pytorch:25.09-py3
%post
pip uninstall -y nvidia-modelopt nvidia-modelopt-core
pip install --no-cache-dir <install-packages>
こちらの環境で学習とcheckpoint saveの検証を行ったところ。loss/grad normへの影響が確認されないこと、そして依然としてmetadataがなくなる問題が解決しなかったため、こちらが問題ではないと結論づけました。
checkpoint save時
checkpointをsaveするときに、metadataに.mcore_dataを付与できているのかを調査しました。mcore_dataを代入しているのは、以下であるので、create_global_planについて、(1)どこで呼び出されているのか(2)呼び出された後の形式は期待されるものなのかの2点を調査しました。
調査したところ、以下で呼び出されていることが判明しました。
そこで、Loggerをそれぞれの呼び出しの前に設定し、問題が発生している学習設定ではどちらのパスを通っているのか、もしくはどちらも通っていないのか検証を行いました。
その結果、以下のパスを通っていること、そして、global_metadataにアクセスすると.mcore_dataがきちんと存在することが判明しました。
global_metadata.mcore_data={'decoder.layers.self_attention.linear_proj.weight': {},
'decoder.layers.self_attention.linear_qkv.weight': {},
'decoder.layers.mlp.linear_fc1.weight': {},
'decoder.layers.mlp.linear_fc2.weight': {},
'optimizer.state.exp_avg.output_layer.weight':
{'nd_reformulated_orig_global_shape': (151936, 4096)},
'optimizer.state.exp_avg_sq.output_layer.weight':
{'nd_reformulated_orig_global_shape': (151936, 4096)},
'optimizer.state.fp32_param.output_layer.weight':
...
save_state_dict_async_planにおいては、正しくmetadataが設定できていることが判明したため、その後、保存されるまでの流れを追い、問題がないのかを検証することにしました。
処理の過程で.mcore_dataが除去されていないか、追跡したところ、以下のwriteを実施する箇所まで.mcore_dataが存在することを確認しました。
実は、この調査は不十分だったことが後に判明するのですが、当初の調査ではcheckpoint save側には問題はないと考え、checkpoint load時に問題が発生しているものとみなし、調査対象を切り替えました。
checkpoint load時
checkpoint loadは.mcore_dataへのアクセスが失敗するget_reformulation_metadata()関数から調査を始めました。
まず、想定されているcheckpointディレクトリをきちんと参照できているかを調査しました。
fs_reader = _get_filesystem_reader(checkpoint_dir)
print(f"DEBUG: fs_reader.path={fs_reader.path}", flush=True)
上記のように、readerが参照しているパスを表示するなど、1つ１つ想定と異なる点がないか確認を行いました。しかし、特に想定外の点は発見できませんでした。
そこで、保存されているデータを、別のスクリプトで読み込み、保存データそれ自体を確認することにしました。以下のような簡単なスクリプトを作成し、.metadataが正しく保存されているか確認を行いました。
import argparse
import pickle
from pprint import pprint
import sys
import os
def main():
parser = argparse.ArgumentParser(
description="Inspect pickle-based .metadata file and print available keys."
)
parser.add_argument(
"--path", required=True, help="Path to the .metadata file"
)
args = parser.parse_args()
if not os.path.exists(args.path):
print(f"Error: File not found: {args.path}", file=sys.stderr)
sys.exit(1)
# Load pickle file
try:
with open(args.path, "rb") as f:
metadata = pickle.load(f)
except Exception as e:
print(f"Failed to load pickle file: {e}", file=sys.stderr)
sys.exit(1)
print(f"Loaded object type: {type(metadata)}")
if isinstance(metadata, dict):
mcore_data = metadata.get("mcore_data", None)
if mcore_data is not None:
print("\nFound 'mcore_data' key in metadata. Dumping its contents:\n")
else:
print("\nMetadata is not a dict. Dumping object directly:\n")
print(metadata.mcore_data)
if __name__ == "__main__":
main()
その結果、checkpoint save側の実装を見る限りでは問題がないと判断したのにも関わらず、.mcore_dataが実際に保存されているデータには存在しないことが判明しました。
これを受けて、checkpoint save側を再度検証することにしました。
追加調査
調査を進めたところ、以下でsuper().finish(metadata, results)を呼び出す直前まではmetadataに異常がないことが判明しました。
そのため、問題が発生しているとするとMegatron-LM側ではなく、PyTorch側ということが判明しました。そのため調査対象をtorch/distributed/checkpoint/filesystem.py::L731-754とし、さらなる調査を行いました。
すると、関数内の処理の1行目である、metadata = dataclasses.replace(metadata, version=CURRENT_DCP_VERSION)で.mcore_dataが欠落することが判明しました。
datasetclasses.replaceで.mcore_dataが欠落する理由は、.mcore_dataはMetadaクラスのデフォルトのフィールドではなく、後から追加されたものだからです。
調査が完了した後、PyTorchのGitHubにてIssueとして報告しようとしたところ、同様の報告がなされていることを発見しました。こちらに詳しく状況がまとまっています。
以上の調査により問題を解決するには、以下のような変更を対象のPyTorchに施せば良いことが判明しました。
さいごに
本記事では、できるだけ省略せずに、日常的に行っているLLM, VLM開発におけるデバッグ作業の様子を紹介しました。LLM、VLM研究開発では、NLP(自然言語処理)、CV(Computer Vision)の知識も必要ですし、最新のモデル動向に関する知識、分散学習の知識等ももちろん必要ですが、それに加えて上記のようなデバッグを行うSoftware Engineering能力も必要です。
研究開発の裏では、このような泥臭い作業が大量に行われており、学習ライブラリが安定稼働しているのはこれらの作業を社内/社外(=OSS Communityなど)の誰かが行っているからに過ぎません。
本ブログ記事から少しでも、現場の感覚を味わっていただけますと幸いです。
Discussion