Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/NicolasHoyosDevss/MaternaQA-es/llms.txt

Use this file to discover all available pages before exploring further.

train_qlora_trl.py implements the canonical Hugging Face QLoRA training path with TRL + PEFT + bitsandbytes. It converts the published conversational JSONL to prompt/completion format in memory, so loss is calculated only over the expected answer without relying on model-specific chat template masks. This makes training robust across Gemma 4 and MedGemma checkpoints, which may not expose assistant turn masks in their tokenizer templates. The base model is never modified — only the LoRA adapter weights are updated and saved.

Full Training Commands

The following commands run the complete fine-tuning over all 5,093 training pairs with the default hyperparameters (2 epochs, batch 1, gradient accumulation 8). Expect 2–4 hours per run depending on GPU speed.
python scripts/train_qlora_trl.py \
  --model-name google/gemma-4-E2B-it \
  --dataset-variant sft_grounded \
  --output-dir outputs/gemma4-grounded

Key Training Defaults

The script ships with conservative defaults tuned for a 16 GB VRAM workstation. All values below come directly from train_qlora_trl.py.
ParameterDefaultDescription
--num-train-epochs2.0Number of full passes over the training split
--per-device-train-batch-size1Per-GPU batch size; kept low to fit within 16 GB VRAM
--gradient-accumulation-steps8Simulates an effective batch size of 8 without extra memory
--learning-rate2e-4Standard learning rate for LoRA / SFT
--lora-r16LoRA rank — controls adapter capacity
--lora-alpha16LoRA scaling factor
--lora-dropout0.05Dropout applied inside LoRA layers
--max-length1024Maximum token length per training example
--warmup-ratio0.03Fraction of steps used for learning rate warm-up
--max-grad-norm0.3Gradient clipping threshold
--seed3407Random seed for reproducibility
--gradient-checkpointingTrueReduces VRAM at the cost of extra compute
--packingFalseDisabled to prevent cross-sample contamination

All CLI Flags

Every flag exposed by parse_args() in train_qlora_trl.py is listed below.
FlagTypeDefaultDescription
--model-namestrgoogle/gemma-4-E2B-itHugging Face model ID or local path of the base model
--model-classauto|causal-lm|image-text-to-textautoForces a specific model loader class; auto selects ImageTextToText for Gemma 4 / MedGemma
--dataset-rootPathdatasets/obstetrics/qa/publicationRoot directory containing the publication variant folders
--dataset-variantsft_closed_book|sft_groundedsft_groundedWhich SFT variant to train on
--output-dirPath(required)Directory where the LoRA adapter and tokenizer are saved
--max-lengthint1024Maximum token length per example
--num-train-epochsfloat2.0Number of training epochs
--max-stepsintNoneIf set, overrides --num-train-epochs and stops after N steps
--per-device-train-batch-sizeint1Training batch size per GPU
--per-device-eval-batch-sizeint1Evaluation batch size per GPU
--gradient-accumulation-stepsint8Steps to accumulate before a weight update
--learning-ratefloat2e-4Initial learning rate
--warmup-ratiofloat0.03Fraction of total steps for LR warm-up
--weight-decayfloat0.0L2 regularization coefficient
--max-grad-normfloat0.3Gradient clipping max norm
--lora-rint16LoRA rank
--lora-alphaint16LoRA alpha scaling
--lora-dropoutfloat0.05LoRA dropout rate
--logging-stepsint10Steps between TensorBoard log entries
--save-stepsintNoneSteps between checkpoint saves (used when --save-strategy steps)
--eval-stepsintNoneSteps between evaluations (used when --eval-strategy steps)
--save-strategyno|steps|epochepochWhen to save checkpoints
--eval-strategyno|steps|epochepochWhen to run evaluation
--train-limitintNoneLimit training examples to N (for smoke tests)
--eval-limitintNoneLimit validation examples to N (for smoke tests)
--seedint3407Global random seed
--assistant-only-lossboolFalseAdvanced: use chat template assistant masks instead of prompt/completion split
--packingboolFalsePack short sequences to improve GPU utilization
--gradient-checkpointingboolTrueTrade compute for lower VRAM usage
--trust-remote-codeboolFalseEnable only if the model requires it
--attn-implementationstrNoneOptional attention backend: eager, sdpa, flash_attention_2
--dry-runflagFalseValidate dataset loading without loading the model or training
--allow-cpuflagFalseAllow --dry-run without CUDA (QLoRA training still requires GPU)

Training Outputs

After training completes, --output-dir contains the following files:
outputs/gemma4-grounded/
├── adapter_config.json           # LoRA configuration and base model reference
├── adapter_model.safetensors     # Trained LoRA adapter weights
├── tokenizer.json                # Saved tokenizer
├── tokenizer_config.json         # Tokenizer configuration
├── special_tokens_map.json       # Special token mappings
├── chat_template.jinja           # Chat template saved alongside tokenizer
├── training_args.bin             # Serialized SFTConfig used during training
└── checkpoint-N/                 # Intermediate checkpoints saved per epoch or step
    ├── adapter_model.safetensors
    ├── adapter_config.json
    ├── optimizer.pt
    ├── scheduler.pt
    └── trainer_state.json
The base model weights are never written to disk — only the adapter files. To run inference, you load the original base model from Hugging Face and mount the adapter on top.

Model Compatibility Notes

Gemma 4 (google/gemma-4-E2B-it) is supported by Unsloth for potentially faster training on compatible hardware. The canonical path in this repository is TRL + PEFT + bitsandbytes QLoRA, which is reproducible without Unsloth-specific patches.MedGemma 1.5 4B IT (google/medgemma-1.5-4b-it) should use the TRL + PEFT + bitsandbytes path. Unsloth has MedGemma artifacts on Hugging Face, but validate with a smoke test before relying on Unsloth for that model.Full fine-tuning is not the target workflow. QLoRA adapters are the default and recommended method for this dataset and hardware profile.16 GB VRAM is the minimum for E2B-class models. The default batch size of 1 with gradient accumulation 8 is intentionally conservative to remain within this limit.
The following four experiments are the core benchmark defined in the dataset publication. Each combination measures a different capability axis.
Model FamilyDataset VariantTraining PathPurpose
Gemma 4 instructsft_groundedTRL + PEFT + bitsandbytes (QLoRA)Measure evidence-guided QA performance
Gemma 4 instructsft_closed_bookTRL + PEFT + bitsandbytes (QLoRA)Measure domain internalization without context
MedGemma 1.5 4B ITsft_groundedTRL + PEFT + bitsandbytes (QLoRA)Measure evidence-guided QA on a medical instruct base
MedGemma 1.5 4B ITsft_closed_bookTRL + PEFT + bitsandbytes (QLoRA)Measure domain adaptation without context on a medical model

TensorBoard Monitoring

The trainer writes TensorBoard events to outputs/<model>/runs/ automatically (configured via report_to: ["tensorboard"] in SFTConfig). Launch TensorBoard to track training loss, evaluation loss, and learning rate schedule in real time:
tensorboard --logdir outputs/
Each model’s run appears as a separate entry in the TensorBoard UI, allowing side-by-side comparison of Gemma 4 and MedGemma training curves across both dataset variants.

Build docs developers (and LLMs) love