Support GRPO-style post-training in Kubeflow Trainer via TrainingRuntime and targeted Torch plugin support

[kapil27]

What you would like to be added?

I would like to propose initial support for GRPO-style post-training in Kubeflow Trainer.
Concretely, a good first step seems to be:

1. add one or more GRPO-focused ClusterTrainingRuntime examples/manifests
2. support a GRPO trainer image + launcher contract
3. extend the Torch plugin only where GRPO needs targeted validation or command mutation
4. keep rollout backend configuration explicit through runtime/config inputs###

Why is this needed?

Kubeflow Trainer v2 already has a strong direction for distributed training, LLM fine-tuning, and runtime-centric extensibility.

At the same time, LLM post-training is moving beyond SFT-only workflows. Techniques like GRPO are becoming relevant for:

• RLHF-style post-training
• reward-driven alignment
• iterative improvement of open-weight models
• experimentation with smaller or more specialized models after supervised fine-tuning

This seems like a natural area for Trainer to support next because:

• Trainer already supports specialized LLM fine-tuning behavior through runtimes and targeted plugin logic
• the LLM Trainer design already leaves room for future fine-tuning techniques such as RLHF
• there does not appear to be a clear upstream path yet for GRPO specifically
• I searched existing issues for GRPO and did not find a dedicated issue already tracking this.

If there is alignment on the direction, I would be happy to follow up with:

• a narrow implementation PR, or
• a proposal/KEP first, if that is the preferred path

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[andreyvelich]

Thanks for creating this @kapil27! Supporting RL runtimes would be really nice, I agree!

Do you know how should we design orchestration for GRPO?
Do we need to use LWS where leader is coordinating rollouts, or we can have dedicated Job in JobSet?
However, we didn't discuss whether JobSet should support cross-Job scaling logic driven by one "leader" Job.

/area rl

cc @astefanutti @tenzen-y

[Sridhar1030]

+1 , TRL GRPO on CustomTrainer + torch-distributed works well;
making it first-class via TrainingRuntime + examples + launcher contract (as proposed here) is a good next step.

[kapil27]

I agree the key question is whether GRPO needs new orchestration semantics or can fit the existing Trainer model.

As mentioned by @Sridhar1030, TRL GRPO on CustomTrainer + torch-distributed already appears viable, I think the best next step is to properly validate a minimal working example in the current Trainer model as part of this issue.

If that proves insufficient and GRPO needs coordinated multi-Job behavior (for example leader-driven rollout/training topology or cross-Job scaling), then I’d be happy to help work through a KEP in more detail.

wdyt?

[andreyvelich]

Sure, that sounds great, thank you @kapil27!