Skip to content

[RFC] Post-training efforts #3948

Description

@theap06

##Summary
TorchRL's LLM post-training stack has grown quickly: GRPO and SFT objectives
(torchrl/objectives/llm), vLLM backends and LLM policy wrappers
(torchrl/modules/llm), an inference server, weight-sync schemes, and
recipes for GRPO, expert iteration, VLA-GRPO and reward-model training. What
it does not yet have is an adoption path for teams that already run their
post-training on other stacks (TRL, NeMo-RL, or in-house loops).

This RFC proposes treating interoperability as the product: making
TorchRL's post-training components individually consumable from existing
workflows, rather than asking users to migrate whole training loops. Nobody
adopts a full stack at once; they adopt one component that solves a problem,
then a second.

Motivation

The teams doing LLM post-training today mostly have a working loop already.
Their pain points are component-level: replay/rollout storage that scales,
collectors that decouple generation from training, losses with correct
token-level bookkeeping, and observability that makes RL runs debuggable.
TorchRL has strong answers to each, but they are currently easiest to use
from inside a TorchRL loop. If a TRL or NeMo-RL user could pull in a single
TorchRL component with a thin adapter and no loop rewrite, each component
becomes an on-ramp to the rest of the stack.

Proposed workstreams

Each is independently useful and separately mergeable; none blocks another.

1. Stable component boundaries (the foundation). Define and document the
small set of interfaces external loops consume: "give me a buffer"
(extend/sample over token or trajectory data), "give me rollouts"
(collector/inference-server protocol, including weight sync), "give me a
loss" (tensordict-in, losses-out with documented key contracts). These
mostly exist implicitly; the work is making them explicit, versioned and
tested as public contracts, so adapters stay thin and don't break on either
side's releases.

2. TRL interop. Thin adapters so TRL users can use TorchRL replay
buffers and collectors inside TRL trainers, and conversely so TorchRL
recipes can consume HF-ecosystem reward models and datasets with zero
friction. Deliverable: an examples page + adapter module + CI-tested
round-trip ("TRL trainer with TorchRL buffer" and "TorchRL GRPO with HF
reward model").

3. NeMo-RL interop. Same shape as (2), scoped after a short discovery
pass on where the natural seams are (rollout storage and weight sync are
the likely candidates). NVIDIA-stack users are heavy post-training adopters;
even a storage-only integration is a meaningful on-ramp.

4. Post-training observability. A standardized logging layer for the
metrics post-training practitioners actually watch: per-token/per-sequence
reward stats, KL to reference, entropy, generation throughput, weight-sync
latency, buffer utilization. One hook set, emitting to the existing logger
backends, used uniformly across the GRPO/expert-iteration/RM recipes, and
usable standalone from external loops. This makes TorchRL runs comparable
and debuggable, and it is the component external users can adopt with the
least commitment.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions