Optimizing Llama 3 fine-tuning for Python code generation using Unsloth

Unsloth is easily the biggest quality-of-life upgrade for anyone fine-tuning Llama 3 on a single GPU, primarily because it kills the OOM (Out of Memory) errors that usually plague Python code datasets. If you're trying to teach Llama 3 a specific internal library or a niche coding style, standard LoRA often feels sluggish or crashes when you hit those long context windows typical of complex Python scripts.

The real trick to getting high-quality code generation isn't just the dataset—it's the memory management. I've been using a 24GB RTX 3090, and with Unsloth's 4-bit quantization, I can push the sequence length to 4096 without sweating. The speedup is roughly 2x compared to vanilla Hugging Face trainers.

To get this running, you need to initialize the model specifically for LoRA. Don't just load the weights; use the FastLanguageModel wrapper to ensure the kernels are optimized.

from unsloth import FastLanguageModel
import torch

max_seq_length = 4096 
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = max_seq_length,
    load_in_4bit = True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Rank 16 is usually the sweet spot for code; 32 is overkill unless the syntax is wildly different
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0, # Optimized kernels prefer 0 dropout
    bias = "none",    
)

One major gotcha with Python code generation is the prompt template. Llama 3 is sensitive. If you don't strictly wrap your code in the correct chat format, the model starts hallucinating indentation or forgetting to close parentheses. I found that using a clear "Instruction -> Input -> Response" structure inside the Llama 3 template prevents the model from "bleeding" the prompt into the code block.

Pro-tips for the dataset config:

Weight the loss on the response only. You don't want the model learning how to write the prompt; you want it learning the Python implementation.

Use a higher learning rate than you think. For code, I've found 2e-4 works better than the standard 5e-5. Code is structural; the model needs to shift its weights more aggressively to capture the logic flow.

Set weight_decay to 0.01. This prevents the model from over-fitting on specific variable names in your training set, which is a common issue when fine-tuning on a small codebase.

When it comes to the actual training loop, the SFTTrainer from the TRL library integrates perfectly with Unsloth. Just make sure you're using the packing = True option. Packing combines multiple short code snippets into a single sequence, which drastically reduces the number of padding tokens and speeds up training by about 30%.

from trl import SFTTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60, 
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
    ),
)

The biggest productivity gain here is the export process. Instead of dealing with the nightmare of merging LoRA adapters manually, Unsloth lets you save to GGUF or 16-bit floats directly. I usually export to GGUF and plug it into Ollama for local testing. It makes the iteration loop—train, test, tweak prompt, repeat—take minutes instead of hours.