こんにちは、イノベーションセンターの鈴ヶ嶺です。普段は AI/ML システムに関する業務に従事しています。 本記事では、CUDA 12.8 から追加された Checkpoint API の概要について解説します。 まず、Checkpoint のユースケースやこれまでの NVIDIA CUDA における Checkpoint の試みなどの背景を説明し、新たに追加された CUDA Checkpointing について解説します。 さらに実際に実装し、torchvision や transformers などの CUDA アプリケーションに対して、Checkpoint の検証をしています。 背景 CUDA Checkpointing 実装と検証 cu_check tool 検証 Pytorch Counter torchvision transformers まとめ 背景 Checkpoint は計算途中の内部状態(メモリなど)をディスクなどに保存し、任意のタイミングで計算を再開できる技術です。 想定されるユースケースとして「障害発生時のバックアップ」、「ライブマイグレーション」、「長時間実行される計算の途中結果保存」、「タスクスケジューリングにおけるプリエンプション」、「フォレンジック分析における証拠保全」などが挙げられます。 特に GPU 分野では、昨今の大規模言語モデルに代表される長期間にわたる学習処理の一時保存や、リソース最適化の一環として、一部の GPU プロセスを別の GPU サーバへ移行するといった用途が考えられます。 しかし、これまでの NVIDIA CUDA における Checkpoint 手法としてさまざまな試みはありましたが、アプリケーションの実行環境自体に変更を加えて、CUDA Driver API の呼び出しを傍受しているため完全に透過的なものではありませんでした。 1 2 3 4 5 6 7 8 9 10 一方、2024 年 7 月 NVIDIA の Technical Blog において、 CUDA Driver API として version 555 から Checkpoint が実験的に実装されたことが発表されました。 オープンソースな Checkpoint utility の CRIU (Checkpoint/Restore in Userspace) との組み合わせについても説明されています。 https://developer.nvidia.com/blog/checkpointing-cuda-applications-with-criu/ 以下にそのツールである cuda-checkpoint が公開されています。リポジトリにはツールのバイナリしか置いておらず、バイナリの文字列を調べると cuGetExportTable 11 以外の API 呼び出しが存在しておらず、おそらくドキュメント化されていない関数を利用する形で実装されていました。 https://github.com/NVIDIA/cuda-checkpoint git clone https://github.com/NVIDIA/cuda-checkpoint.git cd cuda-checkpoint strings ./bin/x86_64_Linux/cuda-checkpoint | grep cu # libcuda.so.1 # cuDriverGetVersion # cuGetExportTable 他にも、 MemVerge 社は早期にこの API を利用した AI の学習 PoC を実施して GTC24 などで発表しています。 CUDA 12.x driver enhancements will enable the open-source CRIU project to checkpoint and restart a GPU-based compute node. We'll provide a technical overview and demonstrate this new capability. https://www.nvidia.com/en-us/on-demand/session/gtc24-p63184/ www.youtube.com そして、2025 年 3 月の現時点では Driver version 570 および CUDA 12.8 から CUDA Checkpointing として正式に API が公開されました。 先ほどの check-checkpoint のリポジトリにも一部その API を利用するコードが公開されていますが、 cuda-checkpoint と併用されており、新たに仕様公開されたものに置き換わっていないと思われます。 12 13 そのため、現在 CUDA 12.8 で新たに公開された API を利用したサンプルコードが見当たらない状況になっていると思われます。 次の章で公開されている Checkpoint API を調査し、どのように利用するのかを確認します。 CUDA Checkpointing CUDA Checkpointing の API 一覧が以下になります。 https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__CHECKPOINT.html CUresult cuCheckpointProcessGetState ( int pid, CUprocessState* state ) CUDA プロセスの現在の状態( CU_PROCESS_STATE_RUNNING , CU_PROCESS_STATE_LOCKED , CU_PROCESS_STATE_CHECKPOINTED , CU_PROCESS_STATE_FAILED )を取得します。 CUresult cuCheckpointProcessLock ( int pid, CUcheckpointLockArgs* args ) CUDA プロセスを lock して、以降の CUDA API 呼び出しをブロックします。 CUresult cuCheckpointProcessCheckpoint ( int pid, CUcheckpointCheckpointArgs* args ) GPU メモリの内容を host memory に保存し、CUDA プロセスを checkpoint します。 CUresult cuCheckpointProcessRestore ( int pid, CUcheckpointRestoreArgs* args ) CUDA プロセスをリストアします。状態は CU_PROCESS_STATE_CHECKPOINTED である必要があります。 CUresult cuCheckpointProcessUnlock ( int pid, CUcheckpointUnlockArgs* args ) CUDA プロセスの lock を解除して、CUDA API Call を再開できるようにします。 CUresult cuCheckpointProcessGetRestoreThreadId ( int pid, int* tid ) CUDA プロセスの Thread ID を取得する。 引用: https://arxiv.org/abs/2502.16631 14 上図を参考にすると、実行中のプロセスの Checkpoint は次のように実行します。 cuCheckpointProcessLock cuCheckpointProcessCheckpoint CRIU dump また、保存したものを Restore するには次のように実行します。 CRIU restore cuCheckpointProcessRestore cuCheckpointProcessUnlock 次の章で実際にこれらを動作するコマンドツールを実装して、Pytorch などのアプリケーションの Checkpoint が可能かを確認します。 実装と検証 cu_check tool 次のようにそれぞれの Checkpoint API をサブコマンドとして、特定のプロセス ID に実行するツールを実装します。 cu_check.c #include <cuda.h> #include <stdio.h> #include <stdlib.h> #include <string.h> #define CHECK_CU (func) \ do { \ CUresult res = (func); \ if (res != CUDA_SUCCESS) { \ const char *errName = NULL ; \ const char *errDesc = NULL ; \ cuGetErrorName (res, &errName); \ cuGetErrorString (res, &errDesc); \ fprintf ( stderr , " %s failed: %s %s\n " , #func, errName, errDesc); \ return - 1 ; \ } \ } while ( 0 ) const char * getCUprocessState (CUprocessState state) { switch (state) { case CU_PROCESS_STATE_RUNNING: return "CU_PROCESS_STATE_RUNNING" ; case CU_PROCESS_STATE_LOCKED: return "CU_PROCESS_STATE_LOCKED" ; case CU_PROCESS_STATE_CHECKPOINTED: return "CU_PROCESS_STATE_CHECKPOINTED" ; case CU_PROCESS_STATE_FAILED: return "CU_PROCESS_STATE_FAILED" ; default : return "OTHER_STATE" ; } } int main ( int argc, char **argv) { if (argc < 3 ) { fprintf ( stderr , "usage: %s [state|lock|checkpoint|restore|unlock] <pid> \n " , argv[ 0 ]); return - 1 ; } const char *subcommand = argv[ 1 ]; int pid = atoi (argv[ 2 ]); CHECK_CU ( cuInit ( 0 )); if ( strcmp (subcommand, "state" ) == 0 ) { CUprocessState state; CHECK_CU ( cuCheckpointProcessGetState (pid, &state)); printf ( "state: %s\n " , getCUprocessState (state)); } else if ( strcmp (subcommand, "thread" ) == 0 ) { int threadId = 0 ; CHECK_CU ( cuCheckpointProcessGetRestoreThreadId (pid, &threadId)); printf ( "thread id: %d\n " , threadId); } else if ( strcmp (subcommand, "lock" ) == 0 ) { CUcheckpointLockArgs args = { .timeoutMs = 600000 // 10min timeout }; CHECK_CU ( cuCheckpointProcessLock (pid, &args)); printf ( "locked successfully \n " ); } else if ( strcmp (subcommand, "checkpoint" ) == 0 ) { CHECK_CU ( cuCheckpointProcessCheckpoint (pid, NULL )); printf ( "checkpointed successfully \n " ); } else if ( strcmp (subcommand, "restore" ) == 0 ) { CHECK_CU ( cuCheckpointProcessRestore (pid, NULL )); printf ( "restored successfully \n " ); } else if ( strcmp (subcommand, "unlock" ) == 0 ) { CHECK_CU ( cuCheckpointProcessUnlock (pid, NULL )); printf ( "unlocked successfully \n " ); } else { printf ( "unknown subcommand: %s\n " , subcommand); return - 1 ; } return 0 ; } gcc -I /usr/local/cuda-12. 8 /include cu_check.c -o cu_check -lcuda # install sudo mv cu_check /usr/local/bin # Usage cu_check state < pid > cu_check lock < pid > cu_check checkpoint < pid > cu_check restore < pid > cu_check unlock < pid > 検証 事前に CRIU をインストールします。 curl -LO " http://github.com/checkpoint-restore/criu/archive/v4.0/criu-4.0.tar.gz " tar xvfz criu-4. 0 .tar.gz cd criu-4. 0 / make -j sudo make install Pytorch Counter CUDA Memory 上に保存し、 1 秒ごとに inc される Counter コードでまず検証します。 torch_counter.py import torch, time counter = torch.tensor( 0 , device= 'cuda' ) while True : print (counter) counter.add_( 1 ) time.sleep( 1 ) 次のように実行します。 1 秒ごとに Counter が継続して出力される様子が確認できると思います。 pip install torch python torch_counter.py & sleep 5 PID = $( pgrep -f ' python torch_counter.py ' ) # checkpoint rm -rf tcnt && mkdir -p tcnt cu_check lock $PID cu_check checkpoint $PID sudo criu dump -j -D tcnt -t $PID du -sh tcnt # 755M # restore sudo criu restore -j -D tcnt & while ! pgrep -f ' python torch_counter.py ' > /dev/null 2 >& 1 ; do sleep 1 ; done sudo cu_check restore $PID sudo cu_check unlock $PID sleep 5 kill -9 $PID torchvision 次に、torchvision の ResNet で検証します。 学習途中の状態が Checkpoint され、再開後に学習途中から実行される様子が確認できると思います。 git clone https://github.com/pytorch/examples.git cd examples/imagenet/ pip install -r requirements.txt python main.py -a resnet152 --dummy -j 0 & sleep 20 PID = $( pgrep -f ' python main.py -a resnet152 --dummy -j 0 ' ) # checkpoint rm -rf resnet && mkdir -p resnet cu_check lock $PID cu_check checkpoint $PID sudo criu dump -j -D resnet -t $PID du -sh resnet # 50G # restore sudo criu restore -j -D resnet & while ! pgrep -f ' python main.py -a resnet152 --dummy -j 0 ' > /dev/null 2 >& 1 ; do sleep 1 ; done sudo cu_check restore $PID sudo cu_check unlock $PID sleep 20 sudo kill -9 $PID transformers 最後に、transformers で検証します。 こちらも同様に継続して学習可能である様子を確認できました。 train_bert.py # ref: https://huggingface.co/docs/transformers/training from datasets import load_dataset from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, ) dataset = load_dataset( "yelp_review_full" )[ "train" ].select( range ( 10000 )) tokenizer = AutoTokenizer.from_pretrained( "google-bert/bert-base-cased" ) def tokenize_function (examples): return tokenizer(examples[ "text" ], padding= "max_length" , truncation= True ) small_train_dataset = dataset.map(tokenize_function, batched= True ) model = AutoModelForSequenceClassification.from_pretrained( "google-bert/bert-base-cased" , num_labels= 5 ) trainer = Trainer( model=model, args=TrainingArguments(num_train_epochs= 10000 ), train_dataset=small_train_dataset, ) trainer.train() 次のように実行します。 pip install transformers datasets accelerate python train_bert.py & sleep 60 PID = $( pgrep -f ' python train_bert.py ' ) # checkpoint rm -rf bert && mkdir -p bert cu_check lock $PID cu_check checkpoint $PID sudo criu dump -j -D bert -t $PID --tcp-established du -sh bert # 5.5G # restore sudo criu restore -j -D bert --tcp-established & while ! pgrep -f ' python train_bert.py ' > /dev/null 2 >& 1 ; do sleep 1 ; done sudo cu_check restore $PID sudo cu_check unlock $PID sleep 20 sudo kill -9 $PID まとめ 本記事では CUDA 12.8 で導入された Checkpoint API について解説しました。 また実際の実装を通して、昨今の大規模言語モデルのデファクトなライブラリである transformers などの CUDA アプリケーションに対して、Checkpoint が動作することを確認しました。 今後これらの技術が活用されることで、長期間にわたる学習処理の一時保存や、効率的なリソース最適化が行われることを期待しています。 Takizawa, Hiroyuki, et al. "CheCUDA: A checkpoint/restart tool for CUDA applications." 2009 International Conference on Parallel and Distributed Computing, Applications and Technologies. IEEE, 2009. ↩ Nukada, Akira, Hiroyuki Takizawa, and Satoshi Matsuoka. "NVCR: A transparent checkpoint-restart library for NVIDIA CUDA." 2011 IEEE International Symposium on Parallel and Distributed Processing Workshops and Phd Forum. IEEE, 2011. ↩ Jiang, Hai, et al. "A checkpoint/restart scheme for cuda programs with complex computation states." International Journal of Networked and Distributed Computing 1.4 (2013): 196-212. ↩ Garg, Rohan, et al. "CRUM: Checkpoint-restart support for CUDA's unified memory." 2018 IEEE International Conference on Cluster Computing (CLUSTER). IEEE, 2018. ↩ Jain, Twinkle, and Gene Cooperman. "CRAC: checkpoint-restart architecture for CUDA with streams and UVM." SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020. ↩ Shukla, Dharma, et al. "Singularity: Planet-Scale, Preemptive and Elastic Scheduling of AI Workloads. arXiv. org (February 2022)." arXiv preprint arXiv:2202.07848. ↩ Eiling, Niklas, et al. "Cricket: A virtualization layer for distributed execution of CUDA applications with checkpoint/restart support." Concurrency and Computation: Practice and Experience 34.14 (2022): e6474. ↩ Nukada, Akira, Taichiro Suzuki, and Satoshi Matsuoka. "Efficient checkpoint/Restart of CUDA applications." Parallel Computing 116 (2023): 103018. ↩ Eiling, Niklas, Stefan Lankes, and Antonello Monti. "Checkpoint/Restart for CUDA Kernels." Proceedings of the SC'23 Workshops of the International Conference on High Performance Computing, Network, Storage, and Analysis. 2023. ↩ Yang, Yanning, et al. "On-demand and Parallel Checkpoint/Restore for GPU Applications." Proceedings of the 2024 ACM Symposium on Cloud Computing. 2024. ↩ https://forums.developer.nvidia.com/t/cugetexporttable-explanation/259109/2 ↩ https://github.com/NVIDIA/cuda-checkpoint/blob/6ec728aff032c18c9fb0794a272d94c6adcce508/src/r570-features.c ↩ https://github.com/checkpoint-restore/criu/blob/4b099510b35f98a1f1d6589b1660470402fc1fef/plugins/cuda/cuda_plugin.c ↩ Stoyanov, Radostin, et al. "CRIUgpu: Transparent Checkpointing of GPU-Accelerated Workloads." arXiv preprint arXiv:2502.16631 (2025). ↩