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

Code before:
        self,
        model: nn.Module,
        optimizer: Optional[Optimizer] = None,
        criterion: Optional[Callable] = None,
        dataloader: Optional[DataLoader] = None,
        lr_scheduler: Optional[LRScheduler] = None,
    ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]:
        param_info = get_param_info(optimizer)
        if not isinstance(model, ModelWrapper):
            # convert model to sync bn
            # FIXME(ver217): gemini does not support sync bn
            # In torch/nn/modules/_functions.py, line 22, ``mean, invstd = torch.batch_norm_stats(input, eps)`` will get fp32 mean and invstd even though the input is fp16.
            # This inconsistency of dtype will cause the error.
            # We have two possible solutions:
            # 1. keep batch norm always in fp32. This is hard for gemini, as it use chunks.
            # 2. patch sync bn or write a new on. This is relatively easy, but we need to test it.
            # model = nn.SyncBatchNorm.convert_sync_batchnorm(model, None)

            # wrap the model with Gemini
            self.dp_group = None
            self.tp_group = None
            try:
                dp_size = dist.get_world_size() // self.tp_size
                assert dp_size > 1, f"the size of DP group should greater than 1. Please reduce the TP group size."
                self.pg_mesh = ProcessGroupMesh(dp_size, self.tp_size)
                self.dp_group = self.pg_mesh.get_group_along_axis(0)
                self.tp_group = self.pg_mesh.get_group_along_axis(1)
                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)
                shardformer = ShardFormer(shard_config)
                model, _ = shardformer.optimize(model)
            except NotImplementedError as e:
                print(f"Tensor Parallelism policy for {model.__class__} is not implemented yet\n.")

            model = GeminiDDP(model, **self.gemini_config, process_group=self.dp_group, verbose=self.verbose)

        if optimizer is not None and not isinstance(optimizer, OptimizerWrapper):
            optimizer = GeminiOptimizer(
                optimizer, model, **self.zero_optim_config, **self.optim_kwargs, param_info=param_info, tp_group=self.tp_group, verbose=self.verbose
            )

        return model, optimizer, criterion, dataloader, lr_scheduler

    def control_checkpoint_io(self) -> bool:
        return True

Review comment (suggestion):
It would be clearer if setting `DP_AXIS` and `TP_AXIS` to 0 and 1 respectively and use these two variables instead. 

Code after:
    def support_no_sync(self) -> bool:
        return False

    def control_precision(self) -> bool:
        return True

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

    def control_device(self) -> bool:
        return True

    def supported_devices(self) -> List[str]:
        return ["cuda"]

    def configure(
        self,
        model: nn.Module,
        optimizer: Optional[Optimizer] = None,
        criterion: Optional[Callable] = None,
        dataloader: Optional[DataLoader] = None,
        lr_scheduler: Optional[LRScheduler] = None,
    ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]:
        optimizer_params_info = get_param_info(optimizer)
        if not isinstance(model, ModelWrapper):
            # convert model to sync bn
            # FIXME(ver217): gemini does not support sync bn
            # In torch/nn/modules/_functions.py, line 22, ``mean, invstd = torch.batch_norm_stats(input, eps)`` will get fp32 mean and invstd even though the input is fp16.
            # This inconsistency of dtype will cause the error.
            # We have two possible solutions:
            # 1. keep batch norm always in fp32. This is hard for gemini, as it use chunks.
            # 2. patch sync bn or write a new on. This is relatively easy, but we need to test it.
            # model = nn.SyncBatchNorm.convert_sync_batchnorm(model, None)

            # wrap the model with Gemini
            if self.enable_tensor_parallelism:
                shardformer = ShardFormer(self.shard_config)
                model, _ = shardformer.optimize(model)

            model = GeminiDDP(model, **self.gemini_config, process_group=self.dp_group, verbose=self.verbose)

        if optimizer is not None and not isinstance(optimizer, OptimizerWrapper):
            optimizer = GeminiOptimizer(
                optimizer,
                model,
                **self.zero_optim_config,
                **self.optim_kwargs,
                tp_group=self.tp_group,
                optimizer_params_info=optimizer_params_info,
                verbose=self.verbose,
            )