Method flatten

colossalai/device/device_mesh.py:481–498

Flatten the logical mesh into an effective 1d logical mesh,

(self)

        return global_pg_ranks

    def flatten(self):
        """
        Flatten the logical mesh into an effective 1d logical mesh,
        """
        if self._is_init_from_process_group:
            raise RuntimeError(
                "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known."
            )

        flatten_mesh_shape_size = len(self._mesh_shape)
        flatten_mesh_shape = [self.num_devices]
        return DeviceMesh(
            self._physical_mesh_id,
            tuple(flatten_mesh_shape),
            mesh_alpha=[max(self.mesh_alpha)] * (flatten_mesh_shape_size - 1),
            mesh_beta=[max(self.mesh_beta)] * (flatten_mesh_shape_size - 1),
            init_process_group=self._init_process_group,
        )

    def all_gather_cost(self, num_bytes, mesh_dim):
        num_devices = self.logical_mesh_id.shape[mesh_dim]