-
Notifications
You must be signed in to change notification settings - Fork 112
docs: add TRL integration guide #537
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
669fd89
docs: add TRL integration guide and rename integration sidebar titles
abidlabs f40031c
Apply suggestion from @abidlabs
abidlabs fbf82b0
Update docs/source/trl_integration.md
abidlabs 3ba9c47
Update examples/trl-integration.py
abidlabs 4c1ab0d
Update examples/trl-integration.py
abidlabs d852324
Update examples/trl-integration.py
abidlabs adaeee3
trim trl example deps to trackio + trl (transitive deps come from trl)
abidlabs cd1747e
Apply suggestion from @qgallouedec
abidlabs File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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", | ||
| ) | ||
| ``` | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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", | ||
| # "huggingface_hub>=1.0.0", | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good catch — removed. Also dropped |
||
| # ] | ||
| # /// | ||
|
|
||
| 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."}, | ||
|
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}") | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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