MedQA en AMD ROCm: fine-tuning clínico sin depender de CUDA

Resumen

MedQA es una prueba práctica y completa de que el ecosistema de Hugging Face —Transformers, PEFT, TRL y Accelerate— puede ejecutar un pipeline de fine-tuning clínico completo sobre hardware AMD sin una sola dependencia CUDA. El proyecto afinó Qwen3-1.7B con LoRA usando un AMD Instinct MI300X y el dataset MedMCQA, entregando tanto la letra de la respuesta correcta como una explicación clínica razonada.

Por qué este experimento importa

En aplicaciones médicas las consecuencias de un error pueden ser graves; un modelo que escoge con confianza una opción equivocada no solo está equivocado, puede ser peligroso. A la vez, la mayoría del trabajo open source en IA asume hardware NVIDIA y CUDA como estándar. Eso limita opciones para organizaciones y centros de investigación, incluyendo muchos en América Latina, que podrían preferir o disponer de infraestructura AMD por costo, disponibilidad o estrategia de compras.

Este proyecto desafía esa suposición: demuestra que se puede realizar fine-tuning útil y reproducible en ROCm sin cambios en el código de entrenamiento, simplemente ajustando variables de entorno. Además, el MI300X ofrece 192 GB de HBM3 en un solo dispositivo, lo que relaja restricciones de memoria que normalmente obligan a recurrir a cuantización agresiva o configuraciones complejas.

Hardware y compatibilidad: AMD ROCm y MI300X

El experimento usó un AMD Instinct MI300X con 192 GB de HBM3. Esa gran capacidad de memoria permite entrenar en fp16 sin tener que pasar a 4-bit u 8-bit, lo que simplifica el pipeline y preserva la estabilidad numérica.

Importante: la misma base de código de Hugging Face funcionó en ROCm con solo tres variables de entorno, sin necesidad de modificar el código, compilar kernels personalizados o instalar “shims” de compatibilidad con CUDA:

• ROCR_VISIBLE_DEVICES = “0”
• HIP_VISIBLE_DEVICES = “0”
• HSA_OVERRIDE_GFX_VERSION = “9.4.2”

Con eso, las librerías y scripts que normalmente corren en CUDA se ejecutaron en ROCm.

Dataset: MedMCQA

El conjunto de datos empleado fue MedMCQA, un dataset de preguntas de opción múltiple derivado de exámenes de ingreso médicos (estilo AIIMS/USMLE). Cada ejemplo contiene:

• Una pregunta clínica
• Cuatro opciones (A–D)
• El índice de la respuesta correcta
• Una explicación de texto libre (campo exp), cuando está disponible

Para este proyecto se utilizó una porción pequeña y deliberada: 2,000 muestras para entrenamiento. Con esa muestra reducida se buscó mostrar que un fine-tuning significativo es posible de manera ágil: el ajuste tomó aproximadamente 5 minutos en el MI300X.

Modelo base y formato de prompt

El modelo base es Qwen/Qwen3-1.7B, un LLM compacto (1.7B parámetros) que equilibra costo y capacidad para tareas de razonamiento clínico. Se cargó con trust_remote_code=True desde Hugging Face Transformers.

La consistencia del formato de prompt fue clave. Todos los ejemplos de entrenamiento y las consultas de inferencia usan la misma plantilla:

Question: {question}

Options: A) {opa} B) {opb} C) {opc} D) {opd}

Answer: {answer_letter}) {answer_text}

Explanation: {explanation}

Durante el entrenamiento el modelo ve la secuencia completa, incluida la respuesta y la explicación. En inferencia se proporciona todo hasta ”### Answer:\n” y se deja que el modelo complete.

Fine-tuning con LoRA (PEFT)

En lugar de ajustar los 1.5B parámetros del modelo, se aplicó LoRA (Low-Rank Adaptation) mediante la biblioteca PEFT. LoRA inserta matrices de adaptación de bajo rango en las capas de atención mientras mantiene fijos los pesos base, reduciendo drásticamente el número de parámetros entrenables.

Configuración utilizada (resumida): r=8, lora_alpha=16, lora_dropout=0.05, target_modules=[“q_proj”,“v_proj”], bias=“none”. Con esa configuración solo ~2.2 millones de parámetros fueron entrenables frente a los 1.543.901.184 parámetros totales (aprox. 0.1443% entrenable). Esto mantiene baja la memoria y acelera el entrenamiento.

Parámetros de entrenamiento

Se emplearon argumentos de Transformers adaptados a un entrenamiento corto y reproducible:

• Epochs: 2
• per_device_train_batch_size: 4
• gradient_accumulation_steps: 4 (batch efectivo = 16)
• learning_rate: 2e-4
• fp16 = True, bf16 = False (se observó NaN con bfloat16 en experimentos previos)
• eval_strategy y save_strategy: “epoch”
• load_best_model_at_end = True
• gradient_checkpointing = True (reduce memoria a costa de más cómputo)
• optim = “adamw_torch”
• lr_scheduler_type = “cosine” con warmup_ratio = 0.05
• report_to = “none”

La elección de fp16 y un scheduler coseno con warmup buscó estabilidad y mejor convergencia en entrenamientos cortos.

Pipeline de entrenamiento y artefactos

El pipeline usó DataCollatorForSeq2Seq, Trainer de Transformers y la integración PEFT. Al finalizar, la carpeta ./outputs contiene los pesos del adaptador LoRA y el tokenizador. Esos archivos son compactos (unos pocos megabytes) frente a un checkpoint completo de varios gigabytes.

El repo, demo y modelo están disponibles públicamente: modelo en Hugging Face Hub (HK2184/medqa-qwen3-lora), demo en HuggingFace Spaces y código en GitHub (MedQA-Medical-AI-on-AMD-ROCm).

Inferencia y despliegue

En tiempo de inferencia se carga el modelo base Qwen3-1.7B en fp16 con device_map=“auto”, se adjunta el adaptador LoRA desde ./outputs y se pasa el prompt hasta ”### Answer:”. La generación emplea decodificación greedy (do_sample=False). El tokenizador se ajusta para que pad_token sea igual a eos_token.

El flujo permite servir el modelo manteniendo un footprint de almacenamiento reducido, ya que solo el adaptador necesita distribuirse además del modelo base accesible desde Hub.

Relevancia para América Latina

Para equipos y laboratorios en América Latina este experimento es relevante por varias razones:

• Reduce la barrera de entrada: demuestra que no es obligatorio depender de CUDA/NVIDIA para proyectos serios de LLM en medicina.
• Flexibilidad de infraestructura: si una institución ya cuenta con servidores AMD, puede aprovecharlos con ROCm sin reescribir pipelines.
• Economía de recursos: usar LoRA permite fine-tunings rápidos y de bajo costo energético, útil para pruebas de concepto en entornos con recursos limitados.

Conclusión

MedQA es una demostración práctica de que se puede realizar fine-tuning clínico efectivo en hardware AMD usando ROCm y LoRA, manteniendo el flujo de trabajo familiar de Hugging Face sin cambios de código. Para proyectos médicos y de salud en la región, esto abre alternativas a la dependencia exclusiva de CUDA y facilita experimentación rápida y reproducible.

Para quienes quieran replicar o profundizar, los artefactos públicos (modelo, demo y código) y la configuración exacta están disponibles en los enlaces del proyecto en Hugging Face y GitHub.