How to optimize LoRA hyperparameters for fine-tuning Llama 3 on medical datasets

Medical datasets are notorious for having high terminology density and strict structural requirements, which makes standard "out-of-the-box" LoRA settings for Llama 3 fail miserably—usually resulting in either catastrophic forgetting or a model that just mimics the training format without actually "learning" the medical logic.

After several failed runs where the model started hallucinating patient IDs or losing its general conversation ability, I've found that the secret is aggressive target module expansion and a very conservative learning rate.

Most people just target q_proj and v_proj. For medical fine-tuning, that's not enough. You need to hit all the linear layers to capture the nuanced relationships in clinical text. I use the following configuration in my PEFT setup:

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

The r=16 and lora_alpha=32 balance is a sweet spot. If you push r to 64 or 128, you'll notice the model starts overfitting to specific phrasing in your medical corpus, losing its ability to reason. Always keep lora_alpha at double the value of r to maintain scaling stability.

The learning rate is where most medical LoRA attempts crash. If you use 2e-4 (the common default), Llama 3 tends to "collapse" on specialized vocab. I've had much better luck with 5e-5 or even 2e-5 combined with a cosine learning rate scheduler.

Here is the training loop logic I've been using with the Hugging Face Trainer:

training_args = TrainingArguments(
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=5e-5,
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    num_train_epochs=3,
    weight_decay=0.01,
    optim="paged_adamw_32bit",
    fp16=True,
)

A huge "gotcha" with medical data is the padding. Medical notes often have weird lengths. If you use standard padding, the model might learn to associate the padding tokens with specific medical outcomes. I highly recommend using packing=True in the SFTTrainer to concatenate examples and minimize waste, which also speeds up training by about 30%.

Key Productivity Gains & Tips:

• Validation Strategy: Don't rely on training loss. Medical models can have low loss but still produce dangerous hallucinations. Create a "Golden Set" of 100 complex medical Q&As and run them through the model every 500 steps.
• Memory Management: If you're hitting OOM on a 24GB card, switch to bitsandbytes 4-bit quantization (NF4). It has almost zero impact on medical accuracy but lets you increase the batch size.
• Prompt Template: Llama 3 is sensitive. Ensure your medical data is wrapped in the exact Instruct format. If you mix formats during LoRA, the model's coherence drops significantly.

If you find the model is becoming too "robotic" and losing its natural flow, drop the lora_dropout to 0.0 and slightly increase the weight_decay. This forces the model to generalize better across the medical terminology without clinging to specific training samples.