MedQA：AMD ROCm上での臨床AIのファインチューニング—CUDAは不要

医療質問応答は、誤った回答が危険を伴う高リスクなタスクです。しかし、オープンソースの医療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分で完了します。