一份完整的 LoRA 微调 Qwen3-1.7B 在 MedMCQA 数据集上的操作指南,基于 AMD MI300X 构建,专为 lablab.ai 上的 AMD 开发者黑客松打造。
项目构想
医学问答是一项风险极高的任务。一个模型如果在临床选择题上自信地选错答案,那不仅仅是错误——而是危险。与此同时,大多数开源医学 AI 项目都假设你拥有一块 NVIDIA GPU。CUDA 是默认选项,其他一切都被视为次要考虑。
本项目挑战了这一假设。
MedQA 是一个基于 LoRA 微调的临床问答模型,完全在 AMD 硬件上使用 ROCm 构建。它接收一道医学选择题,并返回正确的答案选项字母以及基于推理的临床解释。整个训练流程——从数据加载到适配器导出——都在 AMD Instinct MI300X 上运行,完全不依赖 CUDA。
- 🤗 HuggingFace Hub 上的模型:HK2184/medqa-qwen3-lora
- 🚀 在线演示:HuggingFace Spaces
- 💻 GitHub:MedQA-Medical-AI-on-AMD-ROCm
为什么选择 AMD ROCm?
AMD Instinct MI300X 是一款卓越的硬件:单设备拥有 192 GB 的 HBM3 内存。对于大语言模型微调而言,显存往往是关键约束——它决定了批次大小、序列长度,以及你是否需要进行量化。凭借 192 GB 的可用显存,我们在全 fp16 精度下使用 LoRA 训练了 Qwen3-1.7B,无需任何 4-bit 或 8-bit 的量化技巧。
更重要的是,目标是证明 HuggingFace 生态系统——Transformers、PEFT、TRL、Accelerate——能够在 ROCm 上无缝运行。事实证明,确实如此。在 CUDA 上运行的同一套训练代码,只需设置三个环境变量,就能在 ROCm 上运行:
os.environ["ROCR_VISIBLE_DEVICES"] = "0"
os.environ["HIP_VISIBLE_DEVICES"] = "0"
os.environ["HSA_OVERRIDE_GFX_VERSION"] = "9.4.2"
仅此而已。无需修改代码,无需自定义内核,无需 CUDA 兼容性适配层。
数据集:MedMCQA
MedMCQA 是一个大规模医学选择题数据集,源自印度医学入学考试(AIIMS、USMLE 风格)。每条数据包含:
- 一个临床问题
- 四个选项(A–D)
- 正确答案索引
- 可选的自由文本解释(exp 字段)
本项目使用了 2,000 条训练样本——特意选取了一个较小的数据量,以证明有意义的微调可以快速完成。在 MI300X 上,训练耗时约 5 分钟。
模型:Qwen3-1.7B
基础模型是 Qwen/Qwen3-1.7B——阿里巴巴最新推出的小规模语言模型。该模型拥有 17 亿参数,体积紧凑,微调成本低廉,但能力足以生成连贯的临床推理。它支持 `trust_remote_code=True`,并能通过 HuggingFace Transformers 库干净地加载。
提示词格式
提示词格式的一致性对于指令微调至关重要。每个训练样本和每次推理调用都使用相同的模板:
### Question:
{question}
### Options:
A) {opa}
B) {opb}
C) {opc}
D) {opd}
### Answer:
{answer_letter}) {answer_text}
### Explanation:
{explanation}
训练时,模型会看到包含答案和解释的完整序列。推理时,我们提供到 `### Answer:\n` 为止的所有内容,然后让模型从该处开始补全。
使用 LoRA 进行训练
我们没有微调全部 15 亿参数,而是通过 PEFT 库使用了 LoRA(低秩适配)。LoRA 会在注意力层中注入可训练的小型秩分解矩阵,同时保持基础权重冻结不变。
LoRA 配置
from peft import LoraConfig, get_peft_model, TaskType
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=8,
lora_alpha=16,
lora_dropout=0.05,
target_modules=["q_proj", "v_proj"],
bias="none",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 2,228,224 || all params: 1,543,901,184 || trainable%: 0.1443
模型 15 亿参数中,仅有约 220 万个参数被训练。这保持了较低的内存使用量,并实现了快速训练。
训练参数
from transformers import TrainingArguments
args = TrainingArguments(
output_dir="./outputs",
num_train_epochs=2,
per_device_train_batch_size=4,
gradient_accumulation_steps=4, # effective batch size = 16
learning_rate=2e-4,
fp16=True,
bf16=False,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
gradient_checkpointing=True,
optim="adamw_torch",
warmup_ratio=0.05,
lr_scheduler_type="cosine",
report_to="none",
)
几点值得注意:
- `fp16=True, bf16=False`——我们使用标准的 fp16。在早期使用 bfloat16 的实验中,我们遇到了 NaN 损失;切换到 fp16 后问题完全解决。
- `gradient_checkpointing=True`——用计算换内存。鉴于 MI300X 拥有 192 GB 显存,这并非严格必要,但为了在较小 GPU 上实现可复现性,这是良好的实践。
- `gradient_accumulation_steps=4`——物理批次大小为 4,有效批次大小为 16。
- 带预热的余弦学习率调度——对于短时间训练运行,比固定调度收敛更平滑。
完整训练循环
from transformers import DataCollatorForSeq2Seq, Trainer
collator = DataCollatorForSeq2Seq(
tokenizer,
model=model,
padding=True,
pad_to_multiple_of=8,
)
trainer = Trainer(
model=model,
args=args,
train_dataset=train_ds,
eval_dataset=val_ds,
data_collator=collator,
)
trainer.train()
# Save adapter + tokenizer
model.save_pretrained("./outputs")
tokenizer.save_pretrained("./outputs")
训练完成后,`./outputs` 目录中包含 LoRA 适配器权重——仅几 MB 的文件,而非完整的数 GB 模型检查点。
推理
推理时,我们加载基础模型,附加 LoRA 适配器,并可选择合并权重:
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
tokenizer = AutoTokenizer.from_pretrained("./outputs", trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
base_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-1.7B",
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True,
)
model = PeftModel.from_pretrained(base_model, "./outputs")
model.eval()
生成过程使用贪婪解码(`do_sample=False`),并带有重复惩罚以防止模型循环重复:
def generate(prompt, model, tokenizer):
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=200,
do_sample=False,
temperature=1.0,
repetition_penalty=1.1,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.eos_token_id,
)
new_tokens = output[0][inputs["input_ids"].shape[-1]:]
return tokenizer.decode(new_tokens, skip_special_tokens=True)
示例输出
问题:以下哪项是高血压急症的一线治疗方案?
A) Oral amlodipine
B) IV labetalol or IV nitroprusside
C) Sublingual nifedipine
D) IM hydralazine
模型输出:
B) IV labetalol or IV nitroprusside
Explanation:
Intravenous labetalol (beta-blocker) or nitroprusside rapidly reduces blood
pressure in emergency settings. Oral agents act too slowly for hypertensive
emergencies requiring immediate BP control to prevent end-organ damage.
模型不仅输出一个字母选项——它还会解释原因,这正是其临床实用性的体现。
从 HuggingFace Hub 加载
微调后的适配器已公开可用。您无需克隆仓库即可直接加载:
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
tokenizer = AutoTokenizer.from_pretrained(
"Qwen/Qwen3-1.7B", trust_remote_code=True
)
base = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-1.7B",
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
)
model = PeftModel.from_pretrained(base, "HK2184/medqa-qwen3-lora")
model = model.merge_and_unload()
model.eval()
挑战与修复
没有一段"战斗故事"的 AMD ROCm 项目是不完整的。以下是我们遇到的问题:
| 挑战 | 根本原因 | 修复 |
|---|---|---|
| NaN 损失值 | 混合精度不稳定 | 从 bfloat16 切换至 fp16 |
| 未检测到 GPU | 缺少 ROCm 环境变量 | 设置 ROCR_VISIBLE_DEVICES、HIP_VISIBLE_DEVICES、HSA_OVERRIDE_GFX_VERSION |
| bitsandbytes 不受支持 | 没有适用于 ROCm 的 bitsandbytes 构建版本 | 完全放弃量化——MI300X 拥有足够的显存 |
| 推理输出异常 | Tokenizer 填充配置错误 | 设置 pad_token = eos_token 并修正 padding_side |
| 训练器评估错误 | Transformers 版本不匹配 | 固定 transformers>=4.40.0 |
bitsandbytes 的问题值得说明:在 NVIDIA 硬件上,通常需要 4 位量化才能将模型装入内存。而在配备 192 GB HBM3 的 MI300X 上,这完全没有必要。这是真正的硬件优势——训练更干净,没有量化伪影。
结果
| 指标 | 数值 |
|---|---|
| 可训练参数 | 约 220 万(占总参数的 0.15%) |
| MI300X 上的训练时间 | 约 5 分钟 |
| 使用的数据集大小 | 2000 个样本 |
| MedMCQA 基线准确率 | 约 45% |
| 框架 | PyTorch + ROCm 6.1 |
亲自尝试
没有 GPU?没问题。实时 Gradio 演示在 HuggingFace Spaces 上运行(CPU 推理):
拥有 AMD 硬件?克隆仓库并本地运行:
git clone https://github.com/HK2184/MedQA-Medical-AI-on-AMD-ROCm.git
cd MedQA-Medical-AI-on-AMD-ROCm
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
pip install transformers datasets peft accelerate trl gradio
python train.py # ~5 minutes
python infer.py # run sample questions
python app.py # launch Gradio UI
下一步计划
该项目证明了该流程是可行的。接下来的步骤是扩展和加固:
- 更大的数据集——在完整的 MedMCQA 语料库(约 18 万道题)上训练,并加入 PubMedQA
- 置信度评分——在答案旁添加校准后的置信度估计
- RAG 集成——基于实时医学文献检索来支撑答案
- 评估框架——在训练集之外进行严格的留出准确率基准测试
结论
MedQA 表明,在开源 AMD 硬件上构建一个能力强、可解释的医疗 AI 不仅是可行的——而且很简单。HuggingFace 生态系统的 ROCm 兼容性确实很好。MI300X 的内存余量消除了整整一类工程问题。而 LoRA 使得微调一个 17 亿参数的模型只需 5 分钟。
如果你正在 AMD ROCm 上构建应用并遇到瓶颈,上述修复方法应该能为你节省数小时。而如果你正在构建医疗 AI,那么强调可解释性而非单纯准确率的做法值得认真对待。
专为 lablab.ai 上的 AMD 开发者黑客松构建 · 由 AMD ROCm + HuggingFace 生态系统提供支持
*— Harikrishna Sivanand Iyer 与 Srijan Sivaram A