AMD MI300X로 의료 AI 학습… CUDA 없이 ROCm만으로 LoRA 파인튜닝 성공

의료 분야 AI 모델 MedQA가 NVIDIA CUDA 없이 AMD ROCm 환경에서 파인튜닝되어 공개됐다. 다지선다형 임상 문제에 대해 정답과 의학적 추론 설명을 함께 제공하는 모델로, 데이터 로딩부터 LoRA 어댑터 출력까지 전체 학습 파이프라인이 AMD Instinct MI300X에서 단일 CUDA 의존성 없이 실행된다.

AMD Instinct MI300X는 단일 디바이스에 192GB HBM3 메모리를 탑재한다. 대규모 언어모델 파인튜닝에서 VRAM은 배치 크기와 시퀀스 길이를 결정짓는 가장 큰 제약 자원이다. 192GB 메모리 덕분에 Qwen3-1.7B 모델을 4비트나 8비트 양자화 트릭 없이 fp16 풀 정밀도로 LoRA 학습할 수 있었다.

프로젝트의 핵심 메시지는 HuggingFace 생태계인 Transformers, PEFT, TRL, Accelerate가 ROCm 위에서 그대로 동작한다는 점이다. CUDA에서 돌아가던 학습 코드를 ROCm으로 옮기는 데 필요한 작업은 ROCR_VISIBLE_DEVICES, HIP_VISIBLE_DEVICES, HSA_OVERRIDE_GFX_VERSION 세 가지 환경변수를 설정하는 것뿐이다. 코드 수정도, 커스텀 커널도, CUDA 호환 레이어도 필요하지 않다.

학습 데이터셋으로는 인도 의대 입학시험(AIIMS)과 USMLE 스타일의 문제로 구성된 MedMCQA가 사용됐다. 각 샘플은 임상 질문, A~D 4지선다 보기, 정답 인덱스, 그리고 선택적 텍스트 해설로 이뤄져 있다. 이번 프로젝트에서는 2,000개의 학습 샘플만 사용했으며, MI300X에서 약 5분 만에 전체 학습이 완료됐다.

베이스 모델은 알리바바의 최신 소형 언어모델 Qwen/Qwen3-1.7B다. 17억 파라미터 규모로 임상 추론을 생성하기에 충분한 능력을 갖추면서도 파인튜닝 비용이 저렴하다. LoRA 설정은 r=8, lora_alpha=16, lora_dropout 0.05이며 어텐션 레이어의 q_proj와 v_proj 모듈만 학습 대상으로 지정했다. 결과적으로 전체 15억 4,390만 파라미터 중 약 222만 파라미터(0.1443%)만 학습됐다.

학습 하이퍼파라미터는 2 에폭, 디바이스당 배치 크기 4, gradient_accumulation_steps 4(유효 배치 크기 16), 학습률 2e-4로 구성됐다. 초기 bfloat16 실험에서는 NaN loss가 발생해 fp16으로 전환했고, gradient checkpointing을 적용해 더 작은 GPU에서도 재현 가능하도록 만들었다. 학습률 스케줄은 워밍업이 포함된 코사인 방식을 사용했다.

추론 단계에서는 do_sample=False의 그리디 디코딩에 repetition_penalty 1.1을 적용해 모델이 동일 구문을 반복 출력하는 현상을 막았다. 학습된 모델은 단순히 정답 문자만 출력하지 않고 해당 답안이 왜 정답인지 임상적 근거를 함께 설명한다. 예를 들어 고혈압 응급치료 1차 약제를 묻는 질문에 IV labetalol 또는 IV nitroprusside가 즉각적인 혈압 강하에 적합한 이유를 부연 설명하는 식이다.

프로젝트 진행 중 부딪힌 주요 이슈로는 혼합정밀도에서의 NaN loss, GPU 미인식, ROCm용 bitsandbytes 빌드 부재 세 가지가 있었다. 각각 fp16 전환, ROCm 환경변수 설정, 양자화 자체를 사용하지 않는 방식으로 해결했다. 192GB 메모리 덕분에 양자화 없이도 학습이 가능했다는 점이 결정적이었다.

학습된 LoRA 어댑터 가중치는 HuggingFace Hub의 HK2184/medqa-qwen3-lora 저장소에 공개됐다. 베이스 Qwen/Qwen3-1.7B 모델 위에 PeftModel.from_pretrained로 어댑터를 로드하거나 merge_and_unload로 가중치를 병합해 사용할 수 있다. 이번 프로젝트는 NVIDIA 의존성에서 벗어나도 AMD GPU만으로 의료 분야 AI를 충분히 학습·서빙할 수 있음을 보여준 실증 사례다.