Optimising Llama 3

A pre-trained Llama model knows general language. Optimising specialises it for your task, medical question answering, support routing, whatever it is, by continuing training on a smaller, task-specific dataset. But fine-tuning is not always the right first move, and even when it is, you rarely need to retrain every weight. This guide walks the full pipeline, from preparing a dataset to running supervised fine-tuning, evaluating with ROUGE, and then making the whole thing dramatically cheaper with LoRA and quantization.

When to fine-tune at all

It helps to see fine-tuning as one option on a ladder of increasing cost and power. Prompt engineering is free and the right first try when a cleverly worded prompt gets you there. Retrieval-augmented generation is cheap and the answer when you mainly need the model to draw on fresh or private knowledge rather than behave differently. Full fine-tuning, retraining all the weights, is expensive in GPU time and reserved for when the task is structurally different from anything prompting can reach. LoRA fine-tuning gets most of the quality of full fine-tuning for a fraction of the compute. And quantization, separate from all of these, shrinks a model to fit in less memory at inference. The decision is a short chain: if prompting or retrieval solves it, stop there; if you genuinely need a behaviour change and you have abundant GPU, full fine-tune; otherwise reach for LoRA, and add quantization for cheaper inference.

Preparing the data

Most pipelines start by carving a small, fixed slice off a larger dataset, and there is a subtle trap in doing so with Hugging Face datasets.

from datasets import load_dataset, Dataset
ds = load_dataset(dataset_name, split="train")
filtered_ds = Dataset.from_dict(ds[:500])

Slicing with ds[:500] looks like slicing a list, but it does not return another Dataset. It returns a plain dictionary of lists, one list per column, which has none of the methods like .map and .filter you will want next. Wrapping it in Dataset.from_dictrepacks that dictionary into a proper Dataset. Think of it as unpacking a box, taking the first five hundred items, and packing them back into a box of the same type. Miss this and the next step breaks with a confusing error.

The model then needs each example as a single, consistently formatted text string, so you merge the raw columns into a prompt template.

def format_example(row):
    row['formatted_text'] = f"Message: {row['message']}\nCategory: {row['category']}"
    return row
processed_dataset = dataset.map(format_example)
print(processed_dataset[0]['formatted_text'])

The function stamps every row into a fixed layout, like flashcards that always follow the same format, and .map runs it across the whole dataset to produce a new column. The point that catches people out later is that the model learns the format, not the content, so whatever template you train on you must reproduce exactly at inference time, or the model will not recognise the pattern it memorised.

Because preprocessing a large dataset can take a long time, save the result so you never redo it.

from datasets import load_from_disk
ds.save_to_disk("ds_preprocessed")
ds_preprocessed = load_from_disk("ds_preprocessed")
print(ds_preprocessed[0])

This is meal-prep for data. save_to_disk writes the dataset in the fast columnar Arrow format and load_from_disk reads it straight back as a full Dataset, and the print is a quick round-trip check that nothing was corrupted.

Defining the training recipe

There are two common ways to drive Llama fine-tuning, TorchTune with a config, and Hugging Face’s training arguments, and it is worth seeing both. TorchTune uses a configuration dictionary that you can serialise to YAML, describing the model, optimizer, dataset, and runtime.

config_dict = {
    "model": {"_component_": "torchtune.models.llama3_2.llama3_2_1b"},
    "batch_size": 8,
    "device": "cuda",
    "epochs": 15,
    "optimizer": {"_component_": "bitsandbytes.optim.PagedAdamW8bit", "lr": 3e-05},
    "dataset": {"_component_": "custom_dataset"},
    "output_dir": "/tmp/finetune_results"
}

A recipe like this is a complete cooking instruction: which ingredient (the model), how hot the oven (the device), how long to bake (the epochs), and how to stir (the optimizer). The _component_ key is TorchTune’s convention for “import this class from this dotted path and instantiate it.” The paged AdamW optimizer is a regular Adam with a memory manager bolted on, paging optimizer state to CPU when GPU memory fills, much as your operating system pages RAM to disk. Saving the dictionary as YAML separates the training specification from the code, so you can hand the recipe to the TorchTune command line or a colleague, and swapping llama3_2_1b for llama3_2_3b is a one-line change that quadruples model capacity for a quick comparison, with an 8B variant available when you need more again.

The Hugging Face path uses TrainingArguments, and a useful first run is a deliberate smoke test, a high learning rate and a tiny step count just to confirm the loop runs end to end.

from transformers import TrainingArguments
training_arguments = TrainingArguments(
    learning_rate=2e-3,
    warmup_ratio=0.03,
    num_train_epochs=3,
    output_dir='/tmp',
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    save_steps=10,
    logging_steps=2,
    lr_scheduler_type='constant',
    report_to='none'
)

This is the logbook you fill in before a run, defining every dial for speed, memory, and checkpointing. The learning rate of 2e-3 is intentionally large because the goal here is only to prove the loop does not crash, not to produce a good model, and production runs typically sit between 1e-5 and 5e-5. Gradient accumulation of one means no accumulation, one mini-batch per update, and you raise it to simulate a bigger effective batch on a small GPU. The constant scheduler keeps the rate flat after warmup, and report_to='none' is the don’t-call-home flag that keeps the run self-contained rather than pinging an external tracker.

Running supervised fine-tuning

With data and arguments ready, the TRL library’s SFTTrainer wraps the Hugging Face trainer for supervised fine-tuning of causal language models.

from trl import SFTTrainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=training_arguments,
)
trainer.train()

SFTTrainer knows how to handle the formatted prompt strings you built earlier and manages causal-LM specifics like loss masking. Hand it the model, the tokenizer, the data, and the arguments, and trainer.train() runs the whole optimisation loop on its own: forward pass, loss, backward pass, weight update, repeat, checkpoint, and log, until the configured number of epochs is reached.

Measuring quality with ROUGE

To judge a fine-tuned model on a generation task, ROUGE compares generated text against reference text by counting overlapping n-grams.

import evaluate
rouge_evaluator = evaluate.load("rouge")
results = rouge_evaluator.compute(
    predictions=test_answers,
    references=reference_answers
)
print(results["rouge2"])

ROUGE is a librarian checking your model’s output against the reference by counting shared phrases. ROUGE-2 in particular counts matching two-word sequences, so “the red car” against “the red vehicle” scores a match only on “the red.” It is recall-oriented, meaning a higher score says the model captured more of what the reference contained. One call computes all the variants, ROUGE-1 for single words, ROUGE-2 for word pairs, ROUGE-L for the longest common subsequence, and ROUGE-Lsum computed sentence-wise for summaries, and you pick the one your task cares about.

LoRA: fine-tuning without retraining everything

Full fine-tuning updates billions of weights, like repainting an entire building. LoRA instead freezes the base model and slides thin trainable adapter matrices in front of selected weights, painting only those panels and leaving the building untouched, which means roughly ninety-nine percent fewer trainable parameters.

from peft import LoraConfig
lora_config = LoraConfig(
    r=12,
    lora_alpha=8,
    task_type="CAUSAL_LM",
    lora_dropout=0.05,
    bias="none",
    target_modules=['q_proj', 'v_proj']
)
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    tokenizer=tokenizer,
    args=training_arguments,
    peft_config=lora_config,
)

The rank controls how thick those adapter panels are, with higher meaning more expressive but more parameters, while the alpha scales the adapter’s contribution, the effective scale being alpha over r. The task type tells PEFT to expect a decoder-only model like Llama, and the target modules inject adapters into the query and value attention projections, the standard attention-only choice. Passing peft_config to SFTTrainer is all it takes for the trainer to wrap the model with adapters before training begins. A common minimal starting point is a rank of two with alpha set to twice the rank, the thinnest useful adapter, following the empirically safe rule that alpha equals two times r. The practical advice is to start small, measure ROUGE, and increase the rank, four, eight, sixteen, before ever touching alpha.

Quantization: fitting the model in less memory

Quantization is a separate lever that shrinks the stored weights so a model fits in less GPU memory. Eight-bit is the gentle option.

from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "Maykeye/TinyLLama-v0",
    quantization_config=bnb_config,
    low_cpu_mem_usage=True
)

Full-precision weights are sixteen or thirty-two bit floats, high detail but large, and eight-bit rounds each to a coarser value, like compressing a photo from RAW to a high-quality JPEG, losing a little fidelity for roughly half the size. The low_cpu_mem_usage flag avoids the double-memory penalty of loading the full model into CPU RAM before moving it to the GPU, streaming the weights instead.

When eight-bit is not small enough, four-bit is the aggressive option, but it needs two refinements to stay usable.

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(
    "Maykeye/TinyLLama-v0",
    quantization_config=bnb_config,
    low_cpu_mem_usage=True
)

Four-bit storage allows only sixteen possible values per weight. Plain four-bit spreads those values evenly across a range, but network weights are not evenly distributed, they cluster near zero with rare large outliers, so NF4, or Normal Float 4, places the sixteen values according to a normal distribution and puts more of them where the weights actually are, losing less quality. Setting the compute dtype to bfloat16 is the speed half of the trick: the weights stay tiny at four bits, but the actual matrix multiplications run at sixteen-bit precision, which is fast on modern GPUs and more numerically stable than fp16. Stepping down from fp16 to 8-bit roughly halves memory, and 4-bit roughly quarters it, with NF4 plus bfloat16 giving you that quarter-size footprint without the quality collapse plain four-bit would cause.

The pipeline end to end

Stripped back, the whole flow is a chain. Load the dataset and slice it safely with Dataset.from_dict, map your prompt template across it, and save the result to disk so you never reprocess. Configure the run, either as a TorchTune recipe or as Hugging Face TrainingArguments, then hand model, tokenizer, data, and arguments to SFTTrainer and call train. To make it affordable, add a LoraConfig so only adapter weights update, and a BitsAndBytesConfig to load the base model quantized. Finally, generate on a held-out slice and score with ROUGE to see whether the fine-tune actually helped. Start with LoRA at a small rank and quantize for inference, measure, and scale up only the dial that the numbers tell you to.

See you soon.