[D] What framework do you use for RL post-training at scale?

Hi!

I’m sorry if I’m not using the correct tag, I didn’t know which one to pick, and I’m sorry if the question is not aligned with the sub’s purpose, please let me know if that is the case and feel free to block the post as well.

I’m trying to do some post-training at a somewhat large scale, but I’m struggling with some of the known frameworks out there.

For some context, I’m trying to do RL on function calling. This is more of a long-term research project, and I’d like to have the flexibility of writing my own environments and algorithms or modify the existing ones.

I have a preference for FSDP (and other parallelism paradigms but through Pytorch’s `DeviceMesh` and custom code if possible) and vLLM but I can adapt if needed. Ideally the framework can just support the “mainstream” models out of the box (Qwen, Mistral etc.) but I don’t mind writing support for the model I want to use if needed. Currently I have tried this:

verl (from ByteDance): the latest release is from last month but there are fixes almost every day I think. I did spend quite some time in understanding it and its architecture and it should be pretty good but I wanted to try a small “toyish” setup first with just pattern matching of the function call made by the model on the expected call (so a custom reward function), and with a custom agent loop that does not load all of the dataset’s tool but I hit import errors that I had to fix in the repo itself and whatnot and I don’t know how much struggle I’ll have to go through later on. Which doesn’t really bother me but I want to know if there are better alternatives.

torchforge (from meta-pytorch): this seems ideal to me but it is very early in development, I had issues just running their tests and I can do a lot of hacky stuff to get my way through but I’d prefer not and I’m not totally sure I have the capability to get my way through everything since they use Monarch instead of Ray and I’m not familiar with it at all.

OpenRLHF: I haven’t tried it yet, though I’m familiar with Deepspeed, I’m mostly familiar with Pytorch’s FSDP and they don’t seem to support it yet. But it doesn’t bother me, I just haven’t had the chance to look at it yet. But they seem to be lightweight, which I like. It is updated less frequently than verl but I think it’s still up to date.

trl: I used it for SFT quite a lot so I know it’s limitations and I don’t think it’s the right fit for my use case.

– I also looked at NVIDIA’s Gym and RL. It seems like Gym is the infra and RL is the algo / optimization, I’d prefer ideally one library that does both, like the others instead of having to do the pipelining myself. And I don’t like the fact that you can’t just `uv add` them or `pip install`. Granted I can clone the repos and install them in my codebase as editables, but I haven’t tried yet, maybe there will be dependency issues or just CUDA issues, I did struggle a lot in the past with installing NVIDIA repos.

I’d be very grateful if you can share your experience on this. Thanks!

EDIT: What I mean by imports issues in verl are imports of deprecated code from transformers even though verl itself relies on recent releases of transformers. So not issues of my code not importing stuff from verl correctly. I also saw some optional dependency group that relies on an old unmaintained package it seems and I’d just like to avoid having to deal with these issues.

submitted by /u/ReinforcedKnowledge
[link] [comments]

Liked Liked