RLFusion

Minimalist post-training utilities for LLMs with a focus on clarity and ease of learning.

Features

• SFT with role-aware masking (mask non-assistant tokens, train assistant tokens)
• RLVR with GRPO
• On-policy distillation (reverse KL to a teacher on student samples)
• Simple EnvBase for prompts + answers

Install

From the repo root:

uv pip install -e .

If you don't have uv installed:

python -m pip install uv

Optional extras:

# Repo-standard HF CUDA training stack:
# Liger + bitsandbytes/AdamW8bit only
uv pip install -e ".[gpu-train]"

# Optional HF FlashAttention backend (Linux; may build from source)
uv pip install -e ".[flash]" --no-build-isolation

# vLLM support (Linux + CUDA)
uv pip install -e ".[vllm]"

# Full validated stack without FlashAttention
uv sync --extra gpu-train --extra vllm --extra dev --extra test

# Full validated stack with FlashAttention added explicitly
uv sync --extra gpu-train --extra flash --extra vllm --extra dev --extra test

HF CUDA Training Stack (Optional)

If you have a CUDA-enabled GPU and want the repo-standard HF training stack:

uv sync --extra gpu-train --extra vllm --extra dev --extra test

This installs:

• liger-kernel so HF trainers automatically use the Liger model wrappers when available
• bitsandbytes for memory-efficient optimizers such as AdamW8bit

If you also want HF FlashAttention, install it separately:

uv sync --extra gpu-train --extra flash --extra vllm --extra dev --extra test

When flash-attn is installed, RLFusion selects the fastest supported HF attention backend by default on CUDA:

• FlashAttention-3 on Hopper GPUs when a compatible FA3 backend is available
• FlashAttention-2 on Ampere/Ada/Hopper GPUs when a compatible FA2 backend is available
• PyTorch SDPA otherwise

FlashAttention-3 is Hopper-only and requires a compatible FA3 installation. If no compatible FlashAttention backend is installed, the trainers automatically fall back to PyTorch SDPA. When liger-kernel is installed, the HF trainers also switch to the Liger-backed model classes automatically.

Optional FlashAttention caveat:

• The repo-standard vllm 0.17.x stack pins torch 2.10.x.
• Upstream flash-attn 2.8.3 does not publish a matching wheel for every torch 2.10 platform, so Linux installs may compile flash-attn from source on first install.
• On single-architecture Ampere/Ada hosts, setting FLASH_ATTN_CUDA_ARCHS=80 before install can reduce source-build time and memory pressure.

You can override the automatic choice with:

export RLFUSION_ATTN_IMPLEMENTATION=sdpa

vLLM (Optional)

To use vLLM:

uv pip install -e ".[vllm]"

Set engine="vllm" in the Evaluator. GRPO and on-policy distillation default to colocated vLLM rollouts on CUDA and fall back to HF generation on non-GPU devices; pass use_vllm=False to force HF generation, or use_vllm=True to require vLLM explicitly. Optional vllm_args are forwarded to vllm.LLM.

Notes:

• The repo standardizes on vllm 0.17.x on Linux, with a matching torch 2.10.x / torchaudio 2.10.x / torchvision 0.25.x stack from the vllm extra.

• For Liger + AdamW8bit on the same stack, install the gpu-train extra:

  uv sync --extra gpu-train --extra vllm --extra dev --extra test

• Add --extra flash only if you explicitly want HF FlashAttention on top of that stack.

• vllm_enable_sleep=True automatically enables the underlying vLLM sleep mode.

• When combining use_accelerate=True with use_vllm=True, keep tensor_parallel_size=1 and pipeline_parallel_size=1 so each trainer process owns exactly one local vLLM engine. Use multi-GPU vLLM only when not launching the trainer with Accelerate.

vLLM now uses its own default attention-backend auto-selection unless you override it yourself. If you need to force a backend, use:

export VLLM_ATTENTION_BACKEND=FLASH_ATTN

To request FlashAttention 4 explicitly on supported Blackwell GPUs, pass an attention_config through vllm_args:

trainer = GRPOTrainer(
    ...,
    vllm_args={
        "gpu_memory_utilization": 0.5,
        "attention_config": {
            "backend": "FLASH_ATTN",
            "flash_attn_version": 4,
        },
    },
)

Notes:

• FA4 is a vLLM 0.17.x feature exposed through the FLASH_ATTN backend.
• In practice this matters on Blackwell-class GPUs. Hopper prefers FA3; Ampere and Ada stay on FA2.
• If you leave attention_config unset, vLLM keeps its own backend/version auto-selection.

Dev Setup

uv pip install -e ".[dev,test]"

Quickstart

Minimal, inline examples for each trainer are included below.

Distributed Training (Accelerate)

All trainers support multi-GPU distributed training via Hugging Face Accelerate.

1. Set use_accelerate=True on SFTTrainer, GRPOTrainer, or OnPolicyDistillationTrainer.
2. Launch your script with Accelerate:

accelerate launch --num_processes 2 your_script.py

Notes:

• batch_size is per-process; effective batch size is batch_size * num_processes.
• Logging, evaluation, and checkpoint saving are handled on the main process (rank 0).
• Reward-based evaluation (eval_steps + eval_dataset) runs on the main process.

Core Concepts

Environment

EnvBase represents a single sample with a chat-style prompt and an optional answer.

from rlfusion.envs import EnvBase

env = EnvBase(
    prompt=[
        {"role": "system", "content": "Be concise."},
        {"role": "user", "content": "What is 2 + 2?"},
    ],
    answer="4",
)

Trainers

SFT

SFTTrainer consumes (prompt, answer) and masks non-assistant tokens while training assistant tokens. Eval supports two modes:

• Default (eval_sample_completions=False): logs token-level ce_loss and perplexity from labeled assistant tokens.
• Reward sampling (eval_sample_completions=True): generates completions and logs reward metrics. This mode requires eval_dataset environments with get_reward.

from torch.utils.data import Dataset

from rlfusion.envs import EnvBase
from rlfusion.trainers.sft_trainer import SFTTrainer


class ToySFTDataset(Dataset):
    def __init__(self) -> None:
        self.samples = [
            EnvBase(
                prompt=[
                    {"role": "system", "content": "Be concise."},
                    {"role": "user", "content": "What is 2 + 2?"},
                ],
                answer="4",
            )
        ]

    def __getitem__(self, index: int) -> EnvBase:
        return self.samples[index]

    def __len__(self) -> int:
        return len(self.samples)


trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=ToySFTDataset(),
    num_steps=2,
    batch_size=1,
    saving_steps=2,
    logging_steps=1,
)
trainer.train()

For a real SFT corpus, you can use UltraChat:

from rlfusion.datasets import UltraChatSFTDataset

train_dataset = UltraChatSFTDataset(train=True, max_samples=10_000, seed=42)
eval_dataset = UltraChatSFTDataset(train=False, max_samples=1_000, seed=42)

RLVR (GRPO)

GRPOTrainer samples completions from the model, computes rewards via the environment, and optimizes a GRPO objective. To evaluate during training, set eval_steps and pass an eval_dataset of environments with get_reward.

from rlfusion.datasets import MathDataset
from rlfusion.trainers.grpo_trainer import GRPOTrainer

train_dataset = MathDataset(num_samples=200, min_val=0, max_val=50, operand="add")
eval_dataset = MathDataset(num_samples=50, min_val=0, max_val=50, operand="add")

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=train_dataset,
    num_steps=2,
    saving_steps=2,
    logging_steps=1,
    eval_steps=1,
    eval_dataset=eval_dataset,
    group_size=2,
    ppo_steps=1,
    max_new_tokens=64,
    vllm_args={"gpu_memory_utilization": 0.5},
)
trainer.train()

On-policy Distillation

OnPolicyDistillationTrainer follows the on-policy distillation recipe:

• sample trajectories from the student policy
• score sampled completion tokens under student and teacher
• optimize a PPO-style objective with token-level advantage logp_teacher - logp_student_old

This is distillation (no external RL reward in the loss). Reward logging during train/eval is optional and only used for monitoring.

from rlfusion.datasets import MathDataset
from rlfusion.trainers.onpolicy_distillation_trainer import OnPolicyDistillationTrainer

train_dataset = MathDataset(num_samples=200, min_val=0, max_val=50, operand="add")
eval_dataset = MathDataset(num_samples=50, min_val=0, max_val=50, operand="add")

trainer = OnPolicyDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=train_dataset,
    num_steps=2,
    saving_steps=2,
    logging_steps=1,
    eval_steps=1,
    eval_dataset=eval_dataset,
    sampling_temperature=0.7,
    generation_args={"top_p": 0.9},
    ppo_steps=1,
    clip_eps=0.2,
    max_new_tokens=64,
    max_grad_norm=1.0,
    vllm_args={"gpu_memory_utilization": 0.5},
)
trainer.train()

Notes:

• For reward metrics, eval_dataset environments should implement get_reward, and samples should have answer populated.
• In trainer.test(...), num_batches must be >= 1 when set.
• In GRPOTrainer.test(...) and OnPolicyDistillationTrainer.test(...), eval_temperature must be > 0 when set.
• Do not pass temperature via generation_args; use sampling_temperature during training/evaluation, and eval_temperature for RL trainer eval overrides.

Training Guides

• SFT: training_guide/sft.md
• GRPO: training_guide/grpo.md
• On-policy distillation: training_guide/onpolicy_distillation.md
• Reasoning pipeline (dataset-specific): examples/reasoning/README.md
• Qwen3 reasoning SFT on johnolafenwa/reasoning-sft: examples/sft_qwen3_4b_instruct_2507_reasoning.py
• Qwen3-8B-Base reasoning SFT on johnolafenwa/highschool-math-reasoning: examples/reasoning/reasoning_sft_train.py
• Qwen3 reasoning GRPO from that SFT checkpoint: examples/grpo_qwen3_4b_instruct_2507.py
• Qwen3 base-model inference probe: examples/qwen3_4b_instruct_2507_inference.py

Testing

uv run pytest

Linting

uv run ruff check .

Type Checking

uv run ty check src tests

If you prefer not to install it, you can run:

uvx ty check src tests

Weights & Biases Logging

To enable W&B logging, pass enable_wandb=True to a trainer and optionally set:

• wandb_project (default: sft, grpo, or onpolicy_distill)
• wandb_run_name

Login with uv:

uv run wandb login

Or non-interactive:

uv run wandb login $WANDB_API_KEY

Example:

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
    enable_wandb=True,
    wandb_project="rlfusion",
    wandb_run_name="sft-demo",
)

Build a Wheel

uv pip install --upgrade build
uv run python -m build --wheel

Artifacts land in ./dist.

Citation

If you use RLFusion in your work, please cite:

@software{rlfusion,
  title = {RLFusion},
  author = {Olafenwa, John},
  year = {2025},
  url = {https://github.com/johnolafenwa/rlfusion}
}