perf(refit): use NCCL Reshard for all async MXFP8 recipes

[seonjinn]

Summary

Use refit_transport: nccl_reshard in every existing async 1-off MXFP8 rollout performance recipe:

• DeepSeek V3
• Nemotron3 Super
• Qwen3-32B
• Qwen3-30B-A3B
• Qwen3-235B

Both the default collective transport and NCCL Reshard already quantize BF16 trainer weights to MXFP8 on the receiver side, so rollout numerics are unchanged. What this switch changes is the wire path: shard-to-shard resharding of BF16 trainer shards instead of full-tensor broadcast, which reduces refit bandwidth and peak memory. All five recipes are non-colocated async setups, which is exactly the configuration nccl_reshard requires. Nemotron3 Nano is not listed because main does not currently contain a Nano MXFP8 performance recipe.

The tests discover all *-async-1off-mxfp8-rollout.yaml files (with a guard that the glob is non-empty), require this transport, and validate each resolved recipe against check_nccl_reshard_refit_support, so future async MXFP8 recipes cannot silently fall back to the legacy collective or drift into a layout the reshard path rejects.

Test

• python3 -m pytest -q tests/test_mxfp8_rollout_recipes.py (22 passed, 5 skipped locally; the 5 validator cases require the full nemo_rl environment and run in CI unit tests)
• Ran check_nccl_reshard_refit_support against all 5 resolved recipes locally: all pass
• git diff --check

[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.

[seonjinn]

Self-review (adjudicated)

Two-lens review (production config semantics, test quality) with an adversarial spot-check pass against the code at 5bd246b9a (base = current main a43d71b31).

Bottom line: the production change is correct and all 5 recipes are functional as-is on main. Every defect found is in the new test or the PR body wording.

Production semantics — verified clean

• nccl_reshard is a valid transport (VllmRefitSelector, nemo_rl/models/generation/vllm/config.py:29), dispatched through weight_sync/factory.py:162-171 in the non-colocated branch and deep-validated by check_nccl_reshard_refit_support (nccl_reshard_utils.py:562) from grpo.py:1626-1634.
• Receiver-side MXFP8 conversion is end-to-end with no trainer-side cooperation: the sender ships BF16 shards as-is (megatron_policy_worker.py:3075); the receiver selects _bf16_to_mxfp8_receiver_quant_spec purely on the dtype pair (vllm_backend.py:1104-1105, quantization at :1050-1080).
• All 5 recipes resolve to configurations that satisfy every config-decidable validator precondition (non-colocated as required, gen ep=1, etp=1, vpp None, gen pp=1). None will be rejected at startup.
• Scoping is exactly right: resolving colocated.enabled for every *mxfp8-rollout.yaml in the perf directory shows the -async-1off- set is precisely the non-colocated set; every untouched sibling is colocated, where nccl_reshard hard-fails (nccl_reshard_utils.py:593-597). The 30ba3b case is subtle and correct: the async leaf flips colocated.enabled: false itself while its colocated parent is left alone.
• Ignore-pattern consistency: for all 5 recipes, every MXFP8-quantized target rides the bulk path (the only path with the receiver-quant spec); the misc path carries only BF16 targets.
• Base is current main; the recipe test suite passes locally (21 passed); ruff check/format and git diff --check are clean.

Test defects (should fix)

1. The second assertion is unfalsifiable. assert "refit_prequantize" not in generation_config["vllm_cfg"] (tests/test_mxfp8_rollout_recipes.py:346) — refit_prequantize exists nowhere on this branch or on main; a repo-wide grep's only hit is the assert itself (the knob lives on the unmerged perf: Reduce MXFP8 refit weight-transfer overheads #3294 branch). Even if someone set it, VllmRefitConfig is extra="allow" and nccl_reshard returns before refit_cfg validation (config.py:273-275), so it would be silently ignored. Drop the assert or defer it to whenever perf: Reduce MXFP8 refit weight-transfer overheads #3294 merges.
2. Value is asserted, validity is not. The test string-compares refit_transport but covers none of check_nccl_reshard_refit_support's ~14 preconditions. A future edit to a shared base config (e.g. virtual_pipeline_model_parallel_size: 2, a colocated flip, expert_tensor_parallel_size: 2) keeps this test green while the recipe dies at startup after a multi-node allocation (up to 64x4 GPUs). Calling the validator on each resolved config in the test closes this cheaply.
3. Empty glob makes the guard silently dead. An empty parametrize list yields SKIP + exit 0, and run_unit.sh:47-48 additionally treats pytest exit 5 as success. If the perf directory moves or the naming changes, the guard vanishes without a failure. Add a non-empty assertion on the glob result.
4. Coverage is keyed to the filename convention, enforced elsewhere. The "-async-1off-" naming is enforced by a different test (:271-274); relaxing it (e.g. a future -async-8off- MXFP8 recipe) or reordering the suffix (...-mxfp8-rollout-async-1off.yaml) escapes both globs silently.

PR body wording

• "these recipes do not enable trainer-side refit_prequantize" references a setting that does not exist on this branch — reviewers cannot find it because it is on unmerged perf: Reduce MXFP8 refit weight-transfer overheads #3294. Suggest rewording.
• The causal framing ("NCCL Reshard performs the receiver-side MXFP8 conversion, so...") reads as if receiver-side quantization were new. The default collective transport already quantized MXFP8 receiver-side (quantization/fp8.py:557-559); the actual change is full-tensor broadcast → shard-to-shard reshard (bandwidth/peak-memory win). Worth stating precisely since it is the PR's whole point.

Noted, no action required

• No gate couples nccl_reshard with async_grpo.in_flight_weight_updates: true (all 5 set it); the exposure is shared with the collective path via the common _weight_update_lifecycle, so this is inherited risk, not a regression.
• These 5 production perf recipes now default to a feature whose design doc is marked Experimental (docs/design-docs/nccl-reshard-refit.md:1-3). Deliberate, presumably, but worth a maintainer's conscious ack.

CI note: only lightweight checks have run; cicd-main needs a per-SHA /ok to test to exercise the new test (it runs in the L0_Unit_Tests_Other shard).

[seonjinn]

/ok to test 5bd246b

[seonjinn]

/ok to test 4e71df5