Repository: hpcaitech/ColossalAI
File: colossalai/shardformer/shard/shard_config.py (Python)

Code before:
        # get the tensor parallel size
        if not self.enable_tensor_parallelism:
            self._tensor_parallel_size = 1
        else:
            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)

        # get the sequence parallel size
        if not self.enable_sequence_parallelism:
            self._sequence_parallel_size = 1
        else:
            self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)

    def _turn_on_all_optimization(self):
        """
        Turn on all optimization.
        """
        # you can add all the optimization flag here
        try:
            from apex.normalization import FusedLayerNorm as ApexFusedLayerNorm  # noqa

            apex_installed = True
        except ImportError:
            apex_installed = False
            warnings.warn("You set enable_all_optimization=True, but apex is not installed.")
        try:
            import flash_attn  # noqa

            flash_attn_installed = True
        except ImportError:
            flash_attn_installed = False
            warnings.warn("You set enable_all_optimization=True, but flash_attn is not installed.")

        self.enable_fused_normalization = apex_installed
        self.enable_flash_attention = flash_attn_installed
        self.enable_jit_fused = True
        # This can cause non-in-place param sharding when used without ZeRO.
        # It may also slow down training when seq len is small. Plz enable manually.
        # self.enable_sequence_parallelism = True
        # self.enable_sequence_overlap = True

    def _infer(self):
        """
        Set default params for inference.
        """
        # assert self.pipeline_stage_manager is None, "pipeline parallelism is not supported in inference for now"


Review comment (suggestion):
flash_attn is not essential for `enable_flash_attention=True`. if flash_attn is not installed, it will use torch sdpa

Code after:
        # get the tensor parallel size
        if not self.enable_tensor_parallelism:
            self._tensor_parallel_size = 1
        else:
            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)

        # get the sequence parallel size
        if not self.enable_sequence_parallelism:
            self._sequence_parallel_size = 1
        else:
            self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)

    def _turn_on_all_optimization(self):
        """
        Turn on all optimization.
        """
        # you can add all the optimization flag here
        try:
            from apex.normalization import FusedLayerNorm as ApexFusedLayerNorm  # noqa

            apex_avail = True
        except ImportError:
            apex_avail = False
            warnings.warn("You set enable_all_optimization=True, but apex is not installed.")

        self.enable_fused_normalization = apex_avail
        self.enable_flash_attention = True
        self.enable_jit_fused = True
        # This can cause non-in-place param sharding when used without ZeRO.
        # It may also slow down training when seq len is small. Plz enable manually.
        # self.enable_sequence_parallelism = True
        # self.enable_sequence_overlap = True

    def _infer(self):
        """
        Set default params for inference.
        """
        # assert self.pipeline_stage_manager is None, "pipeline parallelism is not supported in inference for now"