A working config, line by line

The whole thing, then the reasons

This is a complete supervised fine-tuning setup for an 8B model on a 24 GB card, using TRL, PEFT and bitsandbytes, all free, all open source.

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
import torch

MODEL = "meta-llama/Llama-3.1-8B-Instruct"

bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                         bnb_4bit_use_double_quant=True,
                         bnb_4bit_compute_dtype=torch.bfloat16)

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, quantization_config=bnb, torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2", device_map={"": 0})

peft_cfg = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none",
    target_modules="all-linear", task_type="CAUSAL_LM")

args = SFTConfig(
    output_dir="out/triage-v3",
    num_train_epochs=2,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_grad_norm=1.0,
    bf16=True,
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    max_length=2048,
    packing=False,
    group_by_length=True,
    assistant_only_loss=True,
    logging_steps=10,
    eval_strategy="steps", eval_steps=100,
    save_steps=200, save_total_limit=3,
    seed=0,
    report_to="tensorboard")

trainer = SFTTrainer(model=model, args=args, peft_config=peft_cfg,
                     train_dataset=train, eval_dataset=val,
                     processing_class=tok)
trainer.train()

The lines that are decisions

Llama-3.1-8B-Instruct, not -8B. Start from the instruct variant when your task is instruction-shaped. It already follows instructions, has a chat template, and needs far less data to specialise. Start from the base only for continued pretraining or when you want no inherited behaviour at all.

attn_implementation="flash_attention_2". Halves memory on the attention path and speeds it up. Requires Ampere or newer; on a T4 omit this and PyTorch’s SDPA backend will be used instead.

device_map={"": 0}. Everything on GPU 0. Using "auto" on a single-GPU machine can silently offload layers to CPU when memory is tight, turning a slow run into a very slow one for reasons that are not logged prominently.

target_modules="all-linear". Architecture-agnostic, so this config survives a change of model family.

num_train_epochs=2. Two passes. Check the resulting step count: with 6,000 examples at an effective batch of 16 that is 750 steps, comfortably in range.

learning_rate=2e-4. Standard for LoRA. For full fine-tuning this line becomes 2e-5, and forgetting that is the single most expensive typo in this file.

gradient_checkpointing_kwargs={"use_reentrant": False}. The non-reentrant implementation is the one that composes correctly with PEFT. With the old default you can get a run where gradients never reach the adapter and the loss is perfectly flat.

packing=False with group_by_length=True. The conservative choice: most of the padding saving, none of the cross-document attention risk. Turn packing on once you have verified that your version separates documents correctly.

assistant_only_loss=True. Loss on the assistant turns only. This requires the model’s chat template to mark generation spans; if it does not, use the completion-only collator instead and pass the exact response template string.

optim="paged_adamw_8bit". Paged, so a spike on one long batch is absorbed by CPU memory rather than ending the run.

seed=0. Recorded, so the run is as reproducible as the hardware allows.

What changes for a full fine-tune

The same file, five lines different:

# drop the BitsAndBytesConfig entirely, and peft_config=None
learning_rate = 2e-5          # not 2e-4
weight_decay  = 0.01          # not 0
optim         = "adamw_bnb_8bit"
num_train_epochs = 1          # full fine-tuning forgets faster

And the memory arithmetic from module three applies, so this configuration belongs on a sharded multi-GPU setup for anything above about 1.5B parameters.

The evaluation callback

One addition worth making from the first run rather than the fifth. Alongside eval_loss, generate on three fixed prompts at every evaluation step and log the text, and run the canary and refusal sets if they are cheap enough:

class Watch(TrainerCallback):
    def on_evaluate(self, args, state, control, model=None, **kw):
        for p in FIXED_PROMPTS:
            print(state.global_step, generate(model, p)[:200])

Every checkpoint then arrives with evidence attached, which is what makes checkpoint selection on a task metric possible rather than aspirational.

The four lines to change first

If something is wrong, in this order: learning_rate, max_length (check how many examples it truncates), num_train_epochs, and whether the masking line is actually taking effect.

Doing this without writing Python

Axolotl and Llama-Factory express the same configuration as a YAML file and handle the wiring:

base_model: meta-llama/Llama-3.1-8B-Instruct
load_in_4bit: true
adapter: qlora
lora_r: 16
lora_alpha: 32
lora_target_linear: true
sequence_len: 2048
micro_batch_size: 2
gradient_accumulation_steps: 8
num_epochs: 2
learning_rate: 0.0002
lr_scheduler: cosine
warmup_ratio: 0.03
bf16: auto
gradient_checkpointing: true
datasets:
  - path: ./train.jsonl
    type: chat_template

Both are free. For a first project a config file is better than a script, because a file diffs cleanly, versions naturally, and is much harder to make subtly inconsistent between two runs.