Skip to main content

Diagnose CUDA OOM, NaN loss, LoRA target modules, tokenizer padding, QLoRA, PEFT data, and adapter merge problems.

Project description

PEFT Doctor: LoRA and QLoRA Fine-Tuning Debugger

PEFT Doctor is a pre-flight checker and troubleshooting toolkit for PEFT, LoRA, and QLoRA fine-tuning. It catches the problems that usually waste a training run: CUDA out of memory, NaN loss, risky learning rates, missing tokenizer padding, wrong LoRA target modules, broken prompt formats, bitsandbytes setup issues, and adapter save/load or merge failures.

It is built for the way people actually fine-tune models today: Hugging Face Transformers, PEFT, TRL, bitsandbytes, Google Colab, local CUDA machines, and common Llama, Mistral, Qwen, Gemma, Phi, GPT-2, Falcon, Bloom, and T5-style model families.

The package works in two ways:

  • Use peft-doctor from the terminal before training.
  • Use diagnose_peft(...) inside your training script with real model, tokenizer, peft_config, training_args, and dataset objects.

Problems PEFT Doctor Helps Fix

Developers often find this package while trying to fix one of these PEFT fine-tuning problems:

  • CUDA out of memory during LoRA or QLoRA training
  • QLoRA 4-bit loading problems with bitsandbytes
  • loss=nan, infinite loss, fp16 overflow, or unstable training loss
  • wrong target_modules for Llama, Mistral, Qwen, Gemma, Phi, GPT-2, Falcon, Bloom, or T5
  • tokenizer padding errors such as tokenizer has no pad_token
  • model not learning after PEFT fine-tuning
  • bad output, repeated text, or prompt template mistakes
  • PEFT adapter not saving, loading, or merging correctly
  • PeftModel.from_pretrained adapter loading issues
  • merge_and_unload() problems when exporting a merged LoRA model
  • Colab PEFT setup problems, missing GPU runtime, or broken install cells
  • dataset format problems for instruction tuning, chat templates, SFT, and prompt/completion data
  • labels fully masked with -100, label/input length mismatches, or bad data collators
  • train/eval leakage, duplicate samples, and long rows getting truncated
  • use_cache=True conflicts with gradient checkpointing
  • tokenizer size larger than model embeddings after adding special tokens
  • too many or zero trainable parameters after applying LoRA
  • missing warmup, scheduler, seed, checkpoint retention, or QLoRA optimizer choices
  • slow long-context training that could use Flash Attention

Install

Minimal install:

python -m pip install peft-doctor

Install with the normal fine-tuning stack:

python -m pip install "peft-doctor[ml]"

In Google Colab:

%pip install -U "peft-doctor[ml]"
!peft-doctor env

Use a GPU runtime in Colab before loading a model: Runtime -> Change runtime type -> T4, L4, A100, or another GPU.

Development install from this repository:

git clone https://github.com/awais-akhtar/peft-doctor.git
cd peft-doctor
python -m pip install -e ".[dev,ml]"

Quick Start

Run a pre-flight check from the terminal:

peft-doctor check \
  --model meta-llama/Llama-3-8B \
  --dataset data.jsonl \
  --batch-size 4 \
  --sequence-length 4096 \
  --learning-rate 2e-4

Use it in Python:

from peft_doctor import diagnose_peft

report = diagnose_peft(
    model=model,
    tokenizer=tokenizer,
    peft_config=peft_config,
    training_args=training_args,
    train_dataset=train_dataset,
    sequence_length=2048,
)

print(report.to_markdown())

Generate safe starter configs:

from peft_doctor import create_safe_lora_config, create_safe_bnb_config, create_safe_training_args

peft_config = create_safe_lora_config(model)
bnb_config = create_safe_bnb_config()
training_args = create_safe_training_args()

What It Checks

Area Common problem Typical fix
GPU memory CUDA out of memory Use QLoRA, batch size 1, gradient checkpointing, shorter sequence length
Target modules LoRA attached to the wrong layers Use model-aware targets like q_proj, v_proj, c_attn, or query_key_value
Prompt format Dataset does not teach the response shape Use instruction/response text or a proper chat template
Learning rate Loss spikes or NaN Try 1e-4, 5e-5, bf16, cleaner samples, and label checks
Tokenizer Padding crash during batching Set tokenizer.pad_token = tokenizer.eos_token when appropriate
Evaluation Eval OOM after training works Disable eval or use a tiny eval batch
Adapter flow Adapter not found after training Use model.save_pretrained() and PeftModel.from_pretrained()
Data quality Duplicate rows, split leakage, masked labels Deduplicate, fix labels, separate train/eval
Model state use_cache, embeddings, trainable params Disable cache, resize embeddings, verify LoRA trainables
Trainer config Missing warmup, seed, scheduler, checkpoint limit Add stable defaults before long runs

Troubleshooting Recipes

For a longer problem-by-problem guide, see docs/troubleshooting.md.

Fix CUDA Out of Memory in PEFT or QLoRA

peft-doctor check \
  --model meta-llama/Llama-3-8B \
  --dataset train.jsonl \
  --batch-size 4 \
  --sequence-length 4096 \
  --learning-rate 2e-4

If the report warns about memory, start with:

training_args = {
    "per_device_train_batch_size": 1,
    "gradient_accumulation_steps": 8,
    "gradient_checkpointing": True,
    "bf16": True,
}

For QLoRA:

from peft_doctor import create_safe_bnb_config

bnb_config = create_safe_bnb_config()

Fix Wrong LoRA Target Modules

peft-doctor targets --model meta-llama/Llama-3-8B
peft-doctor targets --model Qwen/Qwen2.5-7B
peft-doctor targets --family gpt2

Fix NaN Loss in LoRA Fine-Tuning

peft-doctor scan-log trainer_log.jsonl

Common fixes are lower learning rate, bf16 instead of fp16, cleaner samples, valid labels, gradient clipping, and shorter sequences while debugging.

Fix Tokenizer Padding Errors

tokenizer.pad_token = tokenizer.eos_token

PEFT Doctor warns when a causal language model tokenizer has no pad token.

Merge a LoRA Adapter Into the Base Model

peft-doctor adapter-check \
  --base-model meta-llama/Llama-2-7b-hf \
  --adapter your-user/your-lora-adapter \
  --output-dir merged-model

peft-doctor merge-adapter \
  --base-model meta-llama/Llama-2-7b-hf \
  --adapter your-user/your-lora-adapter \
  --output-dir merged-model \
  --dtype fp16

Commands

Full command reference with examples: docs/commands.md.

peft-doctor check

Runs the main pre-flight check.

peft-doctor check --model meta-llama/Llama-3-8B --dataset data.jsonl

Useful options:

peft-doctor check \
  --model Qwen/Qwen2.5-7B \
  --dataset train.jsonl \
  --batch-size 2 \
  --grad-accum 8 \
  --sequence-length 2048 \
  --learning-rate 2e-4 \
  --load-in-4bit \
  --bf16 \
  --gradient-checkpointing

Machine-readable output:

peft-doctor check --model mistralai/Mistral-7B-v0.1 --dataset train.jsonl --output json

Markdown output for issues or pull requests:

peft-doctor check --model gpt2 --dataset train.jsonl --output markdown

peft-doctor targets

Recommends LoRA target_modules.

peft-doctor targets --model meta-llama/Llama-3-8B
peft-doctor targets --family gpt2

Print as JSON:

peft-doctor targets --family qwen --output json

peft-doctor safe-config

Prints a safe LoRA or QLoRA starter config.

peft-doctor safe-config --model meta-llama/Llama-3-8B

Only LoRA:

peft-doctor safe-config --family gpt2 --no-qlora

JSON:

peft-doctor safe-config --family llama --output json

peft-doctor inspect-dataset

Checks a local .json, .jsonl, .csv, or .txt dataset sample.

peft-doctor inspect-dataset data.jsonl

The command looks for common training shapes:

  • messages chat rows with role and content
  • instruction and response style columns
  • single text rows containing instruction/response markers
  • pre-tokenized input_ids and labels

peft-doctor scan-log

Scans a training log for NaN, infinity, CUDA OOM, overflow, and unstable loss jumps.

peft-doctor scan-log trainer_log.jsonl
peft-doctor scan-log run.log --output markdown

peft-doctor adapter-check

Checks a LoRA adapter merge plan without loading the full model.

peft-doctor adapter-check \
  --base-model meta-llama/Llama-2-7b-hf \
  --adapter awaisakhtar/llama-2-7b-summarization-finetuned-on-xsum-lora \
  --output-dir merged-llama

peft-doctor merge-adapter

Merges a PEFT LoRA adapter into the base model and saves a normal Transformers model.

peft-doctor merge-adapter \
  --base-model meta-llama/Llama-2-7b-hf \
  --adapter awaisakhtar/llama-2-7b-summarization-finetuned-on-xsum-lora \
  --output-dir Llama-2-7b-summarization-finetuned-on-xsum \
  --dtype fp16

Push the merged model and tokenizer to the Hugging Face Hub:

huggingface-cli login

peft-doctor merge-adapter \
  --base-model meta-llama/Llama-2-7b-hf \
  --adapter your-user/your-lora-adapter \
  --output-dir merged-model \
  --push-to-hub \
  --hub-model-id your-user/merged-model \
  --dtype fp16

For Colab or private/gated models, store your own Hugging Face token as a secret named HF_TOKEN and read it from the notebook environment. Do not paste access tokens into notebooks, scripts, shell history, or GitHub issues.

For final exports, do not merge from a 4-bit or 8-bit loaded model unless you know your PEFT/Transformers versions support it. The safest export path is fp16, bf16, or fp32, then save_pretrained(..., safe_serialization=True).

peft-doctor scan-notebook

Scans a notebook for common PEFT and Colab mistakes, including exposed Hugging Face tokens.

peft-doctor scan-notebook model_merge.ipynb

peft-doctor env

Checks the local Python, CUDA, and fine-tuning package stack.

peft-doctor env
peft-doctor env --output json

This is especially useful in Colab because many setup problems come from the notebook runtime, not the training script.

peft-doctor colab

Prints a notebook-friendly setup cell.

peft-doctor colab

peft-doctor version

Prints the installed version.

peft-doctor version

Python API

Diagnose a training setup

from peft_doctor import diagnose_peft

report = diagnose_peft(
    model=model,
    tokenizer=tokenizer,
    peft_config=peft_config,
    training_args={
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": 8,
        "learning_rate": 2e-4,
        "num_train_epochs": 3,
        "bf16": True,
        "gradient_checkpointing": True,
    },
    train_dataset=train_dataset,
)

if report.has_errors:
    raise RuntimeError(report.to_markdown())

Generate target modules

from peft_doctor import recommend_target_modules

targets = recommend_target_modules(model_name="meta-llama/Llama-3-8B")

Generate safe LoRA and QLoRA configs

from peft_doctor import create_safe_bnb_config, create_safe_lora_config

peft_config = create_safe_lora_config(model, r=16, lora_alpha=32)
bnb_config = create_safe_bnb_config()

When peft, transformers, and torch are installed, these helpers return real LoraConfig and BitsAndBytesConfig objects. Without those packages, they return plain dictionaries so you can still inspect the recommendation.

Guard training logs

from peft_doctor import NanLossGuard

guard = NanLossGuard()

for log in trainer_state_log_history:
    issues = guard.update(log)
    for issue in issues:
        print(issue.title, issue.fix)

Merge a LoRA adapter

from peft_doctor import merge_lora_adapter

result = merge_lora_adapter(
    base_model="meta-llama/Llama-2-7b-hf",
    adapter="your-user/your-lora-adapter",
    output_dir="merged-model",
    torch_dtype="fp16",
)

print(result.to_dict())

Colab Notebook Pattern

%pip install -U "peft-doctor[ml]"

from peft_doctor import diagnose_peft, create_safe_lora_config, create_safe_bnb_config

peft_config = create_safe_lora_config(model_name="meta-llama/Llama-3-8B")
bnb_config = create_safe_bnb_config()

report = diagnose_peft(
    model_name="meta-llama/Llama-3-8B",
    peft_config=peft_config,
    training_args={
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": 8,
        "learning_rate": 2e-4,
        "bf16": True,
        "gradient_checkpointing": True,
        "load_in_4bit": True,
    },
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)

print(report.to_markdown())

Dependency Note

PEFT Doctor uses open-source Python packages from the normal PyData and Hugging Face fine-tuning stack: torch, transformers, peft, datasets, accelerate, rich, typer, and related optional packages. Model weights and datasets can have their own licenses, so always check the license of the model and data you fine-tune.

Common Safe Config

from peft import LoraConfig

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
)
from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

Publishing

The repository includes GitHub Actions for CI and PyPI publishing.

  1. Push this repository to awais-akhtar/peft-doctor.
  2. On PyPI, create a trusted publisher for package peft-doctor:
    • owner: awais-akhtar
    • repository: peft-doctor
    • workflow: publish.yml
    • environment: pypi
  3. On TestPyPI, create the same trusted publisher with environment testpypi.
  4. Push a version tag to publish to PyPI:
git tag v0.2.0
git push origin v0.2.0

Manual TestPyPI publishing is available from the Publish Python Package workflow in GitHub Actions.

Project Status

This is alpha software. The checks are deliberately conservative: the package should warn early, explain the reason, and give a fix that a developer can actually try.

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

peft_doctor-0.2.0.tar.gz (42.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

peft_doctor-0.2.0-py3-none-any.whl (40.4 kB view details)

Uploaded Python 3

File details

Details for the file peft_doctor-0.2.0.tar.gz.

File metadata

  • Download URL: peft_doctor-0.2.0.tar.gz
  • Upload date:
  • Size: 42.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for peft_doctor-0.2.0.tar.gz
Algorithm Hash digest
SHA256 7d99a7f9d630409e5b430445ca0ba0ce72d36c7262a51e47eed419840757541a
MD5 e86efcf5efab3bada14e5d9b2d1c061f
BLAKE2b-256 baa316ca49e4432e4d62b8cc4416621d510207c0e9d33b8e72b5943f283274f1

See more details on using hashes here.

File details

Details for the file peft_doctor-0.2.0-py3-none-any.whl.

File metadata

  • Download URL: peft_doctor-0.2.0-py3-none-any.whl
  • Upload date:
  • Size: 40.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for peft_doctor-0.2.0-py3-none-any.whl
Algorithm Hash digest
SHA256 6c7aba4fb6df3dbb2d1c3cde1c8bc43773ad8214b1f450b17bff2739e6f827d4
MD5 88f6920613919028d8579ff6f13ac7c4
BLAKE2b-256 8b97e8f287707daaa0c4d43e5682242fc49d97c87eb62fa0da575e1080eb7eee

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page