Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/facebookresearch/LoRe/llms.txt

Use this file to discover all available pages before exploring further.

LoRe has a minimal dependency footprint: a standard Python scientific stack plus PyTorch and Hugging Face Transformers. The only hardware requirement for training is a reasonably modern CUDA GPU; embedding extraction additionally requires enough VRAM to host the Skywork reward model backbone (~16 GB in bfloat16).

Requirements

RequirementMinimum versionNotes
Python3.83.10+ recommended
PyTorch2.3.0Must match your CUDA version
CUDA GPUAny CUDA-capable deviceRequired for embedding generation; strongly recommended for training
Flash Attention 2LatestRequired by prepare.py for RedditTLDR and PersonalLLM

Installing dependencies

Clone the repository, then install all pinned dependencies from requirements.txt:
git clone https://github.com/facebookresearch/LoRe.git
cd LoRe
pip install -r requirements.txt
The full dependency list with pinned versions:
requirements.txt
datasets==2.21.0
matplotlib==3.10.1
numpy==2.2.4
pandas==2.2.3
pydantic==2.11.1
Requests==2.32.3
safetensors==0.5.3
scikit_learn==1.6.1
scipy==1.15.2
torch==2.3.0
torchinfo==1.8.0
transformers==4.46.3
LoRe is not distributed as an installable Python package. All scripts import utils.py by appending the repository root to sys.path at runtime, so you do not need to run pip install -e . or any equivalent.

GPU and Flash Attention 2

Embedding extraction (prepare.py) loads Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 — an 8B parameter model that requires approximately 16 GB VRAM when loaded in bfloat16. This step is mandatory for the RedditTLDR and PersonalLLM datasets and uses attn_implementation="flash_attention_2" for efficient inference.
Flash Attention 2 is not included in requirements.txt because it requires a separate build step. Install it before running prepare.py:
pip install flash-attn --no-build-isolation
This requires a C++ compiler and matching CUDA toolkit headers. See the flash-attn installation guide for troubleshooting.
The train_basis.py scripts also load the reward model backbone to extract the pretrained linear head weights as the reference vector V_final. A GPU is strongly recommended for training but is not strictly required for small K values.

Project structure

After cloning, the repository has the following layout:
LoRe/
├── utils.py                    # Core training, optimization, and evaluation helpers
├── requirements.txt            # Pinned Python dependencies
├── RedditTLDR/                 # Reddit TLDR summarization dataset
│   ├── prepare.py              # Download dataset and extract embeddings
│   ├── train_basis.py          # Train shared reward basis and user weights
│   └── vary_fewshot.py         # Evaluate few-shot personalization
├── PRISM/                      # PRISM multi-turn dialogue dataset
│   ├── prepare.py
│   ├── generate-prism-embeddings.py
│   ├── train_basis.py
│   └── eval_rb2.py             # Evaluate on RewardBench 2
└── PersonalLLM/                # PersonalLLM open-ended response dataset
    ├── prepare.py
    └── train_basis.py
All shared logic lives in utils.py at the root. Dataset-specific scripts import from it using:
import sys, os
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(SCRIPT_DIR))
from utils import *

Verify the installation

After installing dependencies, confirm the core imports work:
import torch
import transformers
import datasets
import sklearn
import scipy

print(f"PyTorch {torch.__version__}")
print(f"Transformers {transformers.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
Expected output on a GPU machine:
PyTorch 2.3.0
Transformers 4.46.3
CUDA available: True
GPU: NVIDIA A100 80GB PCIe
VRAM: 80.0 GB
You can also verify that utils.py is importable from a dataset subdirectory by running the following from, for example, RedditTLDR/:
import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath("."))))
from utils import LoRe, LoRe_regularized, PersonalizeBatch, run, run_regularized
print("LoRe core imports OK")

Environment recommendations

It is recommended to install LoRe dependencies into an isolated environment to avoid conflicts with other PyTorch projects:
python -m venv lore-env
source lore-env/bin/activate   # On Windows: lore-env\Scripts\activate
pip install -r requirements.txt
conda create -n lore python=3.10
conda activate lore
conda install pytorch==2.3.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install -r requirements.txt
Install PyTorch via conda before requirements.txt to ensure CUDA toolkit compatibility.
requirements.txt pins torch==2.3.0, which was built against CUDA 11.8 and 12.1. If your system CUDA version differs, install a compatible PyTorch build from pytorch.org before running pip install -r requirements.txt, and skip the torch line in requirements.txt if needed.

Build docs developers (and LLMs) love