Step-by-Step LoRA Fine-Tune of Llama 3 8B on RTX 4090 24GB

The single most useful thing you can do with a fresh RTX 4090 24GB dedicated server in your first week is run an end-to-end LoRA fine-tune of Llama 3.1 8B. It validates the driver, CUDA toolkit, FlashAttention 2 build, bitsandbytes optimiser path and your data pipeline in one shot, and at the end you have a small adapter you can serve straight back through vLLM. On Ada AD102 silicon with its 16,384 CUDA cores, 1008 GB/s GDDR6X bandwidth and native FP8 tensor cores, you should expect roughly 18,000 training tokens per second, which means a 5-epoch run over 10,000 samples completes in well under two hours. This guide is the recipe we hand to dedicated GPU hosting customers on day one and walks through every step with the reasoning behind each knob.

Contents

• Prerequisites and stack pinning
• Data preparation and chat templating
• Hyperparameter table with rationale
• The training script
• Run, monitor and verify
• Evaluation and overfit detection
• Merge or serve as adapter
• Production gotchas

Prerequisites and stack pinning

Assume the vLLM setup tutorial has been completed (NVIDIA driver 550+, CUDA toolkit 12.4, Python 3.11 venv) and the day-one steps from the RTX 4090 first-day checklist are in place. Pinning versions is not optional: bitsandbytes, peft, trl and transformers all release breaking changes monthly.

source ~/vllm-env/bin/activate
pip install \
  torch==2.4.0 \
  transformers==4.45.2 \
  peft==0.13.2 \
  trl==0.11.4 \
  accelerate==1.0.1 \
  bitsandbytes==0.44.1 \
  datasets==3.0.1 \
  evaluate==0.4.3 \
  flash-attn==2.6.3 --no-build-isolation
export TORCH_CUDA_ARCH_LIST=8.9
export HF_TOKEN=hf_yourtoken
export TOKENIZERS_PARALLELISM=false

Why these versions: PyTorch 2.4 is the first release with stable SDPA on Ada at sequence length 4096, transformers 4.45 introduced the official Llama 3.1 chat template, and FlashAttention 2.6 fixes a numerical drift seen on Ada at BF16. TORCH_CUDA_ARCH_LIST=8.9 tells any source build to compile only for Ada Lovelace, which roughly halves wheel build times. Verify the GPU is visible to PyTorch before going further.

python -c "import torch; print(torch.cuda.get_device_name(0), torch.cuda.is_bf16_supported(), torch.backends.cuda.flash_sdp_enabled())"
# Expected: NVIDIA GeForce RTX 4090 True True

Data preparation and chat templating

The recipe expects a JSONL file with messages in the OpenAI chat format. TRL’s SFTTrainer applies the model’s native chat template automatically when the dataset uses this shape.

{"messages":[{"role":"system","content":"You are a..."},
              {"role":"user","content":"..."},
              {"role":"assistant","content":"..."}]}

10,000 to 30,000 high-quality samples is the right starting range for a domain or style transfer LoRA. Larger datasets help on knowledge injection but increase the risk of overfitting and catastrophic forgetting; for that scale, consider continued pre-training instead. Always reserve 5% as a held-out evaluation split and keep a third tiny set (50-100 samples) for spot-checking generations.

Hyperparameter table with rationale

The training script

Save as train_lora.py. The structure follows the official TRL recipe with explicit settings rather than the auto-defaults so behaviour is reproducible across versions.

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig

model_id = "meta-llama/Llama-3.1-8B-Instruct"
ds = load_dataset("json", data_files="train.jsonl", split="train")
ds = ds.train_test_split(test_size=0.05, seed=42)

tok = AutoTokenizer.from_pretrained(model_id)
tok.pad_token = tok.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
    device_map="cuda:0",
)

lora = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none",
    target_modules=["q_proj","k_proj","v_proj","o_proj",
                    "gate_proj","up_proj","down_proj"],
    task_type="CAUSAL_LM",
)

cfg = SFTConfig(
    output_dir="out/llama3-8b-lora",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    optim="paged_adamw_8bit",
    bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    logging_steps=20,
    save_steps=500,
    eval_steps=200,
    eval_strategy="steps",
    max_seq_length=2048,
    packing=True,
    neftune_noise_alpha=5,
    report_to="tensorboard",
)

trainer = SFTTrainer(
    model=model, tokenizer=tok,
    train_dataset=ds["train"], eval_dataset=ds["test"],
    peft_config=lora, args=cfg,
)
trainer.train()
trainer.model.save_pretrained("out/llama3-8b-lora/adapter")
tok.save_pretrained("out/llama3-8b-lora/adapter")

Two non-obvious additions: use_reentrant=False in gradient checkpointing avoids the deprecated PyTorch path that throws warnings and occasionally hangs on Ada at the epoch boundary; and neftune_noise_alpha=5 adds embedding noise during training, which Hugging Face benchmarks show lifts MT-Bench by 1.5-2 points at zero throughput cost.

Run, monitor and verify

python train_lora.py 2>&1 | tee train.log
# in another shell
nvidia-smi -l 2
nvtop
tensorboard --logdir out/llama3-8b-lora --port 6006 --bind_all

Within the first minute you should see steady-state behaviour. If any signal sits outside these ranges for more than five logging steps, kill the run and check the gotchas section.

For raw throughput numbers across rank, batch and sequence-length sweeps, see the dedicated fine-tune throughput reference. Adding Unsloth (pip install unsloth) gives roughly a 1.8x speed-up with no measurable accuracy loss; the trade-off is a less hackable training loop.

Evaluation and overfit detection

Train and eval loss should both fall for the first epoch. From epoch 2 onward the eval loss flattens; if it starts to climb while train loss keeps falling, you are overfitting. The cure is fewer epochs, more data, higher dropout or lower rank, in that order. After training, run a generation pass on the held-out 50-100 sample set and read the outputs by hand. Automated metrics catch fluency regressions but they miss things like the model picking up dataset artefacts (always replying with bullet lists, repeating a system prompt verbatim).

Merge or serve as adapter

Two options after training, both production-grade:

1. Merge into base: model = model.merge_and_unload(), then save as a normal model directory. Best for single-tenant production where you want maximum throughput; serve with the FP8 deployment guide directly because vLLM will quantise the merged BF16 weights to FP8 at load time.
2. Serve as adapter: vLLM supports LoRA adapters at runtime via --enable-lora --lora-modules my-adapter=/path/to/adapter. Best for the multi-tenant patterns described in the multi-tenant SaaS writeup, where many customers each have their own small adapter swapped in per request.

For larger models that do not fit at BF16, switch from this LoRA recipe to QLoRA: the QLoRA fine-tune guide covers Llama 3 70B on a single 4090 with NF4 plus paged AdamW.

Production gotchas

1. FP16 instead of BF16: paged AdamW underflows in FP16 on Ada and you get NaN losses after ~50 steps. Always set bf16=True on the 4090.
2. FlashAttention compilation failure: flash-attn requires --no-build-isolation and a matching CUDA toolkit. If you see undefined symbol errors at import, your PyTorch and CUDA versions are mismatched.
3. Hugging Face rate limits: gated models like Llama 3 fail with HTTP 401 if HF_TOKEN is missing or the model has not been accepted on the model card. Always test with huggingface-cli whoami first.
4. Disk filling up: each checkpoint at save_steps=500 is ~150 MB for r=16 but ~1.5 GB if you accidentally save the merged model. Set save_total_limit=3 or you will fill 2 TB NVMe by epoch 2 of a long run.
5. Tokenizer pad token: forgetting tok.pad_token = tok.eos_token on Llama 3 produces the cryptic error “Asking to pad but the tokenizer does not have a padding token”, and silently breaks attention masks if you patch it wrong.
6. Grad accumulation off-by-one: TRL versions before 0.11 had a bug where loss was averaged over micro-batches incorrectly, inflating reported loss by ~2x. Pin trl==0.11.4 or newer.
7. Power throttling under poor airflow: 410 W sustained for 90 minutes will heat-soak a poorly-cooled chassis; review thermal performance and power draw efficiency notes if your hot-spot exceeds 90 C.

Verdict

A single dedicated 4090 will fine-tune Llama 3.1 8B end-to-end in under two hours per 3-epoch pass over 10k samples. Stack pinning, BF16, FlashAttention 2, packing and paged AdamW are non-negotiable; everything else is workload-dependent tuning. Compared to renting an A100 by the hour for the same job, you will spend roughly £3-4 of dedicated server time on a £550/month box versus $30-50 on cloud, and you keep the box for the rest of the month for inference. For lighter cards see the 5060 Ti QLoRA guide; for the next tier up see the 4090 or 5090 decision.

See also: QLoRA fine-tune guide, fine-tune throughput, fine-tune box use case, multi-tenant SaaS, vLLM setup, spec breakdown, FP8 deployment.