Fine tuning was performed using PEFT (Parameter Efficient Fine Tuning) with LoRA (Low-Rank Adaptation) on the chat subset of the nvidia/Llama-Nemotron-Post-Training-Dataset dataset.

LoRA Configuration:

lora_config = LoraConfig( task_type="CAUSAL_LM", r=32, lora_alpha=32, lora_dropout=0.1, target_modules=["q_proj", "k_proj", "v_proj"], modules_to_save=["lm_head", "embed_token"], )

Use with Transformers

pip install transformers pip install torch

from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("suwesh/llamatron-1B-peft").to("cuda") tokenizer = AutoTokenizer.from_pretrained("suwesh/llamatron-1B-peft") input_text = "Hello, how are you?" inputs = tokenizer(input_text, return_tensors="pt").to("cuda") outputs = model.generate(**inputs, max_new_tokens=512) response = tokenizer.decode(outputs[0], skip_special_tokens=True)

Or with pipeline

import torch import transformers from transformers import pipeline pipe = pipeline( "text-generation", model=("suwesh/llamatron-1B-peft"), tokenizer=transformers.AutoTokenizer.from_pretrained("suwesh/llamatron-1B-peft"), torch_dtype=torch.bfloat16, device="cuda", ) def to_model(input_text, system_message): messages = [ {"role": "system", "content": system_message}, {"role": "user", "content": input_text} ] outputs = pipe( messages, max_new_tokens=512, temperature=0.6, top_p=0.95 ) return outputs[0]["generated_text"][-1]['content'] response = to_model("Write a joke about windows.", "detailed thinking on")

Load adapter checkpoint for further fine tuning

from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct") tokenizer = AutoTokenizer.from_pretrained("suwesh/llamatron-1B-peft") model = PeftModel.from_pretrained(base_model, "suwesh/llamatron-1B-peft", subfolder="checkpoint-11000")

Training details

Initial Training and Validation losses: 1.69 | 1.67

Checkpoint 11000 Training and Validation losses: 1.06 | 1.09

Evaluation details

We use the nvidia/Llama-3.1-Nemotron-Nano LLM as a Judge for evaluating the responses between the base llama 3.2 1b instruct and our PEFT model. The following are the judge's preference for each prompt to the two models, we also provide the ground truth in the prompt to the judge:

base: 122 peft: 388 tie: 29

system_message = """ You are an expert evaluator comparing two AI responses to a user instruction. Use the following criteria: -Clarity -Factual correctness (compared to the reference answer) -Instruction-following -Depth of reasoning Below is the reference answer, which is the ideal or expected response: """