メモリをスワッピングしながら大規模言語モデル(LLama2)をフルパラメータでファインチューニングできるかどうか?
背景
大規模言語モデルの学習には､大量のメモリ容量が必要になります｡
例えば700億パラメータ(70b)のモデルを16bit変数で読み込むだけでも70 GB程度の容量が必要で､学習にはその10倍程度の容量くらいはあると安心です｡
ただし､学習に必要なGPUは高額(A100の80GBモデルは200万円超)なので､多額の資金が必要になります｡
打開策として､有名なQLoRAのほか､学習に必要なメモリをCPUやSSD(NVME)に移す(オフロード)するテクニックなどがあります｡
しかし数テラバイトのCPUメモリやSSDを増設するにも､それなりの手間とコストがかかります｡
そこで今回は､外付けSSDにメモリスワップさせる形で､大規模言語モデルのメモリをオフロードできないか､無謀な挑戦をしてみます｡
デバイス
今回用いたSSDはこちら｡メモリは4TB､転送速度は最大2GB/sです｡
更に早い規格(Thunderbolt™4: 40Gbps)も存在するようですが､用いるマシンの端子が対応しているか不明だったので､こちらの品にしました｡
GPUなどに比べ､圧倒的にSSDの転送速度が遅い点に注意が必要です｡GPU(や最先端のCPUメモリ)だとテラバイトレベルの帯域があります｡
なので､今回の構成では､読み込み・書き込み速度が1/1000程度に落ちる可能性があります｡
実装
スワップメモリの確保
nvme対応のSSDであれば､単純にdeepspeedのjson設定を"nvme"に変更するだけでOKです(リンク)｡
今回のデバイスはnvme対応ではなかったので､cpuメモリにオフロードしつつ､不足部分をlinuxのスワッピング機能で補う戦略を取ります｡
SSDをフォーマットします｡(今回のケースでは､SSDは/dev/sdb2で認識されました)
sudo mkfs.ext4 /dev/sdb2スワップ領域を作り､有効化します｡
sudo mkswap -c /dev/sdb2
sudo swapon /dev/sdb24TBもあるので､処理にはそれなりの時間がかかります｡
一晩放置したら､無事にスワッピングメモリが追加されていました｡
実装
学習のコード自体は､以下の記事を流用しました｡
最大 150 token程度のテキストを学習させてみます｡
モデルにはLLama2を使いました｡
マシンはRTX3090(24 GB) x1 + CPU RAM 256 GBの構成です｡
メモリの使用状況
・7b (合計メモリ: 約 220 GB) 速度: 13 sec/iter
6 GB(GPU) + 210 GB (CPU) + 0.7 GB (SSD)
→ 7bはメモリスワップなしで動きました｡
ちなみに､1500 tokenだと20GB程度のVRAMを要しました｡
・13b(合計メモリ: 約310 GB) 25 sec/iter
6 GB(GPU) + 220 GB (CPU) + 80 GB (SSD)
→ 256 GBのCPUメモリでは足りず､追加で80GB程度のスワップメモリが必要になりました｡
学習速度は7bと比べると半分程度に落ちました｡
ただ､モデルサイズが大きくなると学習速度は落ちるので､意外と頑張っているしている印象です｡
70b (out of memory)
→cuda out of memoryでした｡VRAMの増強が必要そうです
注意
13bモデルの学習中に､コンピュータがフリーズしてしまいました｡なので､学習は数stepしか試せていません｡
使っているマシンのハードウェアが若干､不安定なことも関連かもしれませ｡
まとめ
メモリスワッピングを使えば､大型のモデルのフルパラメータ ファインチューニングができる可能性を示しました｡
ただし､70bレベルだったりtoken長が大きいと､どうしてもGPUメモリが必要になるようです｡
A100 x2 + 数TBの追加メモリで70bクラスも動くかどうか､気になるところです
スワッピングしすぎるとシステムが不安定になる感じなので､素直にNVMEでオフロードした方が良さそうです｡