Add rollout trace logging with trackio#1360
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements rollout and evaluation trace logging for the Trackio backend, allowing tensor trajectories to be decoded into human-readable traces with associated metadata. Key changes include the addition of the max_rollout_traces_per_step configuration, implementation of trace decoding in StatsLogger, and integration into the RLTrainer training and evaluation loops. Feedback focuses on optimizing performance by moving GPU-to-CPU tensor transfers for input IDs, masks, rewards, and versions outside of the per-sample processing loop to reduce synchronization overhead.
| metadata["reward"] = float(rewards[sample_index].item()) | ||
| if versions is not None: | ||
| sample_versions = ( | ||
| versions[sample_index, :seqlen].detach().cpu().tolist() |
There was a problem hiding this comment.
The versions tensor should also be moved to CPU outside the loop to avoid repeated GPU-to-CPU transfers for each sample.
abidlabs
changed the title
trackio
abidlabs
requested review from
fishcrap,
garrett4wade,
nuzant,
rchardx and
sitabulaixizawaluduo
as code owners
|
Thanks for your contribute! Please run 'pre-commit' before your submit |
|
Ran |
| trackio.Trace( | ||
| messages=[ | ||
| {"role": "user", "content": prompt}, | ||
| {"role": "assistant", "content": completion}, |
There was a problem hiding this comment.
Is there a good way to support multi-turn traces
There was a problem hiding this comment.
I added multi-turn/tool trace support by using structured messages when available and reconstructing user/assistant/tool spans from loss_mask otherwise. Here's what it looks like:
|
Thanks for the review @sitabulaixizawaluduo @PrometheusComing! Addressed all of the comments and reran the pre-commit. All changes have been pushed. |
|
Thanks folks! @sitabulaixizawaluduo ok to merge? |
| space_id: str | None = None | ||
| """HF Space ID for remote dashboard deployment (e.g. "user/my-space"). | ||
| When set, metrics are also pushed to the specified Hugging Face Space.""" | ||
| max_rollout_traces_per_step: int = 32 |
There was a problem hiding this comment.
The default value should be set to disabled, so that the behavior of previous users is not altered.
abidlabs
force-pushed
the
add-trackio-trace-logging
branch
from
feeb891 to
008932a
Compare
|
Thanks @sitabulaixizawaluduo, trace logging is now disabled by default, have updated the docs to reflect that. |
4 participants
Hi folks! This PR adds trace logging via Trackio, the free, local-first experiment tracking library from Hugging Face 🤗
AReaL already has an existing Trackio metrics backend, so this PR extends it to also include logging Traces. specifically I did this:
trackio.Tracerecords whenstats_logger.trackio.modeis enabledstats_logger.trackio.max_rollout_traces_per_stepto cap trace volume per stepHere's what it looks like:
AI assistance was used to prepare this PR.