Repository: hpcaitech/ColossalAI
File: colossalai/booster/plugin/gemini_plugin.py (Python)

Code before:
            initial_scale=initial_scale,
            growth_factor=growth_factor,
            backoff_factor=backoff_factor,
            growth_interval=growth_interval,
            hysteresis=hysteresis,
            min_scale=min_scale,
            max_scale=max_scale,
            max_norm=max_norm,
            norm_type=norm_type,
        )
        self.enable_tensor_parallelism = tp_size > 1
        self.enable_all_optimization = enable_all_optimization
        self.enable_fused_normalization = enable_fused_normalization
        self.enable_flash_attention = enable_flash_attention
        self.enable_sequence_parallelism = enable_sequence_parallelism if self.enable_tensor_parallelism else False
        self.enable_jit_fused = enable_jit_fused
        self.enable_sequence_overlap = enable_sequence_overlap
        self.verbose = verbose

        self.tp_size = tp_size
        self.extra_dp_size = extra_dp_size
        world_size = dist.get_world_size()
        self.zero_size = world_size // (self.tp_size * self.extra_dp_size)
        assert world_size == (self.tp_size * self.extra_dp_size) * self.zero_size, f"The global group size can't be evenly divided by the subgroup size."
        assert self.zero_size > 1, f"Gemini's group size should be greater than 1, please reduce tensor paralllelism group size or extra-dp group size."
        if self.tp_size == 1 and self.extra_dp_size == 1:
            self.extra_dp_group = None
            self.tp_group = None
            self.zero_group = None
        elif self.tp_size == 1:
            self.tp_group = None
            self.pg_mesh = ProcessGroupMesh(self.zero_size, self.extra_dp_size)
            self.zero_group = self.pg_mesh.get_group_along_axis(0)
            self.extra_dp_group = self.pg_mesh.get_group_along_axis(1)
        elif self.extra_dp_size == 1:
            self.extra_dp_group = None
            self.pg_mesh = ProcessGroupMesh(self.zero_size, self.tp_size)
            self.zero_group = self.pg_mesh.get_group_along_axis(0)
            self.tp_group = self.pg_mesh.get_group_along_axis(1)
        else:
            self.pg_mesh = ProcessGroupMesh(self.zero_size, self.extra_dp_size, self.tp_size)
            self.zero_group = self.pg_mesh.get_group_along_axis(0)
            self.extra_dp_group = self.pg_mesh.get_group_along_axis(1)
            self.tp_group = self.pg_mesh.get_group_along_axis(2)

        self.shard_config = ShardConfig(
            tensor_parallel_process_group=self.tp_group,
            enable_tensor_parallelism=self.enable_tensor_parallelism,
            enable_all_optimization=self.enable_all_optimization,
            enable_fused_normalization=self.enable_fused_normalization,
            enable_flash_attention=self.enable_flash_attention,

Review comment (refactor):
This is complicated. Actually, creating an instance of ProcessGroupMesh won't take too long.
Here is pseudo-code:
```python
DP_AXIS, ZERO_AXIS, TP_AXIS = 0, 1, 2
self.pg_mesh = ProcessGroupMesh(self.extra_dp_size, self.zero_size, self.tp_size)
self.zero_group = self.pg_mesh.get_group_along_axis(ZERO_AXIS) if self.zero_size < world_size else _get_default_group()
self.extra_dp_group = self.pg_mesh.get_group_along_axis(DP_AXIS) if self.extra_dp_size > 1 else None
self.tp_group = self.pg_mesh.get_group_along_axis(TP_AXIS) if self.tp_size > 1 else None
```

Code after:
            gpu_margin_mem_ratio=gpu_margin_mem_ratio,
        )
        self.optim_kwargs = dict(
            initial_scale=initial_scale,
            growth_factor=growth_factor,
            backoff_factor=backoff_factor,
            growth_interval=growth_interval,
            hysteresis=hysteresis,
            min_scale=min_scale,
            max_scale=max_scale,
            max_norm=max_norm,
            norm_type=norm_type,
        )
        self.enable_tensor_parallelism = tp_size > 1
        self.enable_all_optimization = enable_all_optimization
        self.enable_fused_normalization = enable_fused_normalization
        self.enable_flash_attention = enable_flash_attention
        self.enable_sequence_parallelism = enable_sequence_parallelism if self.enable_tensor_parallelism else False
        self.enable_jit_fused = enable_jit_fused
        self.enable_sequence_overlap = enable_sequence_overlap
        self.verbose = verbose

        self.tp_size = tp_size
        self.extra_dp_size = extra_dp_size
        world_size = dist.get_world_size()
        self.zero_size = world_size // (self.tp_size * self.extra_dp_size)
        assert world_size == (self.tp_size * self.extra_dp_size) * self.zero_size, f"The global group size can't be evenly divided by the subgroup size."

        self.pg_mesh = ProcessGroupMesh(self.zero_size, self.extra_dp_size, self.tp_size)
        self.zero_group = self.pg_mesh.get_group_along_axis(ZERO_AXIS) if self.zero_size < world_size else _get_default_group()
        self.extra_dp_group = self.pg_mesh.get_group_along_axis(DP_AXIS) if self.extra_dp_size > 1 else None
        self.tp_group = self.pg_mesh.get_group_along_axis(TP_AXIS) if self.tp_size > 1 else None

        self.shard_config = ShardConfig(
            tensor_parallel_process_group=self.tp_group,
            enable_tensor_parallelism=self.enable_tensor_parallelism,
            enable_all_optimization=self.enable_all_optimization,
            enable_fused_normalization=self.enable_fused_normalization,
            enable_flash_attention=self.enable_flash_attention,
            enable_jit_fused=self.enable_jit_fused,
            enable_sequence_parallelism=self.enable_sequence_parallelism,
            enable_sequence_overlap=self.enable_sequence_overlap,
        )

    def support_no_sync(self) -> bool:
        return False

    def control_precision(self) -> bool:
        return True

    def supported_precisions(self) -> List[str]: