[Connor-XY]

Depends on #8752, #8755, #8042 and #8045. Please review the diff from their merge: scoped diff.

• Warn on silent global process-group fallbacks in core helpers #8755 makes log_config_to_disk name its dumps by a given collection, and adds the test file this PR extends.
• Warn when the optimizer falls back to the global process groups #8752 gives Float16OptimizerWithFloat16Params its pg_collection and makes get_megatron_optimizer resolve the collection it passes on.
• Filter optimizer grad-norm duplicates by the collection's TP and GTP groups #8042 adds the ProcessGroupCollection import in optimizer.py that the new annotations use, and its set_duplicate_filter_groups calls sit next to this change. Name FSDP-DTensor experts by the model's expert-parallel layout #8045 changes the distrib_optimizer.py imports next to it.
• Warn when the optimizer falls back to the global process groups #8752 and Warn on silent global process-group fallbacks in core helpers #8755 both carry Require an explicit process-group collection in LanguageModule #6303. The merge of Warn when the optimizer falls back to the global process groups #8752 resolves its conflicts with Filter optimizer grad-norm duplicates by the collection's TP and GTP groups #8042 and Name FSDP-DTensor experts by the model's expert-parallel layout #8045 in optimizer/__init__.py and optimizer.py, and the merges regenerate the process-group allowlist. Each merge commit lists its resolutions.

#8731 changes the import block of fully_sharded_optimizer.py next to this change. Whichever lands second keeps both.

Why. log_config_to_disk names each dump by the ranks in the process groups of the object it logs. Called without them, it warns and falls back to the global parallel state (#8755). The FP32, mixed-precision, distributed and fully sharded optimizers and both Megatron-FSDP wrappers still called it without their collection. Stock training with --config-logger-dir therefore warned, and failed when fallback warnings were errors.

What.

• MixedPrecisionOptimizer, FP32Optimizer, DistributedOptimizer and FullyShardedOptimizer take a keyword-only pg_collection and pass it to log_config_to_disk. The subclasses forward it to MixedPrecisionOptimizer, and Float16OptimizerWithFloat16Params forwards the collection it already takes.
• get_megatron_optimizer passes its collection to these optimizers on the adam/sgd and the emerging-optimizer paths.
• FullyShardedDataParallelV1 and FullyShardedDataParallelV2 pass the collection they hold to log_config_to_disk.

Tests. tests/unit_tests/test_config_logger_process_groups.py adds test_optimizers_pass_their_collection for FP32Optimizer and Float16OptimizerWithFloat16Params: each logs its config once, with its collection.

On 4 GPUs against its base: test_config_logger_process_groups.py passes 7 tests here vs 5 on the base, and the 2 new tests pass. test_process_group_fallback.py (19 passed), transformer/test_cuda_graph_helper_process_groups.py (2 passed) and transformer/test_cuda_graphs.py (83 passed, 3 skipped) give the same results on both sides. None of these files shows a regression. The integration run merges all #6307 PRs onto main and runs the full unit-test suite against main on 4 GPUs and 8 GPUs (two nodes). In it, this PR's test file passes 7 tests; the file is not on main, so main has no results for it. Across the full suite, no test that passes on main fails, apart from one fused-MHC kernel check that failed on one rank and has failed the same way on main. On the same merged tree, end-to-end training smokes on 4 GPUs give losses bit-identical to main and record no process-group fallback warning.

Before the sync with main, the branch ran the test file on 4 CPU ranks with Gloo, with no regressions. The two new tests fail there, because both optimizers allocate a CUDA tensor right after they log. With that allocation stubbed out, both optimizers log once with their collection.

Compatibility. The new arguments are optional and keyword-only. Without them, dumps are named by the global ranks with a warning, as before this PR. On the standard grid the collection holds the global groups, so dump names do not change. Only runs with --config-logger-dir are affected.

• Needs determinism-exempt: optimizer/__init__.py and optimizer/optimizer.py are registered kernel files. The change there only passes the collection along, and numerics do not change.

Synced with main 77391e7 on 2026-11-01 by a merge commit that lists each conflict and its resolution, so existing commits and review links stay valid.

Later on 2026-11-01 the branch also merged main bfab98f. Its #8471 moves FullyShardedDataParallelV1 and FullyShardedDataParallelV2 from mcore_fsdp_adapter.py into mcore_adapter_v1.py and mcore_adapter_v2.py, and keeps mcore_fsdp_adapter.py as the factory that re-exports them. This PR's edits to those classes are applied unchanged where the code moved, and the process-group allowlist follows the move. On 4 GPUs against its new base, the Megatron-FSDP adapter tests and this PR's own tests (10 files) show no regressions: 169 passed here, 167 on the base. The two CUDA-graph end-to-end cases of test_mcore_fully_sharded_data_parallel.py fail on both sides, as they do on main.

Part of #6307.

🤖 Generated with Claude Code

[Connor-XY]

Superseded by #8908, which carries this change as commit 856b017 under the same title, with "Originally #8763." in its message. Closing so review happens in one place. The branch stays on the fork.