医療質問応答は、誤った回答が危険を伴う高リスクなタスクです。しかし、オープンソースの医療AIの多くはNVIDIA GPUとCUDAを前提としています。MedQAプロジェクトはこの前提に挑戦し、AMDハードウェア上で高性能で説明可能な臨床QAモデルを構築する方法を示します。
本プロジェクトでは、AMD Instinct MI300X GPUとROCmソフトウェアスタックを使用し、AlibabaのQwen3-1.7BモデルをLoRAでファインチューニングします。MI300Xは192GBのHBM3メモリを搭載しており、量子化を一切行わずにfp16完全精度でトレーニング可能です。トレーニングデータはMedMCQAデータセット(インドの医学入学試験に基づく多肢選択問題)から2,000サンプルを使用し、MI300X上で約5分で完了します。
ファインチューニングにはLoRA技術を採用し、PEFTライブラリを介してアテンション層に低ランク行列を注入します。LoRA設定はr=8、lora_alpha=16、lora_dropout=0.05、ターゲットモジュールはq_projとv_projです。トレーニングパラメータは220万個(全体の0.14%)のみで、メモリ使用量とトレーニング時間を大幅に削減します。トレーニング引数にはfp16混合精度、勾配チェックポイント、コサイン学習率スケジュール、実効バッチサイズ16などを使用。初期のbfloat16実験でNaN損失が発生しましたが、fp16に切り替えることで完全に解決しました。
推論時には、ベースモデルをロードしてLoRAアダプタをアタッチし、貪欲デコーディング(do_sample=False)と繰り返しペナルティを使用して生成します。例えば、「高血圧緊急症の第一選択治療は?」という質問に対し、モデルは「B) IVラベタロールまたはIVニトロプルシド」と出力し、静脈内投与が迅速に血圧を下げる理由を説明します。単なる選択肢の文字列ではなく、臨床的説明を提供する点が有用です。
プロジェクトではいくつかの課題に直面しました:GPUが検出されない問題は環境変数ROCR_VISIBLE_DEVICES、HIP_VISIBLE_DEVICES、HSA_OVERRIDE_GFX_VERSIONを設定して解決;bitsandbytesの非互換性は、MI300Xの大容量VRAMにより量子化を完全に省略することで回避;トークナイザのパディング設定ミスによるガベージ出力はpad_tokenをeos_tokenに設定して修正;トレーナーの評価エラーはTransformersのバージョン不一致によるもので、バージョンを固定して解決。
結果として、トレーニング可能パラメータは約220万、トレーニング時間は5分、使用データセットサイズは2,000サンプル、ベースライン精度は約45%でした。ユーザーはGitHubリポジトリをクローンし、依存関係をインストールしてトレーニングスクリプトを実行するか、Gradio UIを起動できます。今後の計画には、完全なMedMCQAコーパス(約18万質問)でのトレーニング、信頼度スコアの追加、検索拡張生成(RAG)の統合、評価ベンチマークの構築が含まれます。
MedQAプロジェクトは、AMDハードウェア上で有能で説明可能な医療AIを構築することが可能であり、かつ簡単であることを示しました。HuggingFaceエコシステムのROCm互換性は優れており、MI300Xの大容量メモリは量子化という工学的問題を排除し、LoRAにより1.7Bモデルのファインチューニングが5分で完了します。