Megatron Inference support for GRPO (Dynamic batching)

[nathan-az]

GRPO appears to support megatron as the generation backend, but then will break later in the script when trying to call prepare_refit_info on a NoneType.

I am curious about whether it's possible to use a TRT Model Optimizer quantized model with the Fake Quant layers for training, for QAT GRPO. I assumed this would be simpler with the Megatron backend since no refit would be required (where I imagine this would be more complex for vLLM which wouldn't support those layers).

So my main question is: what is the state of using Megatron as the generation backend? Should this be supported? (If so, I think this doubles as a bug report).

Additional questions are:

• Is there any loose guidance or comparison on generation speed for Megatron vs vLLM?
• Does my intention regarding QAT and being easier with Megatron for generation hold water or am I misunderstanding something?

[nathan-az]

Closing for now since I was able to answer this.

• Doing a hasattr check on megatron allowed me to do inference with non-MoE models (although MoE models had some issues).
• Generation with megatron was much slower than untuned vllm, about 1/4 throughput
• Using checkpoints from modeloptimizer - even with megatron - was not straightforward

I will create separate, specific issues at some point in future.

[euronymous-aithal]

hi @nathan-az : latest version of megatron inference hasn't been integrated into Nemo-RL yet. Currently only static backend from megatron inference has been included. Most advanced features like dynamic batching etc. will be integrated into NemoRL Megatron inference backend in the coming release which should address the perf gap with vllm

[nathan-az]

Most advanced features like dynamic batching etc. will be integrated into NemoRL Megatron inference backend in the coming release which should address the perf gap with vllm

Thanks @euronymous-aithal ! Are there any issues/roadmaps/PRs to follow? It would be great to follow the progress and hopefully see performance comparisons as that release comes.

[shanmugamr1992]

Hello @nathan-az at this time we are still actively looking into several features which give us a boost over VLLM some being (Cuda graphs, native support for data parallel, dynamic batching with chunked kv cache for both prefill and decode and a bunch of other nuanced optimizations ) . We will start integrating all this into nemo-rl in the next couple of weeks, will post an update once the MR's are in place.
If you want to follow the develpoment happening for megatron inference, please take a look at the megatorn-lm .

[nathan-az]

Great, thanks @shanmugamr1992 . In that case I'll re-open this issue.