preft

Prefill-only finetuning for efficient inference.
Apply parameter-efficient adapters during prefill only; discard them for decode.

At serving time this gives substantially higher throughput than standard PEFTs when many user-specific adapters share one base model — see our paper for the throughput / accuracy tradeoff curves.

preft supports

• two adapter mounts: LoRA (on q_proj / v_proj / …) and ReFT (residual-stream LoReFT / DiReFT)
• a peft-style API: PreftConfig, get_preft_model(...), PreftModel.save_pretrained / from_pretrained
• prefill-only gating via the position kwarg ("prefill", "all", "first", "last")
• accelerated multi-adapter serving and weight sync on a forked vLLM — required only for the vLLM path; library training works on stock HuggingFace
• reference scripts for SFT (train_sft.py), GRPO RL (train_rl.py), and an eval harness (eval.py)
• monkey-patches that let GRPO training run against upstream TRL on PyPI

---

Contents

• Install
  • Default mode
  • vLLM mode
• Quick start
• Adapter mounts
• Position semantics
• vLLM integration
• Citation
• License

---

Install

Two modes, picked by whether you want vLLM-accelerated inference:

Mode	pip	uv	What you get	Platform
Default	pip install -e .	uv sync	Library + SFT + RL (HF generation) + eval	CPU, macOS, Linux+CPU, Linux+GPU
vLLM	pip install -e ".[vllm]" then bash setup.sh	bash setup.sh (calls uv sync --extra vllm)	Above + vllm, flash-attn, liger-kernel	Linux + CUDA only

Pick default to use the library on your own or train with stock HuggingFace generation (slower rollouts but works anywhere). Pick vLLM for production RL training with accelerated rollouts and ReFT weight sync.

uv vs pip. Both treat extras as opt-in. uv sync alone installs only core (no vllm, no flash-attn); uv sync --extra vllm installs the vLLM stack. setup.sh runs uv sync --extra vllm and then overlays the forked vLLM, so bash setup.sh is the one-shot for vLLM mode.

Default mode (no vLLM)

git clone <this-repo> && cd preft

# pip
pip install -e .

# OR uv
uv sync                   # creates ./.venv automatically

Confirm:

python -c "import preft; print(preft.PreftConfig(layer_indices=[0], low_rank_dim=4, hidden_size=16))"
# under uv:
uv run --no-sync python -c "import preft; print(preft.PreftConfig(layer_indices=[0], low_rank_dim=4, hidden_size=16))"

torch==2.8.0 is pinned identically in both modes — pip / uv pull the right wheel for your platform (CPU build on macOS, CUDA build on Linux from the default PyPI index). Override with --index-url https://download.pytorch.org/whl/cpu for explicit CPU on Linux.

You can now:

• import preft and use PreftConfig / get_preft_model / PreftModel
• run train_sft.py (drops to HF Trainer)
• run train_rl.py without --use_vllm (HF generation backend — slower but no CUDA needed)
• run eval.py (HF backend by default)

vLLM mode

git clone <this-repo> && cd preft
bash setup.sh

setup.sh performs five steps:

1. uv sync --extra vllm — install core + the [vllm] extra (skipping flash-attn)
2. uv lock --upgrade-package trl — bump TRL to latest
3. FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE uv pip install flash-attn
4. Clone pyvene-ai/vllm (branch dev-v0.11.0) and rsync *.py over the installed vLLM, preserving compiled .so extensions
5. Run configs/check_vllm_fork.py to verify the fork hooks (set_reft_spec, sync_reft_weights, refresh_reft_caches, per-layer reft_adapter*) landed

Once step 5 reports the patches present, the --use_vllm flag and model.build_vllm(...) are available:

uv run --no-sync python train_rl.py \
    --mode loreft --reft_rank 4 \
    --dataset gsm8k --model Qwen/Qwen2.5-0.5B \
    --use_vllm --vllm_mode colocate \
    --batch_size 16 --num_completions 4 --lr 5e-5 \
    --output_dir ./outputs/rl-test

Prefer pip over uv?

pip install -e ".[vllm]"
# then manually run steps 3–5 from setup.sh

Minimal SFT (either mode)

python train_sft.py \
    --mode loreft --rank 4 \
    --model meta-llama/Llama-3.2-1B-Instruct \
    --dataset allenai/tulu-3-sft-mixture \
    --lr 5e-4 --batch_size 2 \
    --output_dir ./outputs/sft-test

train_sft.py --help and train_rl.py --help list every flag. Output directories contain preft_config.json + adapter_model.safetensors.

---

Quick start

import torch
from transformers import AutoModelForCausalLM
from preft import PreftConfig, get_preft_model, PreftModel

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B",
    torch_dtype=torch.bfloat16,
)

config = PreftConfig(
    layer_indices=[8, 16, 24],
    adapter_type="loreft",     # "loreft" or "direft"
    low_rank_dim=8,
    position="prefill",        # paper's default
)
model = get_preft_model(base, config)
model.print_trainable_parameters()

# train however you'd train any HF model:
#   - HuggingFace Trainer
#   - TRL GRPOTrainer
#   - accelerate
#   - raw loop

model.save_pretrained("./my-adapter")

# load later
base2 = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B", torch_dtype=torch.bfloat16,
)
loaded = PreftModel.from_pretrained(base2, "./my-adapter")

---

Adapter mounts

Mount	Config	Saved as	When to pick
ReFT	adapter_type="loreft" or "direft" use_residual=True (default)	preft_config.json + adapter_model.safetensors	residual-stream intervention; works at any layer index
LoRA	use_residual=False use_lora=True lora_target_modules=[…]	peft-format adapter_config.json + adapter_model.safetensors (vLLM-loadable)	weight-level low-rank update on specific projections

ReFT and LoRA are mutually exclusive within one PreftConfig — pick one.

---

Position semantics

position	Where the adapter fires
"prefill"	every prompt token (the paper's prefill-only mode)
"all"	every token, training and decode
"last"	last prompt token only (SFT-style)
"first"	first prompt token only

At decode time, "prefill" / "last" / "first" are no-ops — the adapter costs zero compute beyond the first forward.

---

vLLM integration

After installing vLLM mode (bash setup.sh):

llm = model.build_vllm("Qwen/Qwen2.5-0.5B", gpu_memory_utilization=0.3)
outputs = llm.generate(["What is 2+2?"])

build_vllm exports a per-layer adapter blueprint, brackets the underlying LLM(...) constructor with vllm.reft.set_reft_spec / clear_reft_spec, and pushes adapter weights into the worker via collective_rpc("sync_reft_weights", ...). Subsequent training steps stay in sync via refresh_reft_caches.

---

Citation

@misc{lanpouthakoun2026preftprefillonlyfinetuningefficient,
      title={PreFT: Prefill-only finetuning for efficient inference}, 
      author={Andrew Lanpouthakoun and Aryaman Arora and Zhengxuan Wu and Dhruv Pai and Ben Keigwin and Dan Jurafsky and Christopher Potts},
      year={2026},
      eprint={2605.14217},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.14217}, 
}

---

License

Apache 2.0. See LICENSE.