Motivation
Spin-off from the Post-training efforts RFC (#3948) specifically for Workstream 2 (TRL interop).
Currently, many teams use Hugging Face trl (e.g., PPOTrainer, GRPOTrainer) but face bottlenecks with rollout storage and sampling at scale. Conversely, teams using TorchRL's highly efficient collectors and recipes often want to utilize the vast ecosystem of Hugging Face Reward Models, but writing custom wrappers is tedious.
Nobody adopts a full stack at once. Providing native interoperability between torchrl and trl lowers the barrier to entry, allowing users to adopt TorchRL's high-performance components piece by piece into their existing training loops.
Solution
I recommend introducing a new module (e.g., torchrl.modules.llm.trl_interop) containing thin adapters so that users can seamlessly mix and match torchrl and trl components.
Specifically, I want to implement:
1. TorchRLBufferDataset (TorchRL -> TRL)
A wrapper around TorchRL's ReplayBuffer that implements the standard PyTorch IterableDataset / Hugging Face Dataset interface. This allows trl trainers to directly consume token/trajectory data generated by TorchRL's collectors.
2. HFRewardModelWrapper (TRL -> TorchRL)
A TensorDictModule adapter that wraps a Hugging Face reward model (like those trained via trl.RewardTrainer). It will handle mapping TensorDict token inputs to the expected HF forward arguments and write the output to a reward key in the TensorDict.
3. CI Tests & Examples
Round-trip tests validating a trl trainer reading from a TorchRL buffer, and a TorchRL GRPO step consuming an HF reward model. I will also add a short Sphinx tutorial demonstrating both integration paths.
Alternatives
An alternative is to force users to write their own custom data collators and TensorDict wrappers in their training scripts. However, this creates friction, leads to boilerplate code, and often results in bugs related to token device placement and TensorDict key mismatches. Native adapters solve this reliably for everyone.
Additional context
I'd like to take ownership of the implementation.
Checklist
Motivation
Spin-off from the Post-training efforts RFC (#3948) specifically for Workstream 2 (TRL interop).
Currently, many teams use Hugging Face
trl(e.g.,PPOTrainer,GRPOTrainer) but face bottlenecks with rollout storage and sampling at scale. Conversely, teams using TorchRL's highly efficient collectors and recipes often want to utilize the vast ecosystem of Hugging Face Reward Models, but writing custom wrappers is tedious.Nobody adopts a full stack at once. Providing native interoperability between
torchrlandtrllowers the barrier to entry, allowing users to adopt TorchRL's high-performance components piece by piece into their existing training loops.Solution
I recommend introducing a new module (e.g.,
torchrl.modules.llm.trl_interop) containing thin adapters so that users can seamlessly mix and matchtorchrlandtrlcomponents.Specifically, I want to implement:
1.
TorchRLBufferDataset(TorchRL -> TRL)A wrapper around TorchRL's
ReplayBufferthat implements the standard PyTorchIterableDataset/ Hugging FaceDatasetinterface. This allowstrltrainers to directly consume token/trajectory data generated by TorchRL's collectors.2.
HFRewardModelWrapper(TRL -> TorchRL)A
TensorDictModuleadapter that wraps a Hugging Face reward model (like those trained viatrl.RewardTrainer). It will handle mappingTensorDicttoken inputs to the expected HFforwardarguments and write the output to arewardkey in theTensorDict.3. CI Tests & Examples
Round-trip tests validating a
trltrainer reading from a TorchRL buffer, and a TorchRL GRPO step consuming an HF reward model. I will also add a short Sphinx tutorial demonstrating both integration paths.Alternatives
An alternative is to force users to write their own custom data collators and
TensorDictwrappers in their training scripts. However, this creates friction, leads to boilerplate code, and often results in bugs related to token device placement andTensorDictkey mismatches. Native adapters solve this reliably for everyone.Additional context
I'd like to take ownership of the implementation.
Checklist