Fine-tuning di NVIDIA Nemotron-3-Nano su dati di Q&A di psicologia

NVIDIA Nemotron-3 è la nuova famiglia di modelli open di NVIDIA pensata per ragionamento, coding, chat e workflow di agentic AI. Include diverse dimensioni di modello, come Nano, Super e Ultra, così gli sviluppatori possono scegliere tra modelli più piccoli ed efficienti e modelli più grandi e ad alte prestazioni.

L'aggiornamento chiave di Nemotron-3 è l'attenzione all'efficienza. I modelli sono progettati per offrire prestazioni solide mantenendo più pratici inference e fine-tuning. La versione Nano è particolarmente utile per la sperimentazione pratica perché può essere eseguita su configurazioni GPU più accessibili rispetto ai modelli più grandi.

In questa guida, faremo il fine-tuning di NVIDIA Nemotron-3-Nano-4B su un dataset di domande e risposte di psicologia. Useremo Low-Rank Adaptation (LoRA), Transformers Reinforcement Learning (TRL) e Hugging Face per preparare i dati, addestrare il modello, salvare l'adapter, caricarlo su Hugging Face e confrontare le risposte prima e dopo il fine-tuning.

Per iniziare a trovare i modelli AI open-source più recenti, creare agenti AI e fare fine-tuning di LLM, ti consiglio di iscriverti al nostro skill track Hugging Face Fundamentals.

1. Configurazione dell'ambiente

Nemotron-3 Nano utilizza un'architettura ibrida, quindi i pacchetti correlati a Mamba devono essere installati correttamente. In un notebook Jupyter, per prima cosa rimuoviamo lo stack PyTorch esistente e reinstalliamo la build CUDA 12.8 di PyTorch 2.7.1, che funziona con le versioni bloccate di mamba_ssm e causal_conv1d usate qui.

Installiamo anche le librerie principali per il fine-tuning, tra cui transformers, trl, accelerate, datasets, peft e huggingface_hub.

%%capture
!pip install -U packaging ninja

# Replace the current PyTorch stack with the CUDA 12.8 build that works with these Mamba kernel pins.
!pip uninstall -y torch torchvision torchaudio triton

!pip install "torch==2.7.1" "torchvision==0.22.1" "torchaudio==2.7.1" --index-url https://download.pytorch.org/whl/cu128

!pip install -U "transformers==4.56.2" tokenizers "trl==0.22.2" accelerate datasets peft pandas tqdm huggingface_hub safetensors

!pip install -U --no-build-isolation "mamba_ssm==2.2.5" "causal_conv1d==1.5.2"

Dopo l'installazione dei pacchetti, verifica che CUDA sia disponibile e che PyTorch rilevi la tua GPU. Questo notebook è ottimizzato per una GPU da 24 GB, quindi ti avviserà se la tua GPU ha meno VRAM.

import os
import platform
import torch

print(f"Python: {platform.python_version()}")
print(f"PyTorch: {torch.__version__}")
print(f"PyTorch CUDA build: {torch.version.cuda}")
print(f"CUDA available: {torch.cuda.is_available()}")

if not torch.cuda.is_available():
   raise RuntimeError(
       "CUDA is not available. Select a RunPod PyTorch image with GPU support."
   )

for idx in range(torch.cuda.device_count()):
   props = torch.cuda.get_device_properties(idx)
   total_gb = props.total_memory / 1024**3
   print(
       f"GPU {idx}: {props.name} ({total_gb:.1f} GB VRAM, capability {props.major}.{props.minor})"
   )

if torch.cuda.get_device_properties(0).total_memory < 24 * 1024**3:
   print(
       "Warning: this 4B LoRA notebook is tuned for GPUs with at least 24GB VRAM. Reduce batch sizes on smaller GPUs."
   )

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

Output:

Python: 3.12.3
PyTorch: 2.7.1+cu128
PyTorch CUDA build: 12.8
CUDA available: True
GPU 0: NVIDIA GeForce RTX 3090 (23.6 GB VRAM, capability 8.6)
Warning: this 4B LoRA notebook is tuned for GPUs with at least 24GB VRAM. Reduce batch sizes on smaller GPUs.

Imposta il tuo token di Hugging Face come variabile d'ambiente chiamata HF_TOKEN. Questo permette al notebook di scaricare il modello Nemotron-3 e, più avanti, di caricare l'adapter LoRA fine-tunato su Hugging Face.

from huggingface_hub import login

hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
   raise ValueError(
       "Set HF_TOKEN in the RunPod environment before running this notebook."
   )

login(token=hf_token)
print("Logged in to Hugging Face.")

2. Caricamento e preparazione del dataset

Ora caricheremo il dataset di domande e risposte di psicologia da Hugging Face. Il dataset contiene una colonna question e due colonne di risposta: response_j e response_k. In questa guida useremo response_j come risposta target per il fine-tuning supervisionato.

Per prima cosa carichiamo il dataset, lo mescoliamo per riproducibilità e creiamo split di train, validation e test. 

from datasets import DatasetDict, load_dataset

DATASET_ID = "jkhedri/psychology-dataset"
TRAIN_LIMIT = 8000
VALIDATION_LIMIT = 800
TEST_LIMIT = 300
SEED = 42

raw_dataset = load_dataset(DATASET_ID)
raw_train = raw_dataset["train"].shuffle(seed=SEED)

split_1 = raw_train.train_test_split(test_size=0.15, seed=SEED)
split_2 = split_1["test"].train_test_split(test_size=0.33, seed=SEED)


def maybe_limit(split, limit):
    if limit is None:
        return split
    return split.select(range(min(limit, len(split))))


dataset = DatasetDict(
    {
        "train": maybe_limit(split_1["train"], TRAIN_LIMIT),
        "validation": maybe_limit(split_2["train"], VALIDATION_LIMIT),
        "test": maybe_limit(split_2["test"], TEST_LIMIT),
    }
)

dataset

Output:

DatasetDict({
    train: Dataset({
        features: ['question', 'response_j', 'response_k'],
        num_rows: 8000
    })
    validation: Dataset({
        features: ['question', 'response_j', 'response_k'],
        num_rows: 800
    })
    test: Dataset({
        features: ['question', 'response_j', 'response_k'],
        num_rows: 300
    })
})

Prima di formattare il dataset per l'addestramento, controlla i nomi delle colonne e osserva un esempio. Questo conferma che il dataset è stato caricato correttamente e contiene i campi di domanda e risposta previsti.

dataset["train"].column_names, dataset["train"][0]

Output:

(
    ['question', 'response_j', 'response_k'],
    {
        'question': "I'm experiencing anxiety about social situations and don't know how to cope.",
        'response_j': "Social anxiety can be a difficult and isolating experience, but there are effective treatments available. Let's work on developing coping mechanisms, such as deep breathing and mindfulness, and exposure therapy to gradually confront your fears. We can also explore ways to improve social skills and build self-confidence.",
        'response_k': "Just avoid social situations. It's not worth the anxiety and discomfort. You can also try using alcohol or drugs to help you feel more comfortable in social settings."
    }
)

3. Formattazione del dataset per il fine-tuning con TRL

Ora convertiremo il dataset nel formato prompt-completion previsto da TRL. Ogni esempio includerà un prompt di sistema, la domanda di psicologia dell'utente e la risposta target dell'assistente da response_j.

Il prompt di sistema indica al modello come rispondere: essere di supporto, evitare tracce di ragionamento nascoste, fornire suggerimenti pratici ed evitare di presentarsi come un professionista della salute mentale abilitato.

SYSTEM_PROMPT = """/no_think
You are a supportive psychology question-answering assistant.
Do not include hidden reasoning, thinking traces, <think> tags, or </think> tags in the final answer.
Respond with empathy, practical coping suggestions, and clear next steps.
Give a complete answer in 2-4 short paragraphs or a brief paragraph plus 3-5 practical bullets.
Do not diagnose the user or claim to replace a licensed mental health professional.
If the user may be in immediate danger or crisis, encourage contacting local emergency services or a trusted crisis hotline.
Keep the answer safe, specific, and directly relevant to the user's question without being overly brief."""

CHAT_TEMPLATE_KWARGS = {"enable_thinking": False}
USER_TEMPLATE = "Question:\n\n{question}"


def clean_text(value):
   return " ".join(str(value).strip().split())


def to_prompt_completion(example):
   question = clean_text(example["question"])
   answer = clean_text(example["response_j"])

   return {
       "prompt": [
           {"role": "system", "content": SYSTEM_PROMPT},
           {"role": "user", "content": USER_TEMPLATE.format(question=question)},
       ],
       "completion": [
           {"role": "assistant", "content": answer},
       ],
       "chat_template_kwargs": CHAT_TEMPLATE_KWARGS,
   }


sft_dataset = dataset.map(
   to_prompt_completion, remove_columns=dataset["train"].column_names
)

sft_dataset["train"][0]

Output:

{
   'prompt': [
       {
           'role': 'system',
           'content': "/no_think\nYou are a supportive psychology question-answering assistant.\nDo not include hidden reasoning, thinking traces, <think> tags, or </think> tags in the final answer.\nRespond with empathy, practical coping suggestions, and clear next steps.\nGive a complete answer in 2-4 short paragraphs or a brief paragraph plus 3-5 practical bullets.\nDo not diagnose the user or claim to replace a licensed mental health professional.\nIf the user may be in immediate danger or crisis, encourage contacting local emergency services or a trusted crisis hotline.\nKeep the answer safe, specific, and directly relevant to the user's question without being overly brief."
       },
       {
           'role': 'user',
           'content': "Question:\n\nI'm experiencing anxiety about social situations and don't know how to cope."
       }
   ],
   'completion': [
       {
           'role': 'assistant',
           'content': "Social anxiety can be a difficult and isolating experience, but there are effective treatments available. Let's work on developing coping mechanisms, such as deep breathing and mindfulness, and exposure therapy to gradually confront your fears. We can also explore ways to improve social skills and build self-confidence."
       }
   ],
   'chat_template_kwargs': {'enable_thinking': False}
}

4. Caricamento del modello base Nemotron-3

Successivamente, caricheremo il tokenizer e il modello base NVIDIA Nemotron-3 Nano 4B BF16 da Hugging Face. Impostiamo anche la directory di output per l'adapter LoRA e limitiamo la lunghezza di sequenza a 1024 token per mantenere l'addestramento gestibile su una GPU da 24 GB.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
OUTPUT_DIR = "./nemotron-3-nano-4b-bf16-psychology-qa-lora"
MAX_SEQ_LENGTH = 1024

tokenizer = AutoTokenizer.from_pretrained(
   MODEL_ID,
   token=hf_token,
   trust_remote_code=True,
   use_fast=True,
)

if tokenizer.pad_token is None:
   tokenizer.pad_token = tokenizer.eos_token

tokenizer.padding_side = "right"

base_model = AutoModelForCausalLM.from_pretrained(
   MODEL_ID,
   token=hf_token,
   trust_remote_code=True,
   dtype=torch.bfloat16,
   device_map="auto",
   attn_implementation="eager",
)

base_model.config.use_cache = False
base_model.config.pad_token_id = tokenizer.pad_token_id
base_model.config.eos_token_id = tokenizer.eos_token_id
base_model.generation_config.pad_token_id = tokenizer.pad_token_id
base_model.generation_config.eos_token_id = tokenizer.eos_token_id
base_model.generation_config.use_cache = False
base_model.generation_config.do_sample = False
base_model.generation_config.top_p = None
base_model.generation_config.min_new_tokens = None
base_model.generation_config.repetition_penalty = 1.08
base_model.generation_config.no_repeat_ngram_size = 4

5. Creazione di funzioni di supporto per la generazione

Prima del fine-tuning, creeremo alcune funzioni di supporto per testare le risposte del modello. Queste funzioni costruiscono il prompt di chat, generano una risposta, rimuovono eventuali tag di thinking indesiderati e salvano i risultati in una piccola tabella di confronto.

import gc
import pandas as pd
from tqdm.auto import tqdm


def clear_cuda_cache():
   gc.collect()
   if torch.cuda.is_available():
       torch.cuda.empty_cache()


def build_messages(question, system_prompt=SYSTEM_PROMPT):
   return [
       {"role": "system", "content": system_prompt},
       {
           "role": "user",
           "content": USER_TEMPLATE.format(question=clean_text(question)),
       },
   ]


def remove_thinking_text(text):
   text = text.strip()
   while "<think>" in text and "</think>" in text:
       start = text.find("<think>")
       end = text.find("</think>", start) + len("</think>")
       text = (text[:start] + text[end:]).strip()

   if "</think>" in text:
       text = text.split("</think>")[-1].strip()

   return text.replace("<think>", "").replace("</think>", "").strip()


def generate_answer(
   model, tokenizer, question, system_prompt=SYSTEM_PROMPT, max_new_tokens=180
):
   messages = build_messages(question, system_prompt)
   device = next(model.parameters()).device

   inputs = tokenizer.apply_chat_template(
       messages,
       tokenize=True,
       **CHAT_TEMPLATE_KWARGS,
       add_generation_prompt=True,
       return_dict=True,
       return_tensors="pt",
   )

   inputs = {key: value.to(device) for key, value in inputs.items()}
   input_len = inputs["input_ids"].shape[-1]

   with torch.no_grad():
       outputs = model.generate(
           **inputs,
           max_new_tokens=max_new_tokens,
           do_sample=False,
           use_cache=False,
           repetition_penalty=1.08,
           no_repeat_ngram_size=4,
           pad_token_id=tokenizer.pad_token_id,
           eos_token_id=tokenizer.eos_token_id,
       )

   decoded = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()

   return remove_thinking_text(decoded)


def generate_sample_table(model, tokenizer, examples, output_column):
   rows = []
   model.eval()

   for ex in tqdm(examples, desc=f"Generating {output_column}", leave=False):
       rows.append(
           {
               "question": clean_text(ex["question"]),
               "reference_response_j": clean_text(ex["response_j"]),
               output_column: generate_answer(model, tokenizer, ex["question"]),
           }
       )

   return pd.DataFrame(rows)

6. Esecuzione di una valutazione di esempio pre-fine-tuning

Prima dell'addestramento, genereremo alcune risposte dal modello base Nemotron-3. Questo ci dà una baseline per confrontare in seguito come risponde il modello prima e dopo il fine-tuning con LoRA.

Qui selezioniamo tre esempi dal set di test e generiamo le risposte usando la funzione di supporto creata in precedenza.

sample_examples = [dataset["test"][idx] for idx in range(min(3, len(dataset["test"])))]

pre_samples = generate_sample_table(
   base_model,
   tokenizer,
   sample_examples,
   "base_model_answer"
)

pre_samples

L'output è una piccola tabella con la domanda originale, la risposta di riferimento da response_j e la risposta generata dal modello base. Questa tabella sarà utile più avanti quando la confronteremo con le risposte del modello fine-tunato.

7. Configurazione di LoRA e impostazioni di training

Ora prepareremo il modello per il fine-tuning con LoRA. Abilitiamo il gradient checkpointing per ridurre l'uso di memoria, poi creiamo una configurazione LoRA che mira a tutti i livelli lineari del modello.

from peft import LoraConfig

base_model.gradient_checkpointing_enable()
base_model.config.use_cache = False

lora_config = LoraConfig(
   r=32,
   lora_alpha=64,
   lora_dropout=0.1,
   bias="none",
   task_type="CAUSAL_LM",
   target_modules="all-linear",
)

Poi definiamo le impostazioni di fine-tuning supervisionato usando SFTConfig. Queste impostazioni controllano batch size, learning rate, numero di epoche, frequenza di valutazione, strategia di salvataggio e training in BF16.

from trl import SFTConfig, SFTTrainer

training_args = SFTConfig(
   output_dir=OUTPUT_DIR,
   per_device_train_batch_size=8,
   per_device_eval_batch_size=8,
   gradient_accumulation_steps=8,
   learning_rate=5e-5,
   weight_decay=0.01,
   lr_scheduler_type="linear",
   warmup_ratio=0.05,
   num_train_epochs=2,
   logging_steps=50,
   eval_strategy="steps",
   eval_steps=50,
   save_strategy="steps",
   save_steps=100,
   save_total_limit=2,
   load_best_model_at_end=True,
   metric_for_best_model="eval_loss",
   greater_is_better=False,
   gradient_checkpointing=True,
   bf16=True,
   fp16=False,
   tf32=True,
   max_length=MAX_SEQ_LENGTH,
   packing=False,
   completion_only_loss=True,
   remove_unused_columns=False,
   dataloader_num_workers=4,
   optim="adamw_torch_fused",
   report_to="none",
   seed=SEED,
)

8. Training e salvataggio dell'adapter LoRA

Ora possiamo creare l'SFTTrainer, collegare la configurazione LoRA e iniziare il fine-tuning. Prima dell'addestramento, verifichiamo anche quanti parametri sono addestrabili per confermare che l'adapter LoRA sia stato collegato correttamente.

trainer = SFTTrainer(
   model=base_model,
   args=training_args,
   train_dataset=sft_dataset["train"],
   eval_dataset=sft_dataset["validation"],
   peft_config=lora_config,
   processing_class=tokenizer,
)

trainable_params = sum(
   param.numel() for param in trainer.model.parameters() if param.requires_grad
)

all_params = sum(param.numel() for param in trainer.model.parameters())

if trainable_params == 0:
   raise RuntimeError(
       "No trainable LoRA parameters were attached. Check target_modules before training."
   )

print(f"Trainable LoRA parameters: {trainable_params:,}")
print(f"All parameters visible to trainer: {all_params:,}")
print(f"Trainable percentage: {100 * trainable_params / all_params:.4f}%")

train_result = trainer.train()

trainer.model.eval()
trainer.model.config.use_cache = False
trainer.model.generation_config.use_cache = False

train_result

Durante l'addestramento, il training loss e il validation loss dovrebbero ridursi gradualmente. Questo di solito significa che il modello sta imparando lo stile di risposta dal dataset.

Dopo l'addestramento, salva localmente l'adapter LoRA e il tokenizer:

trainer.model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

Quindi carica l'adapter fine-tunato su Hugging Face:

HUB_REPO_ID = "kingabzpro/nemotron-3-nano-4b-bf16-psychology-qa-lora"

trainer.model.push_to_hub(HUB_REPO_ID, private=False)
tokenizer.push_to_hub(HUB_REPO_ID, private=False)

L'adapter fine-tunato è ora salvato localmente e caricato su Hugging Face sotto HUB_REPO_ID.

Fonte: kingabzpro/nemotron-3-nano-4b-bf16-psychology-qa-lora

9. Confronto delle risposte prima e dopo il fine-tuning

Infine, genereremo risposte dal modello fine-tunato e le confronteremo con gli output del modello base. Questo ci aiuta a vedere se il fine-tuning con LoRA ha migliorato l'allineamento del modello con le risposte di riferimento.

post_samples = generate_sample_table(
   trainer.model,
   tokenizer,
   sample_examples,
   "fine_tuned_answer"
)

comparison = pre_samples[
   ["question", "reference_response_j", "base_model_answer"]
].merge(
   post_samples[["question", "fine_tuned_answer"]],
   on="question",
   how="left",
)

for idx, row in comparison.iterrows():
   print("=" * 100)
   print(f"Sample {idx + 1}")
   print("=" * 100)
   print("\nQUESTION:\n")
   print(row["question"])
   print("\nREFERENCE RESPONSE_J:\n")
   print(row["reference_response_j"])
   print("\nBASE MODEL ANSWER:\n")
   print(row["base_model_answer"])
   print("\nFINE-TUNED ANSWER:\n")
   print(row["fine_tuned_answer"])
   print("\n")

Il modello fine-tunato è diventato più allineato allo stile di risposta di riferimento. È stato più conciso ed è rimasto più vicino alle risposte del dataset. Tuttavia, il modello base a volte ha fornito risposte più dettagliate e pratiche.

Ad esempio, il modello fine-tunato ha migliorato l'allineamento su domande legate alla gestione dello stress e alla concentrazione, ma il modello base ha dato una risposta migliore per l'esempio sul sonno perché includeva dettagli più utili.

Nel complesso, il modello fine-tunato è migliore se il tuo obiettivo è aderire allo stile del dataset di riferimento. Se invece il tuo obiettivo è la massima utilità, il modello base può ancora rendere meglio in alcuni casi perché può fornire risposte più calde e dettagliate.

Se incontri problemi nell'esecuzione del codice sopra, fai riferimento al notebook nel repo Hugging Face: fine-tune-nemotron-nano.ipynb

Considerazioni finali

Anche dopo il fine-tuning di oltre 100 LLM, questo modello ha richiesto più lavoro di setup del previsto. La principale sfida è stata la dipendenza da mamba_ssm, che può facilmente rompersi o entrare in conflitto con un ambiente Python locale esistente.

Per questo, consiglio di usare un ambiente pulito per questo workflow. Nel mio caso, il percorso più semplice è stato ricostruire l'ambiente, installare la versione corretta di PyTorch, fissare le versioni dei pacchetti correlati a Mamba e poi eseguire il notebook da lì.

Un'altra limitazione è la quantizzazione. In questa configurazione, non sono riuscito a caricare semplicemente il modello in 4 bit e fare fine-tuning come in un normale workflow QLoRA, come nel mio tutorial su Qwen3.5 Small. Ho dovuto caricare l'intero modello BF16 e poi fare il fine-tuning con LoRA. Per un modello da 4B è ancora gestibile su una GPU da 24 GB, ma per modelli da 12B in su, l'uso di memoria può diventare rapidamente un problema.

Detto questo, il fine-tuning su GPU consumer è diventato molto più accessibile. Con una scheda da 24 GB come la RTX 3090, è ora possibile adattare modelli open robusti a uno stile o dominio specifico senza bisogno di un grande cluster di training.

Nel complesso, Nemotron-3 Nano è un modello capace, ma richiede un'attenta configurazione dell'ambiente. Una volta che le dipendenze funzionano, si fine-tuna bene e può adattarsi a un nuovo stile di risposta con un numero relativamente piccolo di esempi.