Zespół badaczy z Hugging Face przedstawił MedQA — model sztucznej inteligencji do odpowiadania na pytania medyczne, który został wytrenowany w całości na sprzęcie AMD z wykorzystaniem technologii ROCm, bez konieczności używania CUDA. Szczegóły projektu pokazują, jak osiągnąć efektywne trenowanie modeli AI poza ekosystemem NVIDIA.
MedQA bazuje na modelu Qwen3–1.7B dostrojonym techniką LoRA do odpowiadania na pytania z egzaminów medycznych. Całość została zrealizowana na akceleratorze AMD Instinct MI300X wyposażonym w 192 GB pamięci HBM3, co pozwoliło na trenowanie bez kompresji modelu.
Kluczowe wnioski
- Model MedQA został wytrenowany wyłącznie na sprzęcie AMD ROCm bez użycia CUDA — tylko trzy zmienne środowiskowe wystarczyły do uruchomienia standardowego kodu treningowego.
- Akcelerator AMD Instinct MI300X z 192 GB pamięci HBM3 umożliwił trenowanie w pełnej precyzji fp16 bez konieczności kwantyzacji 4-bit czy 8-bit.
- Wykorzystanie techniki LoRA pozwoliło na trenowanie jedynie 2,2 mln parametrów z 1,5 mld całkowitych, co zajęło około 5 minut na 2000 przykładach treningowych.
- Ekosystem Hugging Face (Transformers, PEFT, TRL, Accelerate) działa bezproblemowo na ROCm bez modyfikacji kodu.
- Model osiąga spójne wyniki w odpowiadaniu na pytania medyczne, zwracając zarówno odpowiedź jak i uzasadnienie kliniczne.
Technologia AMD ROCm jako alternatywa dla CUDA
Projekt celowo wykorzystał sprzęt AMD, aby udowodnić, że trenowanie modeli AI nie musi być uzależnione od ekosystemu NVIDIA. AMD Instinct MI300X oferuje znaczną przewagę pod względem pamięci — 192 GB HBM3 w jednym urządzeniu to więcej niż większość kart NVIDIA dostępnych na rynku.
Kluczowe było to, że cały proces nie wymagał żadnych modyfikacji kodu. Standardowy pipeline treningowy z bibliotek Hugging Face działał natywnie po ustawieniu trzech zmiennych środowiskowych:
ROCR_VISIBLE_DEVICES = "0"HIP_VISIBLE_DEVICES = "0"HSA_OVERRIDE_GFX_VERSION = "9.4.2"
Dataset i architektura modelu
Jako dane treningowe wykorzystano MedMCQA — dataset zawierający pytania wielokrotnego wyboru z indyjskich egzaminów medycznych (AIIMS, w stylu USMLE). Każdy przykład zawiera pytanie kliniczne, cztery opcje odpowiedzi A-D, prawidłową odpowiedź oraz opcjonalne wyjaśnienie.
Model bazowy Qwen3–1.7B został wybrany ze względu na kompaktowy rozmiar (1,7 mld parametrów) przy zachowaniu zdolności do generowania spójnego rozumowania klinicznego. Technika LoRA (Low-Rank Adaptation) pozwoliła na dostrojenie tylko wybranych warstw uwagi, pozostawiając podstawowe wagi modelu zamrożone.
Konfiguracja i wyniki trenowania
Trenowanie wykorzystywało następującą konfigurację LoRA:
- Ranga r=8, alfa=16
- Dropout 0.05
- Moduły docelowe: q_proj, v_proj
- Tylko 2,228,224 trenowalne parametry (0.1443% całego modelu)
Proces treningowy trwał 2 epoki z batch size 4 i akumulacją gradientu 4 (efektywny batch size 16). Wykorzystano precyzję fp16, harmonogram cosine z warmupem oraz gradient checkpointing dla optymalizacji pamięci.
Model MedQA demonstruje, że wysokiej jakości trenowanie modeli medycznych AI jest możliwe poza ekosystemem CUDA, otwierając nowe możliwości dla badaczy pracujących ze sprzętem AMD.
