perf: overlap non-colocated MInf + polict init

[tdene]

Closes #3777.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

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

• ...

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[tdene]

/ok to test 06937b3

[tdene]

/ok to test 93682b9

[cspades]

This PR addresses the performance gap documented in #3777 — megatron was the only non-colocated backend running policy and generation init serially. The fix is architecturally sound (3-way parallel ThreadPoolExecutor for the NeMo Gym path; reuse of initialize_generation_with_policy for the non-gym path) and the atomic checkpoint publication in publish_megatron_conversion is a welcome improvement.

Four findings below: one medium-severity hang on a failure path, and three low-severity suggestions.

Note: The PR description was not filled in — the motivation and approach are clear from the diff and issue #3777, but a one-line summary would help future readers.

Reviewed by 5 specialized agents (rl-expert, bug-finder, test-agent, design-reviewer, comment-reviewer) with devil's advocate pass. Pre-commit: PASS.

Generated by Claude Code

[tdene]

/ok to test a73fd1a

[tdene]

/ok to test 7463c34

[terrykong]

Thanks for this — the change hangs together well, and the part that makes it safe is easy to miss, so worth calling out: the overlap only works because two other things landed with it. megatron_policy_worker.py now skips handle_model_import entirely for refit-fed policies, and publish_megatron_conversion makes the HF->Megatron publish atomic. Without both, the two worker groups would race on the same cache directory — which is exactly what the old "initialize training first" ordering existed to prevent. Switching the completion marker to iter_0000000/run_config.yaml also fixes a real pre-existing bug where an interrupted conversion counted as complete.

One suggestion below, no blockers — a narrow failure in the new atomic-publish path.

Two notes on scope:

• PR #3873 (@wdykas) is merged into this branch rather than into main, so init_megatron_weight_synchronizer, the grpo_megatron_mxfp8_refit_gb200.sh assertion changes and the new test_setup_refits_noncolocated_megatron_while_nemo_gym_waits all land on main through this PR. Flagging it so the right person sees any feedback on those hunks.
• One thing we checked and are explicitly not asking you to change: that script tightens token_mult_prob_error from 2.0 to 1.5, but it is commented out at L1_Functional_Tests_GB200_MXFP8.sh:40 on this branch and on main, so the new bound has never run. Re-enabling needs a 4-GPU runner this PR does not provide, so there is nothing to do here — just be aware the bound is unverified for whoever re-enables it. The metric swap in the same hunk (generation_init_load_time_s -> weight_sync_time_s) is a genuine fix, since the old field is never set on the non-colocated non-Gym path.

We also looked hard at whether the new HF-derived provider for skip_weight_load policies could disagree with the training provider built from the conversion cache, and concluded it cannot — both routes end up building from the same AutoConfig.from_pretrained(...) result. No action needed; noting it so the question does not get re-opened.

Reviewed by 7 agents (RL expert, Megatron-Bridge expert, bug finder, test agent, design reviewer, comment reviewer, devil's advocate). pre-commit could not run in this environment, so no lint results are reported.

Generated by Claude Code

[tdene]

• One thing we checked and are explicitly not asking you to change: that script tightens token_mult_prob_error from 2.0 to 1.5, but it is commented out at L1_Functional_Tests_GB200_MXFP8.sh:40 on this branch and on main, so the new bound has never run. Re-enabling needs a 4-GPU runner this PR does not provide, so there is nothing to do here — just be aware the bound is unverified for whoever re-enables it. The metric swap in the same hunk (generation_init_load_time_s -> weight_sync_time_s) is a genuine fix, since the old field is never set on the non-colocated non-Gym path.

As with several other test-related issues, I'm choosing to defer this to #3570.

Once everything settles down, I vow that #3570 will fix all these assorted test issues.

[tdene]

/ok to test 5362cfb