fix(megatron): shard Megatron rollouts across DP coordinator leaders

[YigongQin]

What does this PR do ?

Shard Megatron inference rollouts across data-parallel coordinator leaders so multi-replica generation scales like vLLM, instead of sending every prompt batch to Ray worker 0 / global rank 0 only.

Issues

None

Usage

No config changes required. Existing Megatron generation recipes automatically use all DP replicas when world_size / (TP × PP × CP × EP) > 1.

Sync rollouts shard the batch across DP leaders:

outputs = policy_generation.generate(generation_input_data, greedy=False)

Async rollouts round-robin single-sample requests across DP leaders (same pattern as vLLM async):

policy:
  generation:
    backend: megatron
    mcore_generation_config:
      async_engine: true

NeMo Gym with multiple DP replicas gets one OpenAI URL per DP leader when expose_http_server: true.

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

Problem: mcore starts a DynamicInferenceEngine + DP coordinator on every replica, but NeMo RL previously routed all rollouts through worker 0 only. Extra DP GPUs did not increase rollout throughput.

Changes:

• MegatronGeneration.generate: shard prompts with shard_by_batch_size(dp_size) and dispatch via run_all_workers_sharded_data (vLLM parity).
• MegatronGeneration.generate_async: round-robin across DP leader workers.
• MegatronGenerationMixin: use parallel_state.get_data_parallel_rank() == 0 for InferenceClient, sleep/wake, HTTP server, and request submission (per-replica, not global rank 0 only).
• Collect NeMo Gym OpenAI base URLs from each DP leader.

Tests added/updated:

• tests/unit/models/generation/test_megatron_generation_dp_shard.py — CPU mock test for driver-side sharding.
• tests/unit/models/generation/test_megatron_generation.py — async test accepts any DP leader.

Suggested local verification:

uv run pytest tests/unit/models/generation/test_megatron_generation_dp_shard.py
uv run pytest tests/unit/models/generation/test_megatron_generation.py -m mcore

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.