Method get_group

internlm/core/context/parallel_context.py:335–345

Returns the group of the current device for `parallel_mode`. Args: parallel_mode: The chosen parallel mode. Returns: torch.distributed.ProcessGroup: The group of the current device for `parallel_mode`.

(self, parallel_mode: ParallelMode)

        return self._world_sizes.get(parallel_mode, 1)

    def get_group(self, parallel_mode: ParallelMode):
        """Returns the group of the current device for `parallel_mode`.

        Args:
            parallel_mode: The chosen parallel mode.

        Returns:
            torch.distributed.ProcessGroup: The group of the current device for `parallel_mode`.
        """
        self._check_parallel_mode(parallel_mode)
        return self._groups[parallel_mode]

    def get_ranks_in_group(self, parallel_mode: ParallelMode):
        """Returns the rank of the current device for `parallel_mode` in the group.