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
| cnt += 1 | ||
| self.eval_rollout.wait(cnt, timeout=None) | ||
| eval_batch = self.eval_rollout.wait(cnt, timeout=None) | ||
| self.stats_logger.log_rollout_traces( |
There was a problem hiding this comment.
The wait method maybe return with None standing for rejected trajectories. Inside _trajectory_to_trackio_traces the first line is trajectory.get("input_ids"), which will AttributeError on any None element. Could you either filter Nones out in log_rollout_traces, or guard at the top of
_trajectory_to_trackio_traces?
There was a problem hiding this comment.
I also agree this should be fixed. It would be best to add test cases for the None trajectory scenario.
There was a problem hiding this comment.
Thanks for catching this, I handled rejected None trajectories from rollout/eval results, updated trajectory type hints, and added tests for rejected trajectories
|
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. |
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.