[ruizhang0101]

Fixes #2452.

Problem

monkey_patch_torch_dist() replaces the torch.distributed collectives with wrappers that unwrap ReloadableProcessGroup to the inner communicator before calling torch. Megatron does not look those functions up at call time. Three modules bind them when they are imported:

• megatron/core/distributed/param_and_grad_buffer.py: dist_reduce_scatter_func, dist_all_gather_func, _coalescing_manager
• megatron/core/tensor_parallel/mappings.py: dist_reduce_scatter_func, dist_all_gather_func
• megatron/core/timers.py: dist_all_gather_func

slime/backends/megatron_utils/actor.py imports megatron.core at module level and applies the patch later, in MegatronTrainRayActor.init, so Megatron keeps the unpatched originals. Every DDP gradient reduce-scatter, distributed-optimizer all-gather and tensor-parallel collective therefore hands the Python subclass group straight to torch, which dispatches it through the PyProcessGroup trampoline instead of the inner NCCL group. torch 2.11 tolerates that detour. torch 2.13 renamed the collectives (reduce_scatter_single, all_gather_single, *_single_coalesced), which the wrapper does not override, and the first optimizer step segfaults in c10d::PyProcessGroup::PyWorkHolder::wait (details and log in #2452).

Change

After the patch has replaced the torch.distributed functions, monkey_patch_torch_dist() now walks the already-imported megatron.* modules and redirects any attribute that is identical to a replaced original to its wrapper. The mapping is taken from a snapshot of torch.distributed made before patching, so it stays correct if the patched set changes. Modules imported after the patch bind the wrappers on their own and need nothing.

The behavior for runs that do not call monkey_patch_torch_dist() (no --offload-train) is unchanged. Megatron's call sites are untouched.

Test

tests/test_reloadable_process_group_megatron_bindings.py registers a fake megatron.core.distributed.param_and_grad_buffer module holding the original functions, applies the patch in a spawned process, and checks that the three bindings now point at the wrappers while an unrelated attribute and a non-megatron module keep the originals.

Verification

• pre-commit run --files on both files: passed (ruff, autoflake, isort, black).
• pytest tests/test_reloadable_process_group_megatron_bindings.py tests/test_reloadable_process_group_memory_check.py: 4 passed.
• GPU: 1x H200 trainer + 1 rollout engine, Megatron Bridge, Qwen2.5-0.5B-Instruct, --use-critic, --offload-train, --update-weight-transport disk, torch 2.13.0+cu130, Slime at 41014d1 with this patch applied to reloadable_process_group.py and no other change. Without the patch the critic actor segfaulted at its first optimizer step in 3 of 3 runs. With it, 3 training steps completed (175 s, 137 s, 144 s), 4 weight updates reached the rollout engine, no segfault. Not covered by that run: tensor parallel > 1 (tensor_parallel/mappings.py holds the same kind of bindings and is rebound by the same loop).