feat: GRPO example for Qwen3 32b context length=128k

[soodoshll]

What does this PR do ?

Address #883.

• Qwen3 32b + 128k context is runnable using TP8+CP4+actckpt on 4x8H100 nodes. Example run: https://wandb.ai/nvidia/nemo-rl/runs/5qhgqcqf?nw=nwuserqidongs
• Roll-out is extremely slow in this case so I use a rather low batch size in this recipe.
• I use YaRN to extend the context length to 128k in vllm. Not 100% sure if megatron is correctly handling this part.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[wangshangsam]

Are there any theoratical differences between this and doing TP8+CP4+actckpt on DTensor, in which case DTensor couldn't work?

[soodoshll]

Are there any theoratical differences between this and doing TP8+CP4+actckpt on DTensor, in which case DTensor couldn't work?

need a profiling to see why. btw, dtensor can run with seqlen=64k

[wangshangsam]

Actually, by exsisting convention, the hardware setup should come before the runtime config, and ideally the max supported context length should come with the model & task, e.g., grpo-math-qwen3-32b-128k-4n8g-megatrontp8cp4.yaml.

Could you also rename the other grpo-math-qwen3-30ba3b-megatron-tp4-32k.yaml into grpo-math-qwen3-30ba3b-32k-4n8g-megatrontp4ep8.yaml? We forgot about this aspect when we reviewed #918 last week (cc @pjin-nvidia)

[wangshangsam]

Are there any theoratical differences between this and doing TP8+CP4+actckpt on DTensor, in which case DTensor couldn't work?

need a profiling to see why. btw, dtensor can run with seqlen=64k

Let's get this PR merged first. The key objective of the issue this PR is addressing is to unblock the Nemotron folks, so if MCore path already works, let's leave it at that. We can dig into the memory profile for very long seq len when addressing #885.

[wangshangsam]

Btw, since you are adding a new recipe, you need to add a bash script that tests this new recipe in the CI too (there's a unit test that check if you have done that or not). Refers to #926 tests/test_suites/llm and tests/test_suites/nightly for an example.

[soodoshll]

Actually, by exsisting convention, the hardware setup should come before the runtime config, and ideally the max supported context length should come with the model & task, e.g., grpo-math-qwen3-32b-128k-4n8g-megatrontp8cp4.yaml.

Could you also rename the other grpo-math-qwen3-30ba3b-megatron-tp4-32k.yaml into grpo-math-qwen3-30ba3b-32k-4n8g-megatrontp4ep8.yaml? We forgot about this aspect when we reviewed #918 last week (cc @pjin-nvidia)

fixed. but overall the recipe naming still lacks consistency. It'd be better if we can have a naming convention.

[wangshangsam]

Actually, by exsisting convention, the hardware setup should come before the runtime config, and ideally the max supported context length should come with the model & task, e.g., grpo-math-qwen3-32b-128k-4n8g-megatrontp8cp4.yaml.
Could you also rename the other grpo-math-qwen3-30ba3b-megatron-tp4-32k.yaml into grpo-math-qwen3-30ba3b-32k-4n8g-megatrontp4ep8.yaml? We forgot about this aspect when we reviewed #918 last week (cc @pjin-nvidia)

fixed. but overall the recipe naming still lacks consistency. It'd be better if we can have a naming convention.

Naming convention: https://github.com/NVIDIA-NeMo/RL/tree/main/tests/test_suites#naming

[terrykong]

wouldn't this fail according to the wandb link you shared?

I think with longer generations, we'll probably run into outliers that skew the mean. This one is run for so few steps, it's probably hard to write something robust. Maybe:

Suggested change

	'mean(data["train/token_mult_prob_error"]) < 1.1' \
	'min(data["train/token_mult_prob_error"]) < 1.1' \

and then add a comment above why you use min for this particular test

I made an issue to track this: #1039

[terrykong]

which models need this? is it possible to handle this in code?

In the past when we had stuff like this, the consensus was to handle it in code since we knew which model types needed it, ex: fdb565c

regardless of if this is handled in code or yaml, it should probably have an entry in the model-quirks.md so we have documentation

[soodoshll]

It is used by qwen3 + long context length: https://huggingface.co/Qwen/Qwen3-32B#processing-long-texts.

Since it's a optional configuration that can be changed by the user, I tend to explicitly put it in yaml.

Will update model-quirks.md to reflect this.

[soodoshll]

This PR is not merged because it depends on the support of YaRN in the Megatron backend to be ready.
Related MR in Megatron-LM: #3854

[guyueh1]

I think mcore path cannot parse/handle the rope_scaling.type="yarn" even if you pass these arguments, this might succeed but your model arch is still rope, this seems confusing.

[guyueh1]

maybe we should error out in mcore path if seeing rope_scaling.type=="yarn" until necessary support is added;

@yaoyu-33 to confirm if the qwen and llama bridge in mbridge can parse and handle this field. I believe this is only supported in deepseek model type.