Modelpublic

BrainHealthAI/MedQA-Llama3.1-8B-SFT-Big

sourceHugging Facellama3.1updated 5mo agoView on Hugging Face

0likes6downloads

Model Card

MedQA-Llama3.1-8B-SFT-Big

QLoRA fine-tune of Llama-3.1-8B-Instruct on 50,000 trilingual medical Q&A (EN, FR, Moroccan Darija) augmented with the Dorosz Causal Knowledge Graph. Part of the BRAIN HEALTH / Operation HELIX-FT project.

Output format: model wraps the final answer in <answer>...</answer> and replies in the same language as the question (EN, FR, or Darija).

Training data

Item	Value
Source	`Williamsanderson/MedQA-Darija-MultiLingual`
Total trilingual-complete rows	102,057
Training subset (this run)	50,000 stratified by `specialty_id`
Specialties covered	57 (~877 samples / specialty, balanced)
Held-out eval split	1,000 (random 2%)
Languages	English, French, Moroccan Darija (round-robin per row)
KG augmentation	Dorosz Causal KG — 348 drug terms, 34,344 edges (route, indication, contraindication, dosage) prepended to user message when drug mentions are detected in the gold answer

Training recipe (QLoRA)

Setting	Value
Base model	`meta-llama/Llama-3.1-8B-Instruct`
Quantization	4-bit NF4 + double quant (Dettmers et al. 2023)
Compute dtype	bfloat16
LoRA rank	64
LoRA α / dropout	128 / 0.1
LoRA target modules	qproj, kproj, vproj, oproj, gateproj, upproj, down_proj
Trainable parameters	167,772,160 / 8,198,033,408 (2.05%)
Effective batch size	16 (perdevice 2 × gradaccum 8)
Learning rate	2e-4, cosine schedule
Warmup ratio	0.03
Max gradient norm	1.0
Optimizer	pagedadamw8bit
Max sequence length	2048
Epochs (planned / actual)	3 / 2.20 (early-stopped on eval_loss patience=3)
Hardware	RunPod L40S 48 GB (Ada, 18176 CUDA cores)
Total training time	13.86 hours (49895s)
Throughput	2.95 samples/sec, 0.184 steps/sec
Total FLOPs	2.41e+18

Training metrics

Metric	Value
Train samples seen	49,000
Eval samples (held-out)	1,000
Best eval loss ⭐	0.7685 (epoch 1.96, step ~5860)
Final eval loss	0.8210 (overfit indicator)
Final train loss	0.4443
Average train loss	0.7405
Final epoch reached	2.20 / 3 (early stop)
Total optimizer steps	6750

The eval loss reached its minimum around step 5860 (epoch 1.96) at 0.7685. Beyond epoch 2.0 the model showed gradual overfitting (evalloss ↑ to 0.82 at epoch 2.20). The `loadbestmodelat_end=True` callback ensures the published adapter is the best checkpoint, not the final one.

Training curves

Training curves (loss / lr / grad_norm)

Evaluation metrics over training

GPU utilization & memory

Usage

python

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

base = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16, device_map="auto",
)
model = PeftModel.from_pretrained(base, "BrainHealthAI/MedQA-Llama3.1-8B-SFT-Big")
tok = AutoTokenizer.from_pretrained("BrainHealthAI/MedQA-Llama3.1-8B-SFT-Big")

SYSTEM_FR = (
    "Vous êtes un assistant médical rigoureux. Répondez TOUJOURS en français. "
    "Raisonnez d'abord entre <think>...</think>, puis donnez la réponse finale "
    "entre <answer>...</answer>."
)

msgs = [
    {"role": "system", "content": SYSTEM_FR},
    {"role": "user",   "content": "Question : Quels sont les symptômes du diabète de type 2 ?"},
]
inputs = tok.apply_chat_template(msgs, return_tensors="pt", add_generation_prompt=True).to(model.device)
out = model.generate(inputs, max_new_tokens=512, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

For Darija (Arabic-script) questions, the model replies in Arabic-script Darija. For English questions, in English. The system prompt MUST instruct the language explicitly to avoid drift.

Companion model

For comparison, see `Williamsanderson/MedQA-Llama3.1-8B-SFT-Small` — trained on a different (smaller, single-language) corpus without KG augmentation.

Limitations

—Prototype R&D only — not a certified medical device. Always defer to a qualified physician.
—Possible hallucinations despite fine-tuning. Confidence calibration not yet evaluated.
—Darija outputs use Arabic script. Quality is uneven across the long tail of rare specialties (only ~877 trilingual samples / specialty).
—Training stopped at 2.20 epochs (early stop). A full 3-epoch run might marginally improve eval performance.

References

—QLoRA: Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. NeurIPS 2023. arXiv:2305.14314
—LoRA: Hu et al. (2022). arXiv:2106.09685
—Llama-3.1: Grattafiori et al. (2024). arXiv:2407.21783
—BRAIN HEALTH / Operation HELIX-FT — internal project document, MVP 2026.

Citation

bibtex

@misc{medqa_sft_big_2026,
  title  = { MedQA-Llama3.1-8B-SFT-Big: Trilingual medical QA via QLoRA SFT on Llama-3.1-8B with Dorosz KG },
  author = { BRAIN HEALTH project — Operation HELIX-FT },
  year   = { 2026 },
  url    = { https://huggingface.co/BrainHealthAI/MedQA-Llama3.1-8B-SFT-Big }
}