Relocate GroupInfo into shared process-group utilities

[fzyzcjy]

Move the GroupInfo dataclass out of training_utils/parallel.py into a shared
process_group_utils module that also provides multi-process-group helpers
(GeneralPGUtil / MultiPGUtil) for collectives over single or hierarchical
process groups, used by the cross-replica / effective-DP code paths.

• process_group_utils.py: GroupInfo + GeneralPGUtil / MultiPGUtil (+ tests).
• training_utils/parallel.py: import GroupInfo from the shared module.
• megatron / fsdp parallel.py: import GroupInfo from the shared module.
• distributed_utils.py: use GeneralPGUtil for masked-whiten all-reduce.

[gemini-code-assist]

Code Review

This pull request refactors process group utilities by moving GroupInfo to a new module miles.utils.process_group_utils and introducing GeneralPGUtil and MultiPGUtil to support both native PyTorch and torchft process groups. The feedback identifies two critical issues: first, _NativePGUtil methods will crash if group is None because dist.get_global_rank does not support None groups, so the source/destination rank should default to 0. Second, collective_bool_and creates a tensor on the CPU by default, which will cause a runtime error when used with an NCCL process group; the tensor should be placed on the appropriate device (e.g., CUDA).

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In _NativePGUtil, several methods (reduce, broadcast, gather, gather_object) call dist.get_global_rank(group, 0). If group is None (which is a common fallback in parallel states when certain parallelisms like pipeline or expert parallel are disabled), dist.get_global_rank(None, 0) will raise an AttributeError or TypeError and crash.

Since the global rank for group=None (the default WORLD group) is simply the group rank itself, we can safely default the destination/source rank to 0 when group is None.

    def reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None:
        dst = 0 if group is None else dist.get_global_rank(group, 0)
        dist.reduce(tensor, dst=dst, op=op, group=group)

    def broadcast(self, tensor: torch.Tensor, group: dist.ProcessGroup) -> None:
        src = 0 if group is None else dist.get_global_rank(group, 0)
        dist.broadcast(tensor, src=src, group=group)

    def barrier(self, group: dist.ProcessGroup) -> None:
        dist.barrier(group=group)

    def all_gather(
        self, output_tensors: list[torch.Tensor], input_tensor: torch.Tensor, group: dist.ProcessGroup
    ) -> None:
        dist.all_gather(output_tensors, input_tensor, group=group)

    def gather(
        self,
        input_tensor: torch.Tensor,
        gather_list: list[torch.Tensor] | None,
        group: dist.ProcessGroup,
    ) -> None:
        dst = 0 if group is None else dist.get_global_rank(group, 0)
        dist.gather(input_tensor, gather_list=gather_list, dst=dst, group=group)

    def gather_object(self, obj: Any, object_gather_list: list[Any] | None, group: dist.ProcessGroup) -> None:
        dst = 0 if group is None else dist.get_global_rank(group, 0)
        dist.gather_object(obj, object_gather_list, dst=dst, group=group)

[Shi-Dong]

Please apply.

[gemini-code-assist]

In collective_bool_and, the tensor is created on the CPU by default. If the provided group is an NCCL process group (which is standard for GPU training), calling all_reduce on a CPU tensor will result in a runtime crash (e.g., RuntimeError: Tensors must be CUDA tensors).

To prevent this, we should place the tensor on the appropriate device (e.g., using cuda if CUDA is available).

Suggested change

	def collective_bool_and(*, value: bool, group: dist.ProcessGroup) -> bool:
	"""Make a bool `and` operation on all ranks in this process group"""
	tensor = torch.tensor([1.0 if value else 0.0], dtype=torch.float32)
	GeneralPGUtil.create(group).all_reduce(tensor, group, op=dist.ReduceOp.MIN)
	return tensor.item() > 0.5
	def collective_bool_and(*, value: bool, group: dist.ProcessGroup) -> bool:
	"""Make a bool `and` operation on all ranks in this process group"""
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	tensor = torch.tensor([1.0 if value else 0.0], dtype=torch.float32, device=device)
	GeneralPGUtil.create(group).all_reduce(tensor, group, op=dist.ReduceOp.MIN)
	return tensor.item() > 0.5

[Shi-Dong]

Please apply.

[Shi-Dong]

Please hoist this import to the beginning

[Shi-Dong]

Please hoist this import to the beginning

[Shi-Dong]

Posted a couple of nits, but please also apply Gemini's suggestions as they look valid.

[Shi-Dong]

Please switch relative import to absolute import.

[Shi-Dong]

Please switch relative import to absolute import.

[Shi-Dong]

Please apply.

[Shi-Dong]

Please apply.