Fine-Tuning Magistral: Hướng dẫn từng bước

Mistral AI gần đây đã ra mắt một họ mô hình lập luận mới có tên Magistral. Những mô hình này được thiết kế để vượt trội trong các lĩnh vực cụ thể, cung cấp khả năng lập luận minh bạch và hỗ trợ đa ngôn ngữ.

Trong hướng dẫn này, chúng ta sẽ tập trung fine-tune Magistral Small trên bộ dữ liệu lập luận MCQs y khoa. Magistral Small là biến thể mã nguồn mở của họ Magistral với 24 tỷ tham số. Mô hình này dựa trên phiên bản 3.1 (2503) và có năng lực lập luận được tăng cường.

Dưới đây là các bước chúng ta sẽ thực hiện trong hướng dẫn này:

• Thiết lập môi trường và cài đặt tất cả các phụ thuộc cần thiết.
• Tải mô hình và tokenizer Magistral Small đã được huấn luyện trước để chuẩn bị tùy biến.
• Tải bộ dữ liệu MCQs y khoa và xử lý để sẵn sàng cho huấn luyện.
• Thực hiện suy luận mô hình trước khi fine-tune để thiết lập đường cơ sở so sánh.
• Cấu hình mô hình và tham số huấn luyện, sau đó bắt đầu quá trình fine-tune.
• So sánh kết quả trước và sau khi fine-tune để đánh giá cải thiện.
• Lưu mô hình và tokenizer đã fine-tune, và tải chúng lên kho Hugging Face để chia sẻ và triển khai.

1. Thiết lập

Ban đầu tôi thử chạy script fine-tune trên Kaggle, nhưng ngay cả với lượng tử hóa 4-bit, GPU miễn phí vẫn không đủ cho tác vụ này. Do những hạn chế đó, tôi chuyển sang RunPod và chọn phiên bản GPU A100 SXM, vốn cho hiệu năng nhanh hơn và VRAM lớn hơn.

1. Vào “Pods” trên RunPod, chọn image PyTorch mới nhất, chọn máy A100 SXM, và triển khai pod của bạn.

Nguồn: My Pods

2. Chỉnh sửa thiết lập Pod để tăng dung lượng đĩa của container, và thêm access token Hugging Face của bạn để xác thực.

Nguồn: My Pods

3. Khi pod chạy, bấm nút Connect và khởi chạy một phiên JupyterLab để phát triển tương tác.

4. Cài đặt tất cả các gói Python cần thiết

%%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Được hỗ trợ bởi

5. Đăng nhập vào Hugging Face bằng access token của bạn (đã lưu trước đó). Điều này cho phép bạn tải các mô hình bị giới hạn và sau đó tải lên mô hình cùng tokenizer đã fine-tune:

from huggingface_hub import login
import os

hf_token = os.environ.get("HF_TOKEN")
login(hf_token)Được hỗ trợ bởi

2. Tải mô hình và tokenizer

Chúng ta sẽ dùng phiên bản Magistral-Small lượng tử hóa 4-bit của Unsloth để tiết kiệm cả dung lượng lưu trữ và VRAM. Cách tiếp cận này nhanh và hiệu quả hơn nhiều so với tải mô hình độ chính xác đầy đủ rồi tự lượng tử hóa thủ công. 

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Được hỗ trợ bởi

Sau khi tải mô hình, hãy kiểm tra mức sử dụng bộ nhớ GPU của bạn.

!nvidia-smiĐược hỗ trợ bởi

Như bạn thấy, mô hình chỉ dùng 23GB VRAM. 

3. Tải và xử lý bộ dữ liệu

Định nghĩa mẫu prompt hướng dẫn mô hình trả lời bằng một trong các lựa chọn được cung cấp, đưa phần lập luận trong thẻ <analysis></analysis> và câu trả lời cuối cùng trong thẻ <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:
{}"""
Được hỗ trợ bởi

Tạo hàm Python để định dạng mỗi ví dụ theo phong cách prompt và nối thêm token EOS nếu thiếu.

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}Được hỗ trợ bởi

Sau đó, tải và nạp bộ dữ liệu mamachang/medical-reasoning từ Hugging Face. Áp dụng hàm định dạng để tạo cột "text" mới với prompt có cấu trúc.

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])Được hỗ trợ bởi

Cột "text" chứa system prompt, câu hỏi, phần phân tích và câu trả lời, tất cả theo đúng định dạng yêu cầu.

Vì SFTTrainer mới không chấp nhận truyền trực tiếp tokenizer, hãy chuyển tokenizer thành data collator cho language modeling.

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)Được hỗ trợ bởi

4. Suy luận mô hình trước khi fine-tune

Trước khi fine-tune, điều quan trọng là đánh giá hiệu năng của mô hình gốc trên bộ dữ liệu của bạn. Đường cơ sở này sẽ giúp bạn so sánh khách quan sau khi fine-tune và đo lường cải thiện.

Định nghĩa mẫu prompt cho suy luận gồm system prompt, phần câu hỏi (với placeholder), và phần phản hồi. 

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>
"""Được hỗ trợ bởi

Chọn câu hỏi thứ 11 từ bộ dữ liệu của bạn, loại bỏ tiền tố không cần thiết, và tokenize để đưa vào mô hình

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])Được hỗ trợ bởi

Trong bài kiểm tra đường cơ sở này, phần phân tích do mô hình tạo ra quá dài và thường không đưa ra câu trả lời đúng hoặc súc tích. 

5. Thiết lập mô hình

Ở phần này, chúng ta sẽ thiết lập mô hình để fine-tune. 

Chúng ta sẽ sử dụng LoRA (Low-Rank Adaptation) để fine-tune mô hình, giúp giảm đáng kể bộ nhớ sử dụng và tăng tốc quá trình huấn luyện vì chúng ta chỉ tinh chỉnh một phần nhỏ tham số của mô hình.

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)Được hỗ trợ bởi

Lớp TrainingArguments cho phép bạn cấu hình batch size, optimizer, learning rate và các siêu tham số huấn luyện khác. SFTTrainer từ thư viện TRL đơn giản hóa supervised fine-tuning bằng cách tích hợp mô hình, bộ dữ liệu, data collator, tham số huấn luyện và cấu hình LoRA vào một quy trình thống nhất.

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,
)Được hỗ trợ bởi

6. Huấn luyện mô hình

Để tránh lỗi Out of Memory (OOM), điều quan trọng là dọn bộ gom rác của Python và làm trống bộ nhớ đệm CUDA trước khi bắt đầu huấn luyện. Điều này đảm bảo mọi bộ nhớ không dùng đến từ các thao tác trước được giải phóng.

Sau đó, bắt đầu quá trình huấn luyện bằng SFTTrainer.

import gc, torch
gc.collect()
torch.cuda.empty_cache()
model.config.use_cache = False
trainer.train()Được hỗ trợ bởi

Trong quá trình huấn luyện, mức sử dụng GPU đạt 100%, và vẫn còn đủ VRAM để có thể tải thêm một mô hình song song.

Quá trình huấn luyện mất khoảng 1 giờ 10 phút, với loss huấn luyện giảm đều theo từng bước, cho thấy việc học hiệu quả.

7. Suy luận mô hình sau khi fine-tune

Chúng ta sẽ đánh giá hiệu năng của mô hình đã fine-tune bằng cách dùng lại câu hỏi mẫu thứ 11 từ bộ dữ liệu. Cách này cho phép so sánh với mô hình gốc và đánh giá mức cải thiện đạt được sau fine-tune.

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])Được hỗ trợ bởi

Hiện mô hình tạo ra phần phân tích ngắn gọn, chính xác và câu trả lời đúng.

<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>Được hỗ trợ bởi

Để kiểm chứng thêm hiệu năng, hãy thử suy luận trên một mẫu khác.

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])Được hỗ trợ bởi

Mô hình tiếp tục đưa ra câu trả lời chính xác và phân tích có cấu trúc tốt, thể hiện cải thiện ổn định ngay cả chỉ sau một epoch fine-tune.

<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>Được hỗ trợ bởi

8. Lưu mô hình 

Sau khi fine-tune, bạn có thể dễ dàng chia sẻ mô hình bằng cách đẩy lên Hugging Face Hub. Điều này giúp người khác có thể sử dụng hoặc tiếp tục fine-tune.

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)Được hỗ trợ bởi

Mô hình đã fine-tune hiện có tại: kingabzpro/Magistral-Small-Medical-QA

Nguồn: kingabzpro/Magistral-Small-Medical-QA

9. Tải mô hình đã fine-tune từ Hugging Face

Sau khi fine-tune và lưu LoRA của bạn lên Hugging Face Hub, bạn có thể dễ dàng tải lại để phân tích thêm, suy luận, hoặc triển khai như một máy chủ suy luận. Quy trình gồm xóa phiên hiện tại, tải mô hình gốc và bộ chuyển tiếp LoRA, rồi chạy suy luận trên dữ liệu mới. Bạn có thể tìm hiểu thêm về Làm việc với Hugging Face Hub trong khóa học của chúng tôi. 

1. Trước khi tải mô hình mới, nên xóa mọi đối tượng mô hình và trainer hiện có và dọn bộ nhớ GPU để tránh lỗi bộ nhớ.

2. Tải mô hình gốc và bộ chuyển tiếp LoRA đã fine-tune của bạn từ Hugging Face Hub, sau đó hợp nhất chúng. Đồng thời tải tokenizer. 

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)Được hỗ trợ bởi

3. Chuẩn bị một prompt mẫu, tokenize, sinh phản hồi và giải mã đầu ra. 

# 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])Được hỗ trợ bởi

4. Phản hồi được sinh ra chính xác và phù hợp chặt chẽ với phong cách, cấu trúc của bộ dữ liệu của bạn, cho thấy mô hình đã fine-tune sẵn sàng cho suy luận và triển khai trong thực tế.

<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>
Được hỗ trợ bởi

Nếu bạn gặp vấn đề khi chạy đoạn mã trên, vui lòng tham khảo notebook fine-tuning-magistral.ipynb trên Hugging Face.

Kết luận

Trong một hướng dẫn fine-tune riêng, chúng tôi đã xem xét hiệu năng của chế độ distilled DeepSeek-R1-0528. Trong khi DeepSeek-R1-0528 gặp khó khi thích ứng với phong cách cụ thể của bộ dữ liệu và độ chính xác, Magistral cho thấy khả năng thích ứng mượt mà và mang lại kết quả mạnh mẽ với nỗ lực tối thiểu. 

Trải nghiệm này củng cố niềm tin của tôi rằng Mistral AI đang nhanh chóng nổi lên như một người dẫn đầu trong lĩnh vực AI, không chỉ ở mô hình ngôn ngữ mà còn trên nhiều giải pháp rộng hơn, bao gồm cả thị giác máy tính.

Trong hướng dẫn này, chúng ta đã fine-tune thành công mô hình lập luận mới nhất của Mistral AI, Magistral, trên bộ dữ liệu MCQs y khoa. 

Tuy nhiên, cần lưu ý rằng fine-tune các mô hình lớn như Magistral đòi hỏi phần cứng đáng kể, thường là GPU với ít nhất 40GB VRAM, vì ngay cả RTX 4090 cũng có thể không đủ cho một số cấu hình. Tôi cũng đã thử nghiệm framework Unsloth và phiên bản mô hình 4-bit của nó, nhưng vẫn gặp vấn đề bộ nhớ GPU.

Nếu bạn đang tìm một khóa học thực hành để làm quen với fine-tune, hãy xem Fine-Tuning With Llama 3.