LAwFTrainer

Fine-tune your LLM using minimal data (down to 1 example) and computing power with the Learning Anchors without Forgetting trainer.

Usage

Installation

Clone this repository, and then:

pip install "trl[peft]"
pip install flash-attn  # Recommended
pip install -e ./lawf-trainer

Dataset format

The LAwFTrainer supports both conversational and instruction dataset formats.

[
  {
    "prompt": [
      {"role": "system", "content": "You are a helpful assistant."},
      {"role": "user", "content": "Replace the background color in App.tsx with #f5f5f5."}
    ],
    "teacher_prompt": [
      {"role": "system", "content": "You are a helpful assistant.\n\nWhen modifying a code file, you need to first review the contents of the code file before deciding how to modify it."},
      {"role": "user", "content": "Replace the background color in App.tsx with #f5f5f5."}
    ],
    "completion": [{"role": "assistant", "content": "<Output generated using teacher_prompt>"}],
    "tools": []
  }
]

The tools field is optional, and its format is identical to the dataset format used by SFTTrainer in TRL.

WARNING: The completion field MUST be generated using the model to fine-tune and teacher_prompt (For OpenAI Python SDK, [completion.choices[0].message.model_dump()] is recommended). Using other models or directly modifying the completion field is pointless.

CLI

Fine-tune using CLI:

lawf-trainer \
  --model_name_or_path qwen/Qwen3-32B \
  --dataset ./path/to/dataset.json \
  --save_dir ./outputs \
  --attn_implementation flash_attention_2

References

Kalle Kujanpää, Harri Valpola & Alexander Ilin (2024).
Efficient Knowledge Injection in LLMs via Self-Distillation.
arXiv preprint arXiv:2412.14964.
https://arxiv.org/abs/2412.14964