Fine-tuning di Magistral: guida passo passo

Mistral AI ha recentemente lanciato una nuova famiglia di modelli di ragionamento chiamata Magistral. Questi modelli sono progettati per eccellere in domini specifici, fornire un ragionamento trasparente e supportare più lingue.

In questo tutorial ci concentreremo sul fine-tuning di Magistral Small sul dataset di ragionamento Medical MCQs. Magistral Small è una variante open-source della famiglia Magistral con 24 miliardi di parametri. Si basa sulla versione 3.1 (2503) e presenta capacità di ragionamento migliorate.

Ecco i passaggi che seguiremo in questo tutorial:

• Configura l'ambiente e installa tutte le dipendenze necessarie.
• Carica il modello e il tokenizer preaddestrati di Magistral Small per prepararli alla personalizzazione.
• Carica il dataset di MCQ medici e processalo per assicurarti che sia pronto per l'addestramento.
• Esegui l'inferenza del modello prima del fine-tuning per stabilire un baseline di confronto.
• Configura il modello e i parametri di training, poi avvia il processo di fine-tuning.
• Confronta i risultati prima e dopo il fine-tuning per valutare i miglioramenti.
• Salva il modello e il tokenizer ottimizzati e caricali su un repository Hugging Face per la condivisione e il deployment.

1. Configurazione

Ho inizialmente provato a eseguire lo script di fine-tuning su Kaggle, ma anche con la quantizzazione a 4 bit le GPU gratuite non erano sufficienti per questo compito. A causa di queste limitazioni, sono passato a RunPod e ho selezionato un'istanza GPU A100 SXM, che offre prestazioni più veloci e una VRAM maggiore.

1. Vai su “Pods” in RunPod, scegli l'ultima immagine PyTorch, seleziona la macchina A100 SXM e distribuisci il tuo pod.

Fonte: My Pods

2. Modifica le impostazioni del Pod per aumentare la dimensione del disco del container e aggiungi il tuo token di accesso Hugging Face per l'autenticazione.

Fonte: My Pods

3. Una volta che il pod è in esecuzione, clicca sul pulsante Connect e avvia un'istanza JupyterLab per lo sviluppo interattivo.

4. Installa tutti i pacchetti Python necessari

%%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 bitsandbytesOfferto da

5. Accedi a Hugging Face usando il tuo token di accesso (che hai salvato in precedenza). Questo ti permette di caricare modelli con accesso limitato e in seguito caricare il tuo modello e tokenizer ottimizzati:

from huggingface_hub import login
import os

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

2. Caricamento del modello e del tokenizer

Useremo la versione quantizzata a 4 bit di Magistral-Small di Unsloth per risparmiare spazio di archiviazione e VRAM. Questo approccio è molto più veloce ed efficiente rispetto a scaricare il modello in piena precisione e quantizzarlo manualmente. 

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 = 1Offerto da

Dopo il caricamento del modello, controlla l'utilizzo della memoria GPU.

!nvidia-smiOfferto da

Come possiamo vedere, usa solo 23GB di VRAM. 

3. Caricamento e processamento del dataset

Definisci un template di prompt che istruisca il modello a rispondere con una delle opzioni fornite, includa il ragionamento tra i tag <analysis></analysis> e la risposta finale tra i tag <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:
{}"""
Offerto da

Crea una funzione Python per formattare ciascun esempio secondo lo stile del prompt e aggiungi il token EOS se mancante.

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}Offerto da

Poi scarica e carica il dataset mamachang/medical-reasoning da Hugging Face. Applica la funzione di formattazione per creare una nuova colonna "text" con il prompt strutturato.

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])Offerto da

La colonna "text" contiene il system prompt, la domanda, l'analisi e la risposta, tutti nel formato richiesto.

Poiché il nuovo SFTTrainer non accetta direttamente un tokenizer, converti il tokenizer in un data collator per il language modeling.

from transformers import DataCollatorForLanguageModeling

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

4. Inferenza del modello prima del fine-tuning

Prima del fine-tuning, è importante valutare le prestazioni del modello base sul tuo dataset. Questo baseline ti aiuterà a confrontare oggettivamente i risultati dopo il fine-tuning e a misurare i miglioramenti.

Definisci un template di prompt per l'inferenza che includa un system prompt, una sezione domanda (con un placeholder) e una sezione risposta. 

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>
"""Offerto da

Seleziona l'undicesima domanda dal tuo dataset, rimuovi eventuali prefissi non necessari e tokenizzala per l'input del modello

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])Offerto da

In questo test di baseline, l'analisi generata dal modello era eccessivamente lunga e spesso non riusciva a produrre una risposta corretta o concisa. 

5. Configurazione del modello

In questa parte imposteremo il modello per il fine-tuning. 

Useremo LoRA (Low-Rank Adaptation) per fare il fine-tuning del modello, riducendo sensibilmente l'impronta in memoria e velocizzando il training, dato che ottimizziamo solo un piccolo sottoinsieme dei parametri del modello.

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)Offerto da

La classe TrainingArguments ti consente di configurare batch size, ottimizzatore, learning rate e altri iperparametri di training. Il SFTTrainer della libreria TRL semplifica il supervised fine-tuning integrando modello, dataset, data collator, argomenti di training e configurazione LoRA in un unico flusso.

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,
)Offerto da

6. Addestramento del modello

Per evitare errori di Out of Memory (OOM), è importante svuotare il garbage collector di Python e la cache CUDA prima di iniziare l'addestramento. Questo assicura che qualsiasi memoria non utilizzata da operazioni precedenti venga rilasciata.

Poi, avvia il processo di training usando il SFTTrainer.

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

Durante il training, l'utilizzo della GPU ha raggiunto il 100% e c'era abbastanza VRAM disponibile per potenzialmente caricare in parallelo un altro modello.

Il processo di training ha richiesto circa 1 ora e 10 minuti, con la training loss che diminuiva costantemente a ogni step, segno di un apprendimento efficace.

7. Inferenza del modello dopo il fine-tuning

Valuteremo le prestazioni del modello ottimizzato usando la stessa undicesima domanda del dataset. Questo approccio ci consente di confrontare i risultati con il modello base e valutare i miglioramenti ottenuti con il fine-tuning.

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])Offerto da

Ora il modello genera un'analisi concisa e accurata e la risposta è corretta.

<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>Offerto da

Per convalidare ulteriormente le prestazioni, prova l'inferenza su un altro esempio.

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])Offerto da

Il modello continua a fornire risposte accurate e analisi ben strutturate, dimostrando miglioramenti consistenti anche dopo un solo epoch di fine-tuning.

<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>Offerto da

8. Salvataggio del modello 

Dopo il fine-tuning, puoi condividere facilmente il tuo modello caricandolo su Hugging Face Hub. In questo modo sarà accessibile ad altri per l'uso o per ulteriore fine-tuning.

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)Offerto da

Il modello ottimizzato è ora disponibile qui: kingabzpro/Magistral-Small-Medical-QA

Fonte: kingabzpro/Magistral-Small-Medical-QA

9. Caricamento del modello ottimizzato da Hugging Face

Dopo il fine-tuning e il salvataggio della tua LoRA su Hugging Face Hub, puoi ricaricarla facilmente per ulteriori analisi, inferenza o deployment come server di inferenza. Questo processo prevede la pulizia della sessione corrente, il caricamento del modello base e dell'adapter LoRA, e l'esecuzione dell'inferenza su nuovi dati. Puoi saperne di più su Working with the Hugging Face Hub nel nostro corso. 

1. Prima di caricare un nuovo modello, è buona pratica eliminare eventuali oggetti model e trainer esistenti e svuotare la cache della GPU per evitare problemi di memoria.

2. Scarica il modello base e il tuo adapter LoRA ottimizzato da Hugging Face Hub, quindi uniscili. Carica anche il 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)Offerto da

3. Prepara un prompt di esempio, tokenizzalo, genera una risposta e decodifica l'output. 

# 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])Offerto da

4. La risposta generata è accurata e si avvicina molto allo stile e alla struttura del tuo dataset, dimostrando che il modello ottimizzato è pronto per l'inferenza e il deployment nel mondo reale.

<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>
Offerto da

Se riscontri problemi nell'esecuzione del codice sopra, fai riferimento al Jupyter notebook fine-tuning-magistral.ipynb su Hugging Face.

Conclusione

In una guida di fine-tuning separata, abbiamo visto come si comporta la modalità distillata DeepSeek-R1-0528. Mentre DeepSeek-R1-0528 ha faticato ad adattarsi allo stile specifico del dataset e all'accuratezza, Magistral ha mostrato un adattamento fluido e ha fornito ottimi risultati con il minimo sforzo. 

Questa esperienza rafforza la mia convinzione che Mistral AI stia emergendo rapidamente come un attore di primo piano nello spazio dell'IA, non solo nei modelli linguistici ma anche in un ampio spettro di soluzioni, tra cui la computer vision.

In questa guida, abbiamo effettuato con successo il fine-tuning dell'ultimo modello di ragionamento di Mistral AI, Magistral, su un dataset di MCQ medici. 

Tuttavia, è importante notare che il fine-tuning di modelli di grandi dimensioni come Magistral richiede hardware sostanziale, tipicamente GPU con almeno 40GB di VRAM, poiché anche la RTX 4090 potrebbe non essere sufficiente per alcune configurazioni. Ho anche sperimentato il framework Unsloth e la sua versione del modello a 4 bit, ma ho comunque riscontrato problemi di memoria GPU.

Se cerchi un corso pratico per prendere confidenza con il fine-tuning, dai un'occhiata a Fine-Tuning With Llama 3.