Finjustering av NVIDIA Nemotron-3-Nano på psykologi-data för frågor och svar

NVIDIA Nemotron-3 är NVIDIAs nya öppna modelfamilj byggd för resonerande, kodning, chatt och agentbaserade AI-arbetsflöden. Den omfattar olika modellstorlekar, som Nano, Super och Ultra, så att utvecklare kan välja mellan mindre, effektiva modeller och större, högpresterande modeller.

Den viktigaste nyheten med Nemotron-3 är fokus på effektivitet. Modellerna är utformade för att ge stark prestanda samtidigt som inferens och finjustering hålls mer praktiska. Nano-versionen är särskilt användbar för praktiska experiment eftersom den kan köras på mer tillgängliga GPU‑uppsättningar jämfört med större modeller.

I den här guiden kommer vi att finjustera NVIDIA Nemotron-3-Nano-4B på en psykologidataset för frågor och svar. Vi använder Low-Rank Adaptation (LoRA), Transformers Reinforcement Learning (TRL) och Hugging Face för att förbereda data, träna modellen, spara adaptern, pusha den till Hugging Face och jämföra svar före och efter finjustering.

För att komma igång med att hitta de senaste öppna AI-modellerna, bygga AI‑agenter och finjustera LLM:er rekommenderar jag att du anmäler dig till vår Hugging Face Fundamentals-färdighetsspår.

1. Konfigurera miljön

Nemotron-3 Nano använder en hybridarkitektur, så Mamba-relaterade paket behöver installeras korrekt. I en Jupyter‑notebook tar vi först bort den befintliga PyTorch‑stacken och installerar om CUDA 12.8‑bygget av PyTorch 2.7.1, som fungerar med de låsta versionerna av mamba_ssm och causal_conv1d som används här.

Vi installerar också kärnbiblioteken för finjustering, inklusive transformers, trl, accelerate, datasets, peft och 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"

Efter installation av paketen, kontrollera att CUDA är tillgängligt och att PyTorch kan upptäcka ditt GPU‑kort. Den här noteboken är anpassad för ett 24 GB‑GPU, så den varnar om ditt GPU har mindre 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

Utdata:

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.

Ställ in din Hugging Face‑token som en miljövariabel med namnet HF_TOKEN. Detta låter noteboken ladda ner Nemotron‑3‑modellen och senare pusha den finjusterade LoRA‑adaptern till 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. Ladda och bearbeta datasetet

Nästa steg är att ladda psykologi‑datasetet för frågor och svar från Hugging Face. Datasetet innehåller en kolumn question och två svarskolumner: response_j och response_k. I den här guiden använder vi response_j som målsvar för övervakad finjustering.

Vi laddar först datasetet, blandar det för reproducerbarhet och skapar tränings‑, validerings‑ och testsplits. 

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

Utdata:

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
    })
})

Innan vi formaterar datasetet för träning, kontrollera kolumnnamnen och titta på ett exempel. Detta bekräftar att datasetet laddades korrekt och innehåller de förväntade fälten för fråga och svar.

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

Utdata:

(
    ['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. Formatera datasetet för TRL‑finjustering

Nu konverterar vi datasetet till det prompt‑kompletteringsformat som TRL förväntar sig. Varje exempel kommer att inkludera en systemprompt, användarens psykologi‑fråga och det målsvar från response_j som assistenten ska ge.

Systemprompten talar om för modellen hur den ska svara: vara stöttande, undvika dolda resonemangsspår, ge praktiska förslag och undvika att agera som legitimerad vårdpersonal inom psykisk hälsa.

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]

Utdata:

{
   '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. Ladda basmodellen Nemotron‑3

Därefter laddar vi NVIDIA Nemotron‑3 Nano 4B BF16‑tokenizer och basmodell från Hugging Face. Vi sätter också utdatakatalogen för LoRA‑adaptern och begränsar sekvenslängden till 1024 token för att hålla träningen hanterbar på ett 24 GB‑GPU.

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. Skapa hjälpfunktioner för generering

Innan finjustering skapar vi några hjälpfunktioner för att testa modellens svar. Dessa funktioner bygger chat‑prompten, genererar ett svar, tar bort oönskade thinking‑taggar och sparar resultaten i en liten jämförelsetabell.

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. Kör en utvärdering före finjustering

Innan träning genererar vi några svar från basmodellen Nemotron‑3. Detta ger oss en baslinje så att vi senare kan jämföra hur modellen svarar före och efter LoRA‑finjustering.

Här väljer vi tre exempel från testmängden och genererar svar med hjälp av hjälpfunktionen vi skapade tidigare.

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

Utdata är en liten tabell med den ursprungliga frågan, referenssvaret från response_j och svaret som genererades av basmodellen. Denna tabell kommer att vara användbar senare när vi jämför med den finjusterade modellens svar.

7. Konfigurera LoRA och träningsinställningar

Nu förbereder vi modellen för LoRA‑finjustering. Vi aktiverar gradientcheckpointing för att minska minnesanvändningen och skapar sedan en LoRA‑konfiguration som riktar sig mot alla linjära lager i modellen.

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",
)

Därefter definierar vi inställningarna för övervakad finjustering med SFTConfig. Dessa inställningar styr batchstorlek, inlärningshastighet, antal epoker, utvärderingsfrekvens, sparstrategi och BF16‑träning.

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. Träna och spara LoRA‑adaptern

Nu kan vi skapa SFTTrainer, koppla LoRA‑konfigurationen och starta finjusteringen. Innan träning kontrollerar vi också hur många parametrar som är träningsbara för att bekräfta att LoRA‑adaptern kopplats in korrekt.

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

Under träningen bör träningsförlust och valideringsförlust gradvis minska. Det brukar innebära att modellen lär sig svarsstilen från datasetet.

Efter träningen, spara LoRA‑adaptern och tokenizern lokalt:

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

Pusha sedan den finjusterade adaptern till 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)

Den finjusterade adaptern är nu sparad lokalt och uppladdad till Hugging Face under HUB_REPO_ID.

Källa: kingabzpro/nemotron-3-nano-4b-bf16-psychology-qa-lora

9. Jämföra svar före och efter finjustering

Slutligen genererar vi svar från den finjusterade modellen och jämför dem med utdata från basmodellen. Detta hjälper oss se om LoRA‑finjusteringen förbättrade modellens anpassning till referenssvaren.

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")

Den finjusterade modellen blev mer anpassad till referenssvarens stil. Den var mer koncis och höll sig närmare datasetsvaren. Basmodellen gav dock ibland mer detaljerade och praktiska svar.

Till exempel förbättrades anpassningen för frågor om stresshantering och koncentration, men basmodellen gav ett starkare svar för exemplet om sömn eftersom den inkluderade mer hjälpsamma detaljer.

Överlag är den finjusterade modellen bättre om ditt mål är att matcha referensdatasetets stil. Om målet är maximal hjälpsamhet kan basmodellen fortfarande prestera bättre i vissa fall eftersom den kan ge varmare och mer detaljerade svar.

Om du får problem med att köra koden ovan, se noteboken i Hugging Face‑repot: fine-tune-nemotron-nano.ipynb

Avslutande tankar

Även efter att ha finjusterat 100+ LLM:er krävde den här modellen mer uppsättning än väntat. Den största utmaningen var beroendet av mamba_ssm, som lätt kan gå sönder eller krocka med en befintlig lokal Python‑miljö.

Därför rekommenderar jag att använda en ren miljö för detta arbetsflöde. I mitt fall var den enklaste vägen att bygga om miljön, installera rätt PyTorch‑version, låsa Mamba‑relaterade paket och sedan köra noteboken därifrån.

En annan begränsning är kvantisering. För den här konfigurationen kunde jag inte helt enkelt ladda modellen i 4‑bit och finjustera den som ett standard‑QLoRA‑arbetsflöde, som i min Qwen3.5 Small-guide. Jag var tvungen att ladda hela BF16‑modellen och sedan finjustera den med LoRA. För en 4B‑modell är detta fortfarande hanterbart på ett 24 GB‑GPU, men för 12B‑modeller och uppåt kan minnesanvändningen snabbt bli ett problem.

Med det sagt har finjustering på konsument‑GPU blivit mycket mer tillgänglig. Med ett 24 GB‑kort som RTX 3090 är det nu möjligt att anpassa starka öppna modeller till en viss stil eller domän utan att behöva ett stort träningskluster.

Sammantaget är Nemotron‑3 Nano en kapabel modell, men den kräver noggrann miljökonfiguration. När beroendena väl fungerar finjusteras den bra och kan anpassa sig till en ny svarsstil med ett relativt litet antal exempel.