跳至主要内容

微调 Magistral:分步指南

使用 Transformers 框架在医疗选择题(MCQs)数据集上微调 Mistral 推理模型的分步指南。
· 10分钟 阅读

使用 AI 探索

Mistral AI 最近发布了名为 Magistral 的新一代推理模型系列。这些模型旨在在特定领域表现出色,提供透明的推理过程,并支持多种语言。

在本教程中,我们将重点在医疗选择题推理数据集上微调 Magistral Small。Magistral Small 是 Magistral 系列的开源变体,拥有 240 亿参数。它基于 3.1 版本(2503),具备增强的推理能力。

本教程将按照以下步骤进行:

  • 搭建环境 并安装所有必要的依赖。
  • 加载预训练的 Magistral Small 模型和分词器,为定制化做准备。
  • 加载医疗选择题数据集 并进行处理,确保其可用于训练。
  • 在微调前进行模型推理,以建立对比基线。
  • 配置模型和训练参数,然后启动微调过程。
  • 对比结果,评估微调前后的改进。
  • 保存微调后的模型和分词器,并上传至 Hugging Face 仓库以便分享与部署。

1. 环境搭建

我最初尝试在 Kaggle 上运行微调脚本,但即使采用 4 位量化,免费 GPU 也难以胜任。出于这些限制,我转而使用 RunPod,并选择了 A100 SXM GPU 实例,其性能更快、显存更大。

1. 在 RunPod 的“Pods”页面,选择最新的 PyTorch 镜像,选用 A100 SXM 机器并部署您的 pod。

Runpod 部署

来源:My Pods

2. 编辑 Pod 设置,增大容器磁盘大小,并添加您的 Hugging Face 访问令牌用于认证。

在 Runpod 中编辑 Pod

来源:My Pods

3. 当 pod 运行后,点击Connect 按钮,启动 JupyterLab 进行交互式开发。

4. 安装所有必要的 Python 包

%%capture
%pip install -U transformers==4.52.1
%pip install -U datasets 
%pip install -U accelerate 
%pip install -U peft 
%pip install -U trl 
%pip install -U bitsandbytes

5. 使用您的访问令牌(之前已保存)登录 Hugging Face。这样您可以加载受限模型,并在之后上传微调后的模型和分词器:

from huggingface_hub import login
import os

hf_token = os.environ.get("HF_TOKEN")
login(hf_token)

2. 加载模型与分词器

我们将使用 Magistral-Small 的 Unsloth 4 位量化版本,以节省存储空间和显存。相比下载全精度模型再手动量化,这种方式更快、更高效。 

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load tokenizer & model

model_dir = "unsloth/Magistral-Small-2506-bnb-4bit"

tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",  
    torch_dtype=torch.bfloat16,             
)

model.config.use_cache = False
model.config.pretraining_tp = 1

从 Hugging Face 下载分词器和模型文件

加载模型后,检查 GPU 内存使用情况。

!nvidia-smi

如图所示,仅占用 23GB 显存。 

GPU 统计

3. 加载与处理数据集

定义一个提示模板,指示模型从所给选项中作答,将推理写在 <analysis></analysis> 标签内,并将最终答案写在 <answer></answer> 标签内。

train_prompt_style = """
Please answer with one of the options in the bracket. Write reasoning in between <analysis></analysis>. Write the answer in between <answer></answer>.
### Question:
{}

### Response:
{}"""

创建一个 Python 函数,将每个样本按提示风格进行格式化,并在缺失时追加 EOS 标记。

EOS_TOKEN = tokenizer.eos_token  # Must add EOS_TOKEN

def formatting_prompts_func(examples):
    inputs = examples["input"]
    outputs = examples["output"]
    texts = []
    for question, response in zip(inputs, outputs):
        # Remove the "Q:" prefix from the question
        question = question.replace("Q:", "")
        
        # Append the EOS token to the response if it's not already there
        if not response.endswith(tokenizer.eos_token):
            response += tokenizer.eos_token
            
        text = train_prompt_style.format(question, response)
        texts.append(text)
    return {"text": texts}

然后,从 Hugging Face 下载并加载 mamachang/medical-reasoning 数据集。应用上述格式化函数,创建带有结构化提示的全新 "text" 列。

from datasets import load_dataset

dataset = load_dataset(
    "mamachang/medical-reasoning",
    split="train",
    trust_remote_code=True,
)
dataset = dataset.map(
    formatting_prompts_func,
    batched=True,
)
print(dataset["text"][10])

“text” 列包含系统提示、问题、分析和答案,均符合所需格式。

格式化后的数据

由于新版 SFTTrainer 不直接接收分词器,需要将分词器转换为语言建模的数据整理器(data collator)。

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

4. 微调前的模型推理

在微调之前,评估基座模型在您的数据集上的表现非常重要。该基线将帮助您在微调后进行客观对比并衡量改进幅度。

为推理定义一个提示模板,其中包含系统提示、问题部分(带占位符)和回答部分。 

inference_prompt_style = """
Please answer with one of the options in the bracket. Write reasoning in between <analysis></analysis>. Write the answer in between <answer></answer>.

### Question:
{}

### Response:
<analysis>
"""

选择数据集中的第 11 个问题,去除多余前缀,并进行分词以供模型输入

question = dataset[10]['input']
question = question.replace("Q:", "")

inputs = tokenizer(
    [inference_prompt_style.format(question) + tokenizer.eos_token],
    return_tensors="pt"
).to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(response[0].split("### Response:")[1])

在该基线测试中,模型生成的分析过长,且常常无法给出正确或简洁的答案。 

微调前的模型推理

5. 模型设置

本部分将对模型进行微调前的设置。 

我们将使用 LoRA (低秩适配) 来微调模型,仅对模型参数的一小部分进行训练,从而显著降低内存占用并加快训练过程。

from peft import LoraConfig, get_peft_model

# LoRA config
peft_config = LoraConfig(
    lora_alpha=16,                           # Scaling factor for LoRA
    lora_dropout=0.05,                       # Add slight dropout for regularization
    r=64,                                    # Rank of the LoRA update matrices
    bias="none",                             # No bias reparameterization
    task_type="CAUSAL_LM",                   # Task type: Causal Language Modeling
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],  # Target modules for LoRA
)

model = get_peft_model(model, peft_config)

通过 TrainingArguments 类可以配置批大小、优化器、学习率等训练超参数。TRL 库的 SFTTrainer 将模型、数据集、数据整理器、训练参数与 LoRA 配置整合到同一工作流中,从而简化有监督微调。

from trl import SFTTrainer
from transformers import TrainingArguments


# Training Arguments
training_arguments = TrainingArguments(
    output_dir="Magistral-Medical-Reasoning",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=2,
    optim="paged_adamw_32bit",
    num_train_epochs=1,
    logging_steps=0.2,
    warmup_steps=10,
    logging_strategy="steps",
    learning_rate=2e-4,
    fp16=False,
    bf16=False,
    group_by_length=True,
    report_to="none"
)

# Initialize the Trainer
trainer = SFTTrainer(
    model=model,
    args=training_arguments,
    train_dataset=dataset,
    peft_config=peft_config,
    data_collator=data_collator,
)

对数据集进行分词。

6. 模型训练

为避免显存不足(OOM)错误,在开始训练前清理 Python 垃圾回收器并清空 CUDA 缓存非常重要。这样可以释放之前操作中未使用的内存。

随后,使用 SFTTrainer 启动训练过程。

import gc, torch
gc.collect()
torch.cuda.empty_cache()
model.config.use_cache = False
trainer.train()

训练期间,GPU 利用率达到 100%,且显存充足,理论上还能并行加载另一个模型。

训练过程中的系统统计。

整个训练大约耗时 1 小时 10 分钟,训练损失在每一步都稳步下降,表明学习有效。

训练损失随步数变化

7. 微调后的模型推理

我们将使用数据集中的同一第 11 个样本问题来评估微调后模型的表现。这有助于与基座模型进行对比,评估通过微调所取得的提升。

question = dataset[10]['input']
question = question.replace("Q:", "")

inputs = tokenizer(
    [inference_prompt_style.format(question,) + tokenizer.eos_token],
    return_tensors="pt"
).to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(response[0].split("### Response:")[1])

模型现在能够生成简洁且准确的分析,答案正确。

<analysis>
Analysis:

This is a prospective cohort study looking at the relationship between childhood diet and cardiovascular disease in adulthood. The key issue with the validity of this study is confounding. The researchers did not account for other factors that could influence cardiovascular disease risk in adulthood besides childhood diet. 

To address confounding, the researchers should have stratified the analysis by potential confounders like family history, physical activity levels, and smoking status. This would allow them to see if the relationship between diet and disease persists even when accounting for these other variables. 

Blinding, crossover, matching, and randomization do not address the main validity concern in this study design.
</analysis>
<answer>
D: Stratification
</answer>

为进一步验证表现,再对另一条样本进行推理。

question = dataset[100]['input']
question = question.replace("Q:", "")

inputs = tokenizer(
    [inference_prompt_style.format(question) + tokenizer.eos_token],
    return_tensors="pt"
).to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(response[0].split("### Response:")[1])

模型持续给出准确答案与良好结构的分析,即便仅训练一个 epoch 也展现了稳定的改进。

<analysis>

This is a clinical vignette describing a 55-year-old man with burning and shooting pain in his feet and lower legs that worsens at night. He has a history of type 2 diabetes mellitus and hypertension. 

The key findings are:
- Burning and shooting pain in feet and lower legs 
- Pain worsens at night
- History of type 2 diabetes mellitus

This presentation is most consistent with diabetic peripheral neuropathy. The pain distribution, timing, and history of diabetes point towards a distal symmetric sensorimotor polyneuropathy as the etiology. The other options can be ruled out based on the clinical presentation.
</analysis>
<answer>
D: Distal symmetric sensorimotor polyneuropathy
</answer>

8. 保存模型 

微调完成后,您可以将模型推送到 Hugging Face Hub,方便他人使用或继续微调。

new_model_name = "kingabzpro/Magistral-Small-Medical-QA"

trainer.model.push_to_hub(new_model_name)
trainer.processing_class.push_to_hub(new_model_name)

上传 LoRA 与分词器。

微调后的模型现已发布于:kingabzpro/Magistral-Small-Medical-QA

Hugging Face Hub 上的 kingabzpro/Magistral-Small-Medical-QA 模型仓库

来源:kingabzpro/Magistral-Small-Medical-QA

9. 从 Hugging Face 加载微调模型

将 LoRA 微调结果保存到 Hugging Face Hub 之后,您可以轻松重新加载模型进行进一步分析、推理,或部署为推理服务。该流程包括清理当前会话、加载基座模型与 LoRA 适配器并合并,然后在新数据上进行推理。您可以在我们的课程中进一步了解如何使用 Hugging Face Hub。 

1. 在加载新模型前,建议删除已有的模型与训练器对象,并清空 GPU 缓存,以避免内存问题。

2. 从 Hugging Face Hub 下载基座模型与您微调的 LoRA 适配器,然后进行合并。同时加载分词器。 

del model
del trainer
torch.cuda.empty_cache()


from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import torch

# Base model
base_model_id = "unsloth/Magistral-Small-2506-bnb-4bit"

# Your fine-tuned LoRA adapter repository
lora_adapter_id = "kingabzpro/Magistral-Small-Medical-QA"

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

# Attach the LoRA adapter
model = PeftModel.from_pretrained(
    base_model,
    lora_adapter_id,
    device_map="auto",
    trust_remote_code=True,
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

3. 准备一个示例提示,进行分词、生成回答并解码输出。 

# Inference example
prompt = """
Please answer with one of the options in the bracket. Write reasoning in between <analysis></analysis>. Write the answer in between <answer></answer>.

### Question:
A research group wants to assess the relationship between childhood diet and cardiovascular disease in adulthood.
A prospective cohort study of 500 children between 10 to 15 years of age is conducted in which the participants' diets are recorded for 1 year and then the patients are assessed 20 years later for the presence of cardiovascular disease.
A statistically significant association is found between childhood consumption of vegetables and decreased risk of hyperlipidemia and exercise tolerance.
When these findings are submitted to a scientific journal, a peer reviewer comments that the researchers did not discuss the study's validity.
Which of the following additional analyses would most likely address the concerns about this study's design? 
{'A': 'Blinding', 'B': 'Crossover', 'C': 'Matching', 'D': 'Stratification', 'E': 'Randomization'},
### Response:
<analysis>

"""

inputs = tokenizer(
    [prompt + tokenizer.eos_token],
    return_tensors="pt"
).to("cuda")

outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=1200,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(response[0].split("### Response:")[1])

4. 生成的回答准确,且与数据集的风格与结构高度一致,表明该微调模型已可用于真实场景推理与部署。

<analysis>

Analysis:
This is a prospective cohort study looking at the relationship between childhood diet and cardiovascular disease in adulthood. The key concern from the peer reviewer is about the study's validity. 

To address concerns about validity, the researchers could perform additional analyses to control for confounding. Matching and stratification would help control for known confounders like socioeconomic status or family history. Crossover and blinding are not applicable to this observational study design. Randomization would not be possible since the study is observational.
</analysis>
<answer>
D: Stratification
</answer>

如果在运行上述代码时遇到问题,请参阅 Hugging Face 上的 fine-tuning-magistral.ipynb Jupyter 笔记本。

结语

在另一份微调指南中,我们考察了 DeepSeek-R1-0528 蒸馏模式的表现。尽管 DeepSeek-R1-0528 难以适应数据集的特定风格和准确性要求,Magistral 却能顺畅适配,并以极小成本产出强劲结果。 

这次实践进一步强化了我对 Mistral AI 的看法:它正迅速成为 AI 领域的领军者,不仅在语言模型方面,也覆盖包括计算机视觉在内的更广泛解决方案。

在本指南中,我们成功在医疗选择题数据集上微调了 Mistral AI 最新的推理模型 Magistral。 

但需要注意的是,微调像 Magistral 这样的大模型需要相当强劲的硬件,通常至少 40GB 显存的 GPU,一些配置下甚至 RTX 4090 也可能不足。我也尝试了 Unsloth 框架及其 4 位模型版本,但仍遇到了 GPU 内存不足的问题。

如果您在寻找上手微调的实践课程,欢迎查看 Fine-Tuning With Llama 3。

Abid Ali Awan

作者

Abid Ali Awan

作为一名持证的数据科学家,我热衷于利用前沿技术打造创新的机器学习应用。凭借在语音识别、数据分析与报告、MLOps、对话式人工智能以及自然语言处理方面的扎实背景,我不断打磨构建智能系统的能力,力求带来切实影响。除技术专长外,我也擅长沟通,能够将复杂概念提炼为清晰、简明的表述。因此,我成为数据科学领域备受关注的博主,与不断壮大的数据专业人士社区分享洞见与实践经验。目前,我专注于内容创作与编辑,借助大语言模型打造有力且吸引人的内容,帮助企业与个人更好地发挥数据价值。

主题