TRLoom

TRLoom weaves a single YAML config into an end-to-end Hugging Face TRL fine-tuning job.

Configure the model, dataset, trainer, Weights & Biases, and optional Modal GPU execution — then run one command.

Docs: https://saqlain2204.github.io/trloom/

Features

• TRL-native — discovers trainers/configs from your installed TRL version (SFT, DPO, GRPO, KTO, Reward, RLOO, and experimental methods)
• YAML-first — model, dataset, training args, W&B, and Modal all live in one file
• Datasets — Hugging Face Hub, local files (json/jsonl/csv/parquet/…), saved datasets directories, and mixtures
• W&B — enable and configure logging entirely from YAML
• Modal — launch the same YAML on remote GPUs with volume-backed outputs
• CLI + Python API — trloom run config.yaml or FineTuneJob.from_yaml(...)

Installation

Requires Python 3.10+ and a working TRL / PyTorch environment for actual training.

One-command bootstrap (any OS)

Creates .venv if needed, then installs TRLoom:

# Local only (no Modal auth)
python scripts/bootstrap.py

# Local + Modal extra + `modal setup`
python scripts/bootstrap.py --modal

Optional flags: --wandb, --dev, --all, --skip-modal-setup.

Then activate the venv:

# macOS / Linux
source .venv/bin/activate

# Windows PowerShell
.venv\Scripts\Activate.ps1

Manual install

pip install -e .

# Optional extras
pip install -e ".[wandb]"
pip install -e ".[modal]"
pip install -e ".[docs]"
pip install -e ".[all]"
pip install -e ".[dev]"

Quickstart

1. Write a config

# sft.yaml
method: sft

model:
  model_name_or_path: Qwen/Qwen2.5-0.5B-Instruct
  use_peft: true
  lora_r: 16
  lora_alpha: 32

dataset:
  path: trl-lib/Capybara
  train_split: train

training:
  output_dir: ./outputs/sft
  learning_rate: 2.0e-4
  num_train_epochs: 1
  per_device_train_batch_size: 2
  gradient_accumulation_steps: 4
  report_to: none

wandb:
  enabled: false

modal:
  enabled: false

2. Run

trloom validate sft.yaml
trloom run sft.yaml

Or from Python:

from trloom import FineTuneJob, run_from_yaml

# One-liner
run_from_yaml("sft.yaml")

# Or step through the API
job = FineTuneJob.from_yaml("sft.yaml")
job.run()

Documentation

Full guides and API reference:

• Online: saqlain2204.github.io/trloom
• Local preview:

pip install -e ".[docs]"
mkdocs serve

Coverage includes configuration reference, datasets, W&B, Modal, CLI, and the Python API.

Configuration reference

Section	Purpose
method	TRL method key: sft, dpo, grpo, kto, reward, rloo, …
model	Model id + PEFT/quantization (aligned with TRL ModelConfig)
dataset	Hub repo, local path, or datasets: mixture
training	Forwarded to the TRL *Config class (SFTConfig, DPOConfig, …)
wandb	Weights & Biases project/entity/tags/mode
modal	Remote GPU execution on Modal
reward_funcs	Names or import paths for GRPO/RLOO-style rewards
trainer_kwargs	Extra kwargs passed to the Trainer constructor
push_to_hub / hub_model_id	Optional Hub upload after training

List methods available in your environment:

trloom methods

Dataset examples

Hub

dataset:
  path: trl-lib/Capybara
  train_split: train

Local JSONL

dataset:
  path: ./data/train.jsonl
  train_split: train

Mixture

dataset:
  train_split: train
  datasets:
    - path: stanfordnlp/imdb
      split: train
      weight: 0.5
    - path: ./data/extra.jsonl
      weight: 0.5

Weights & Biases

training:
  report_to: wandb   # optional; set automatically when wandb.enabled is true

wandb:
  enabled: true
  project: my-project
  entity: my-team
  run_name: qwen-sft-01
  tags: [sft, lora]
  mode: online       # online | offline | disabled

Modal

1. Install and authenticate: pip install 'trloom[modal]' && modal setup
2. Create secrets named in the config (default: huggingface, wandb)
3. Set modal.enabled: true (or pass --modal)

trloom run examples/grpo_modal.yaml --modal
# or generate a standalone script
trloom modal-script examples/grpo_modal.yaml -o run_modal.py
modal run run_modal.py

Python API

from trloom import FineTuneJob, available_methods, load_config

print(available_methods())

config = load_config("sft.yaml")
job = FineTuneJob(config)
trainer = job.build()   # construct TRL trainer
job.run()               # train + save (+ optional Hub push)

Method	Description
FineTuneJob.from_yaml(path)	Load YAML into a job
FineTuneJob.from_dict(data)	Load an in-memory config
job.build()	Construct the TRL trainer
job.train() / job.run()	Run training end-to-end
run_from_yaml(path)	Load + run (honors modal.enabled)
load_config(path)	Validate and return FineTuneConfig
available_methods()	List TRL methods for this install

Examples

See the examples/ directory:

• modal_smoke/ — complete end-to-end Modal walkthrough (tiny GPT-2, 3 steps, T4)
• sft_hub.yaml — SFT from the Hub
• sft_local.yaml — SFT from local JSONL
• dpo_wandb.yaml — DPO with W&B
• grpo_modal.yaml — GRPO on Modal

Development

pip install -e ".[dev]"
pytest

# Documentation
pip install -e ".[docs]"
mkdocs serve
mkdocs build --strict

License

Apache-2.0