Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions docs/source/_toctree.yml
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,11 @@
title: How-to guides
- sections:
- local: transformers_integration
title: Transformers Trainer
title: Running with Transformers
- local: trl_integration
title: Running with TRL
- local: rapidfireai_integration
title: RapidFire AI
title: Running with RapidFire AI
title: Integration
- sections:
- local: api
Expand Down
37 changes: 37 additions & 0 deletions docs/source/trl_integration.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
# TRL Integration

Trackio integrates natively with [TRL](https://github.com/huggingface/trl) so you can log metrics from any TRL trainer (`SFTTrainer`, `DPOTrainer`, `GRPOTrainer`, etc.) with minimal setup. Ensure you have the latest version of `trl` installed (version 1.2.0 or higher).

```python
from datasets import Dataset
from trl import SFTConfig, SFTTrainer

# Create a small fake dataset
prompts = ["The capital of France is", "Hamlet was written by"] * 12
completions = [" Paris.", " Shakespeare."] * 12
dataset = Dataset.from_dict({"prompt": prompts, "completion": completions})

# Train a model using the TRL SFTTrainer API
trainer = SFTTrainer(
model="Qwen/Qwen3-0.6B",
args=SFTConfig(report_to="trackio", run_name="Qwen3-0.6B-sft"),
train_dataset=dataset,
)
trainer.train()
```

## Configuring Project and Space

Set the project and space ID directly in your TRL config (e.g. [`~trl.SFTConfig`], [`~trl.DPOConfig`], [`~trl.GRPOConfig`]):

```python
from trl import SFTConfig

args = SFTConfig(
report_to="trackio",
run_name="my-run",
project="my-project",
trackio_space_id="username/space_id",
)
```

53 changes: 53 additions & 0 deletions examples/trl-integration.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
# /// script
# dependencies = [
# "trackio>=0.23.0",
# "trl>=1.2.0",
# "datasets>=4.4.0",
# "transformers[torch]>=5.0.0rc2",
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure this is needed: trl already depends on datasets, transformers and torch

# "huggingface_hub>=1.0.0",
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — removed. Also dropped datasets and transformers[torch] since trl already pulls them in transitively.

# ]
# ///

import random

from datasets import Dataset
from trl import SFTConfig, SFTTrainer

suffix = random.randint(100000, 999999)
project_name = f"trackio-trl-demo-{suffix}"

prompts = [
{"role": "user", "content": "What is the capital of France?"},
{"role": "user", "content": "Who wrote Hamlet?"},
{"role": "user", "content": "What is 2 + 2?"},
{"role": "user", "content": "What color is the sky?"},
{"role": "user", "content": "Name a primary color."},
{"role": "user", "content": "What is the largest planet?"},
] * 4
completions = [
{"role": "assistant", "content": "Paris."},
{"role": "assistant", "content": "Shakespeare."},
{"role": "assistant", "content": "4."},
{"role": "assistant", "content": "Blue."},
{"role": "assistant", "content": "Red."},
{"role": "assistant", "content": "Jupiter."},
Comment thread
abidlabs marked this conversation as resolved.
Outdated
] * 4
dataset = Dataset.from_dict({"prompt": prompts, "completion": completions})

trainer = SFTTrainer(
model="Qwen/Qwen3-0.6B",
args=SFTConfig(
output_dir="./model_output",
num_train_epochs=1,
per_device_train_batch_size=4,
learning_rate=2e-5,
logging_steps=1,
report_to="trackio",
project=project_name,
),
train_dataset=dataset,
)

trainer.train()

print(f"Run complete. Open with: trackio show --project {project_name}")
Loading