unsloth-fine-tune-gemma-4

Imported from advaitjain/unsloth-fine-tune-gemma-4 (AGENTS.md). Install upstream with npx skills add advaitjain/unsloth-fine-tune-gemma-4. Copyright stays with the author.

AGENTS.md

Guidance for AI coding agents (Claude, Gemini, Cursor, etc.) working in this repo. Read this before making changes.

What this repo is

Worked examples of Unsloth for supervised fine-tuning (SFT) in full 16-bit precision (FP16) on Gemma 4 E2B models, followed by optimized edge deployment and evaluation using LiteRT-LM. Sized end-to-end for GPUs with at least 16 GB of VRAM. The reader's experience is the product — keep examples small, runnable, and explicit about VRAM costs.

User-facing docs live in README.md. Experiment design and historical results live in feature-specific experiments.md files. This file is for agents.

Codebase Map

• GSM8K SFT (Arithmetic Reasoning): gsm8k-math/finetune_gsm8k.py (runs structured math tuning templates in full FP16 precision), gsm8k-math/eval_gsm8k.py (greedy decoding adapter test runner), gsm8k-math/inference_demo.py (interactive side-by-side reporter).
• CUAD SFT (Legal Contract Extraction): cuad/finetune_cuad.py (PEFT adapter fine-tuning loop for governing law extraction), cuad/eval_cuad.py (standalone greedy metrics EM/F1 validation), cuad/inference_demo.py.
• Vision SFT (LaTeX OCR): latex_ocr/run_master_sweeps.py (sequential sweeps coordinator), examples/finetune_vision.py (VLM supervised fine-tuning loop), examples/eval_vision.py (Exact Match and Normalized Edit Distance vision verification pipeline).
• Compiled Edge Execution (LiteRT-LM): litert-lm/README.md (conversion and runtime docs), litert-lm/convert_and_eval.sh (automated end-to-end pipeline script), litert-lm/merge_adapter.py (merges adapters to full precision safetensors), litert-lm/eval_litert_gsm8k.py (evaluates compiled .litertlm models), and examples/litert_lm_inference.py.
• Weight Merging: examples/merge_lora.py (merges adapters to full 16-bit precision safetensors).
• Experimental Sandbox: experimental/README.md (task studies: mapping experimental/train_regex.py for custom formats and experimental/train_emotion.py for boundary semantic matching studies).

Hardware Requirement: 16 GB VRAM

Every default in the repo is chosen to fit within at least 16 GB of VRAM. Before adding or changing anything, verify:

• Full 16-bit precision fine-tuning (load_in_4bit=False, dtype=None or explicit bfloat16/fp16).
• This repo focuses exclusively on Gemma 4 E2B (unsloth/gemma-4-E2B-it).
• Standard 16-bit full-precision SFT adapter training requires approximately 11.7 GB of VRAM.
• Base inference in 16-bit requires approximately 10.3 GB of VRAM.
• Avoid loading pre-quantized 4-bit variants or proposing QLoRA unless explicitly requested, as the core paradigm centers on full-precision (FP16) PEFT adapters suitable for downstream compilation to LiteRT-LM.
• Don't propose full fine-tuning or parameter scaling exceeding the 16 GB constraint without flagging the constraint and asking.

Tooling

• uv is the package manager. Always run Python via uv run python …. Never pip install directly. Never invent a requirements.txt.
• Dependencies live in pyproject.toml. unsloth transitively pulls in torch, transformers, peft, trl, bitsandbytes, accelerate, xformers, triton — don't re-pin those unless you've hit a real conflict.
• The hf CLI ships with huggingface_hub ≥ 0.34. Pre-download models with:

  HF_HUB_ENABLE_HF_TRANSFER=1 uv run hf download <hf-id>

  Same HF cache (~/.cache/huggingface) is reused by all tools.

Code conventions (match existing style)

Look at examples/inference.py and gsm8k-math/finetune_gsm8k.py for the canonical shape. In particular:

• Module-level docstring describing what the script does and any non-obvious CLI usage. Keep it short — a few lines, not a tutorial.
• Constants like DEFAULT_MODEL, DATASET_NAME near the top of the file.
• argparse for CLI args with type=int|float, default=…, help="…". Use both short and long flags only where the existing scripts do (-m/--model, -a/--adapter, -p/--prompt).
• def main() -> None: plus if __name__ == "__main__": main().
• Type hints on function signatures and helper return types.
• Model Loading Conventions:
  • Use generic unsloth.FastModel for standard loading and merging scripts, as it covers both text and vision models natively.
  • Use unsloth.FastVisionModel for vision training entry-points. Avoid using standard transformers loaders to guarantee Unsloth optimizations.
• Multimodal / Vision Configuration:
  • The target visual sequence defaults to 280 soft tokens.
  • To customize soft visual budgets dynamically, you must implement synchronization overrides directly on both the Model and Processor config objects:

    # 1. Modify Model Config
    model.config.vision_soft_tokens_per_image = args.vision_tokens
    model.config.vision_config.default_output_length = args.vision_tokens
    # 2. Modify Processor Config
    processor.image_processor.image_seq_length = args.vision_tokens
    processor.image_processor.max_soft_tokens = args.vision_tokens

• Chat Templates & Turn Marker Alignment:
  • Chat messages utilize the multimodal list structure:

    [{"role": "user", "content": [{"type": "text", "text": prompt}]}]

  • Role configurations use string "model", never "assistant".
  • Gemma 4 turn templates utilize <|turn>user\n and <|turn>model\n.
  • Ensure correct response masking triggers with train_on_responses_only configured for Gemma 4 turn markers:

    instruction_part = "<|turn>user\n"
    response_part = "<|turn>model\n"

• LiteRT-LM Compiled Execution:
  • Load pre-compiled .litertlm models via litert_lm.Engine.
  • Handle GPU vision fallback context safely to guard execution against FFI initialization failures:

    try:
        engine = litert_lm.Engine(model_path=path, backend=litert_lm.Backend.CPU, vision_backend=litert_lm.Backend.GPU)
    except RuntimeError:
        engine = litert_lm.Engine(model_path=path, backend=litert_lm.Backend.CPU, vision_backend=litert_lm.Backend.CPU)

  • Standard terminal math renders via texicode.pipeline.render_tex using raw context:

    import texicode.pipeline as tp
    rendered = tp.render_tex(latex_string, False, True, "raw", {"fonts": "normal"})

• Sampling protocols: Default sampling for chat loops recommendation is temperature=1.0, top_p=0.95, top_k=64. Logical reasoning and structured OCR mathematical evaluations require deterministic greedy settings: do_sample=False.
• Comments only when the why is non-obvious. Don't narrate the obvious.
• Don't add docs files, READMEs in subdirectories, or planning notes.

Quantitative Evaluation & Normalization

To calculate precise comparative scores, predictions and gold strings must run through strict normalization procedures before validation tests are checked. Baseline evaluations use:

1. Regex-Based Math isolations: Extract exact target calculations under robust patterns. Isolate specific #### <answer> formats, trailing measurements (e.g. weeks, hours, boxes), and format decimal conversions cleanly to avoid exact match rejection due to decimal notation style disparities (refer to examples/eval_gsm8k_automated.py:extract_answer).
2. Visual LaTeX OCR Normalizations: Space formatting characters and common commands must be standardized:
   • Strip all white-spaces.
   • Map shortcuts cleanly: \le(?!q) -> \leq, \ge(?!q) -> \geq, \to -> \rightarrow, \epsilon -> \varepsilon.
   • Standardize braces subscript notation mappings: _([a-zA-Z0-9]|\\[a-zA-Z]+) -> _{\1} (e.g., x_i to x_{i}).
   • Validate transcription consistency mathematically using Exact Match (EM) and Normalized Edit Distance (NED) (based on customized Levenshtein calculations) in examples/eval_vision.py.

Experimental Insights & Hyperparameter Rules

Review validated learning takeaways from past task studies in experiments.md and experimental/README.md before planning updates:

• The Alignment Tax constraint: Applying training scripts over restrictive domains reduces baseline zero-shot capability parameters across overall cognitive outputs. Keep learning targets focused and verify scores relative to zero-shot baselines.
• The Generalization Bottleneck: Model configurations under parameter limits lack capacity elements to formulate customized, brand new semantic syntaxes (such as LRegex logic structures). In these settings, training limits cause overfitting over prompt variations, while baseline English PCRE pre-trained priors dominate evaluation queries. Do not request small adapter updates for large architectural/syntactic structural transitions.
• Low-LR boundaries protecting reasoners: When running semantic categorization tasks over noisy data categories, utilize low learning rates (such as $2e-5$) rather than standard settings ($2e-4$). This shields base weights from incorrect target designations while stabilizing capability boundaries (see experimental/README.md scientific audit).

How to run things

# One-shot inference smoke test (default Gemma 4 E2B)
uv run python examples/inference.py

# GSM8K fine-tune end-to-end in full FP16 precision
uv run python gsm8k-math/finetune_gsm8k.py

# Vision SFT fine-tune over LaTeX OCR (enforces default 2x Alpha rule for rank)
uv run python examples/finetune_vision.py --lora-rank 16 --vision-tokens 280 --output-dir lora_vision

# Greedy eval of a saved adapter
uv run python gsm8k-math/eval_gsm8k.py --adapter gsm8k-math/lora_exp17/

# Automated pipeline: Convert adapter to compiled LiteRT-LM format and execute CPU evaluation
./litert-lm/convert_and_eval.sh

# VLM exact score evaluation (exact score comparisons across N samples)
uv run python examples/eval_vision.py --model lora_vision/ --eval-rows 50

New execution variants and hyperparameter tuning operations should expose CLI commands rather than hardcoding variables into scripts.

Verifying changes

There is no test suite. Verify in this order:

1. Parse check (no GPU needed):

   uv run python -c "import ast; ast.parse(open('gsm8k-math/finetune_gsm8k.py').read()); print('OK')"

2. End-to-end smoke run with reduced settings:

   uv run python gsm8k-math/finetune_gsm8k.py --max-steps 10 --train-rows 100 \
     --output-dir /tmp/lora_smoke

   Expected: BEFORE eval prints, training loss decreases, AFTER eval prints, adapter saved. Total ~2 min on a 16 GB GPU.
3. Reload check: uv run python gsm8k-math/eval_gsm8k.py --adapter /tmp/lora_smoke.

For long training runs, launch via Bash with run_in_background=true and wait for the completion notification. Don't poll. Don't run a parallel Monitor task — there's an observed (not fully diagnosed) cascade where killing a monitor can take down the training process. If you need progress updates, just read the log file periodically.

Training artifacts

After a training run, --output-dir contains adapter_config.json, adapter_model.safetensors, the saved tokenizer, an auto-generated README.md, and checkpoint-N/. See experiments.md for what each file means.

lora_*/ is .gitignored. Never commit training artifacts. Stage files explicitly with git add <file>; never git add -A/..

Don'ts (caught the hard way)

• Don't use temperature=1.0 to evaluate math/reasoning. Sampling noise dominates whatever the model actually learned. Greedy decoding alone took the GSM8K eval from 0/3 to 1/3 on the same adapter.
• Don't truncate eval generation at 256 tokens — GSM8K reasoning chains exceed that. The default in generate() is 512 for a reason.
• Don't recommend bigger LoRA rank or longer training as the first move to improve accuracy on base models without auditing learning dynamics. The base model's underlying capability acts as a ceiling for complex logical extrapolation — see task experiments.
• Don't commit lora_*/ directories or any large weights. They're build output, not source.
• Don't bypass uv (no system pip, no python … directly).
• Don't add --no-verify, --force, or skip-hooks flags to git unless the user asks. Same for any destructive git operation.

Where to find more

• README.md — user-facing setup and getting-started.
• experiments.md — documented hyperparameter sweeps + results table + artifact glossary + full CLI flag reference.
• pyproject.toml — dependency list and Python version range.
• examples/inference.py — minimal FastModel reference; copy its style.
• experimental/README.md — documentation and reproducing guides on custom parser sandbox models.