Fine-Tuning NVIDIA Nemotron-3-Nano pada Data Tanya Jawab Psikologi

NVIDIA Nemotron-3 adalah keluarga model terbuka baru dari NVIDIA yang dibuat untuk penalaran, coding, chat, dan alur kerja AI agen. Keluarga ini mencakup berbagai ukuran model, seperti Nano, Super, dan Ultra, sehingga pengembang dapat memilih antara model yang lebih kecil dan efisien atau model yang lebih besar dan berkinerja tinggi.

Pembaruan utama pada Nemotron-3 adalah fokusnya pada efisiensi. Model dirancang untuk memberikan performa kuat sambil menjaga inferensi dan fine-tuning tetap praktis. Versi Nano sangat berguna untuk eksperimen langsung karena dapat berjalan pada setelan GPU yang lebih mudah diakses dibandingkan model yang lebih besar.

Dalam panduan ini, kita akan melakukan fine-tune NVIDIA Nemotron-3-Nano-4B pada dataset tanya jawab psikologi. Kita akan menggunakan Low-Rank Adaptation (LoRA), Transformers Reinforcement Learning (TRL), dan Hugging Face untuk menyiapkan data, melatih model, menyimpan adapter, mendorongnya ke Hugging Face, serta membandingkan respons sebelum dan sesudah fine-tuning.

Untuk memulai mencari model AI open-source terbaru, membangun agen AI, dan melakukan fine-tuning LLM, saya menyarankan mendaftar ke track keterampilan Hugging Face Fundamentals kami.

1. Menyiapkan Lingkungan

Nemotron-3 Nano menggunakan arsitektur hibrida, sehingga paket terkait Mamba perlu diinstal dengan benar. Dalam notebook Jupyter, pertama-tama kita menghapus tumpukan PyTorch yang ada dan memasang ulang build CUDA 12.8 dari PyTorch 2.7.1, yang bekerja dengan versi mamba_ssm dan causal_conv1d yang dipin di sini.

Kita juga memasang pustaka inti untuk fine-tuning, termasuk transformers, trl, accelerate, datasets, peft, dan 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"

Setelah memasang paket, periksa bahwa CUDA tersedia dan PyTorch dapat mendeteksi GPU Anda. Notebook ini disetel untuk GPU 24GB, jadi akan memberi peringatan jika GPU Anda memiliki VRAM lebih sedikit.

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

Keluaran:

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.

Tetapkan token Hugging Face Anda sebagai variabel lingkungan bernama HF_TOKEN. Ini memungkinkan notebook mengunduh model Nemotron-3 dan nantinya mendorong adapter LoRA yang telah di-fine-tune ke 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. Memuat dan Memroses Dataset

Selanjutnya, kita akan memuat dataset tanya jawab psikologi dari Hugging Face. Dataset ini berisi kolom question dan dua kolom respons: response_j dan response_k. Untuk panduan ini, kita akan menggunakan response_j sebagai jawaban target untuk supervised fine-tuning.

Kita terlebih dahulu memuat dataset, mengacaknya untuk reprodusibilitas, dan membuat split train, validasi, serta uji. 

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

Keluaran:

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

Sebelum memformat dataset untuk pelatihan, periksa nama kolom dan lihat satu contoh. Ini memastikan bahwa dataset dimuat dengan benar dan berisi field pertanyaan serta respons yang diharapkan.

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

Keluaran:

(
    ['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. Memformat Dataset untuk TRL Fine-Tuning

Sekarang kita akan mengonversi dataset ke format prompt-completion yang diharapkan oleh TRL. Setiap contoh akan menyertakan prompt sistem, pertanyaan psikologi dari pengguna, dan respons asisten target dari response_j.

Prompt sistem memberi tahu model bagaimana harus merespons: bersikap suportif, hindari jejak penalaran tersembunyi, berikan saran praktis, dan hindari bertindak sebagai profesional kesehatan mental berlisensi.

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]

Keluaran:

{
   '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. Memuat Model Dasar Nemotron-3

Selanjutnya, kita akan memuat tokenizer dan model dasar NVIDIA Nemotron-3 Nano 4B BF16 dari Hugging Face. Kita juga menetapkan direktori keluaran untuk adapter LoRA dan membatasi panjang sekuens hingga 1024 token agar pelatihan tetap dapat dikelola pada GPU 24GB.

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. Membuat Fungsi Bantu untuk Generasi

Sebelum fine-tuning, kita akan membuat beberapa fungsi bantu untuk menguji respons model. Fungsi-fungsi ini membangun prompt chat, menghasilkan jawaban, menghapus tag pemikiran yang tidak diinginkan, dan menyimpan hasilnya dalam tabel perbandingan kecil.

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. Menjalankan Evaluasi Sampel Pra-Fine-Tuning

Sebelum pelatihan, kita akan menghasilkan beberapa respons dari model dasar Nemotron-3. Ini memberi kita baseline sehingga nanti dapat dibandingkan bagaimana model merespons sebelum dan sesudah fine-tuning LoRA.

Di sini, kita memilih tiga contoh dari set uji dan menghasilkan jawaban menggunakan fungsi bantu yang kita buat sebelumnya.

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

Keluaran berupa tabel kecil dengan pertanyaan asli, jawaban acuan dari response_j, dan jawaban yang dihasilkan oleh model dasar. Tabel ini akan berguna nanti saat kita membandingkannya dengan respons model yang sudah di-fine-tune.

7. Mengonfigurasi LoRA dan Pengaturan Pelatihan

Sekarang kita akan menyiapkan model untuk fine-tuning LoRA. Kita mengaktifkan gradient checkpointing untuk mengurangi penggunaan memori, lalu membuat konfigurasi LoRA yang menargetkan semua layer linear dalam model.

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

Selanjutnya, kita mendefinisikan pengaturan supervised fine-tuning menggunakan SFTConfig. Pengaturan ini mengontrol ukuran batch, laju pembelajaran, jumlah epoch, frekuensi evaluasi, strategi penyimpanan, dan pelatihan 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. Melatih dan Menyimpan Adapter LoRA

Sekarang kita dapat membuat SFTTrainer, memasang konfigurasi LoRA, dan memulai fine-tuning. Sebelum pelatihan, kita juga memeriksa berapa banyak parameter yang dapat dilatih untuk memastikan adapter LoRA terpasang dengan benar.

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

Selama pelatihan, training loss dan validation loss seharusnya berkurang secara bertahap. Ini biasanya berarti model mempelajari gaya respons dari dataset.

Setelah pelatihan, simpan adapter LoRA dan tokenizer secara lokal:

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

Lalu dorong adapter yang telah di-fine-tune ke 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)

Adapter yang telah di-fine-tune sekarang tersimpan secara lokal dan diunggah ke Hugging Face di bawah HUB_REPO_ID.

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

9. Membandingkan Respons Sebelum dan Sesudah Fine-Tuning

Terakhir, kita akan menghasilkan jawaban dari model yang telah di-fine-tune dan membandingkannya dengan keluaran model dasar. Ini membantu kita melihat apakah fine-tuning LoRA meningkatkan keselarasan model dengan jawaban acuan.

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

Model yang di-fine-tune menjadi lebih selaras dengan gaya jawaban acuan. Responsnya lebih ringkas dan tetap dekat dengan jawaban dalam dataset. Namun, model dasar terkadang memberikan respons yang lebih detail dan praktis.

Sebagai contoh, model yang di-fine-tune meningkatkan keselarasan pada pertanyaan terkait manajemen stres dan konsentrasi, tetapi model dasar memberikan respons yang lebih kuat untuk contoh terkait tidur karena menyertakan detail yang lebih membantu.

Secara keseluruhan, model yang di-fine-tune lebih baik jika tujuan Anda adalah mencocokkan gaya dataset acuan. Jika tujuan Anda adalah kemampuan membantu semaksimal mungkin, model dasar mungkin masih lebih baik dalam beberapa kasus karena dapat memberikan jawaban yang lebih hangat dan detail.

Jika Anda mengalami masalah saat menjalankan kode di atas, lihat notebook di repo Hugging Face: fine-tune-nemotron-nano.ipynb

Pemikiran Akhir

Bahkan setelah melakukan fine-tuning pada 100+ LLM, model ini memerlukan kerja penyiapan lebih banyak dari yang diperkirakan. Tantangan utama adalah dependensi mamba_ssm, yang mudah rusak atau konflik dengan lingkungan Python lokal yang ada.

Karena itu, saya menyarankan menggunakan lingkungan yang bersih untuk alur kerja ini. Dalam kasus saya, jalur termudah adalah membangun ulang lingkungan, memasang versi PyTorch yang benar, mem-pin paket terkait Mamba, lalu menjalankan notebook dari sana.

Batasan lain adalah kuantisasi. Untuk setelan ini, saya tidak bisa begitu saja memuat model dalam 4-bit dan melakukan fine-tuning seperti alur kerja QLoRA standar, seperti dalam tutorial Qwen3.5 Small saya. Saya harus memuat model BF16 penuh lalu melakukan fine-tuning dengan LoRA. Untuk model 4B, ini masih dapat dikelola pada GPU 24GB, tetapi untuk model 12B ke atas, penggunaan memori dapat dengan cepat menjadi masalah.

Meski begitu, fine-tuning pada GPU konsumen kini jauh lebih mudah diakses. Dengan kartu 24GB seperti RTX 3090, kini memungkinkan untuk menyesuaikan model open-source yang kuat ke gaya atau domain tertentu tanpa memerlukan kluster pelatihan besar.

Secara keseluruhan, Nemotron-3 Nano adalah model yang mumpuni, tetapi membutuhkan penyiapan lingkungan yang cermat. Setelah dependensi berfungsi, model ini dapat di-fine-tune dengan baik dan beradaptasi ke gaya respons baru dengan jumlah contoh yang relatif sedikit.