This repository provides a PyTorch implementation of the self-rewarding language model training framework from MetaAI, along with an implementation of SPIN (Self-Play Fine-Tuning). It is designed for researchers and practitioners looking to fine-tune language models using preference data and self-generated rewards, enabling stronger model performance without external human feedback.

How It Works

The library implements training pipelines for both Self-Rewarding Language Models and SPIN. Self-rewarding models leverage a language model to generate rewards for its own responses, which are then used to fine-tune the model via a Direct Preference Optimization (DPO) approach. SPIN uses self-play to generate preference data, also fine-tuning the model with DPO. The framework is highly modular, allowing users to define custom reward prompts and orchestrate arbitrary sequences of SFT, SPIN, and self-rewarding DPO stages.

Quick Start & Requirements

• Install via pip: pip install self-rewarding-lm-pytorch
• Requires PyTorch.
• Example usage and detailed configuration options for custom reward prompts and fine-tuning pipelines are available in the README.

Highlighted Details

• Implements both Self-Rewarding LM and SPIN training frameworks.
• Supports custom reward prompt templates and reward parsing via regex.
• Allows flexible ordering of fine-tuning stages (SFT, SPIN, Self-Reward DPO, External Reward DPO).
• Integrates with x-transformers for model architecture.

Maintenance & Community

• Sponsored by A16Z Open Source AI Grant Program and Huggingface.
• The author is lucidrains, known for various PyTorch implementations of AI research papers.
• Discussion and development likely occur within the open-source community channels associated with the author.

Licensing & Compatibility

• The repository does not explicitly state a license in the provided README. Users should verify licensing for commercial or closed-source use.

Limitations & Caveats

• The implementation of KV cache handling is noted as a future task.
• An environment flag for auto-clearing checkpoints is also pending.
• The library is primarily focused on research experimentation and may require significant adaptation for production deployment.