We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/shardformer/shard/shard_config.py`
Content:
```
1 import warnings
2 from dataclasses import dataclass, field
3 from typing import Any, Dict, Optional
4 
5 import torch.distributed as dist
6 from torch.distributed import ProcessGroup
7 
8 from colossalai.pipeline.stage_manager import PipelineStageManager
9 
10 from .grad_ckpt_config import GradientCheckpointConfig
11 
12 __all__ = ["ShardConfig"]
13 SUPPORT_SP_MODE = ["split_gather", "ring", "all_to_all"]
14 
15 
16 @dataclass
17 class ShardConfig:
18     r"""
19     The config for sharding the huggingface model
20 
21     Args:
22         tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.
23         pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.
24         enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.
25         enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.
26         enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.
27         enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.
28         enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.
29         enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.
30         gradient_checkpoint_config (Optional[GradientCheckpointConfig]): The gradient checkpoint config. Defaults to None.
31         enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.
32     """
33     tensor_parallel_process_group: Optional[ProcessGroup] = None
34     sequence_parallel_process_group: Optional[ProcessGroup] = None
35     pipeline_stage_manager: Optional[PipelineStageManager] = None
36     enable_tensor_parallelism: bool = True
37     enable_all_optimization: bool = False
38     enable_fused_normalization: bool = False
39     enable_flash_attention: bool = False
40     enable_jit_fused: bool = False
41     enable_sequence_parallelism: bool = False
42     sequence_parallelism_mode: str = None
43     enable_sequence_overlap: bool = False
44     parallel_output: bool = True
45     make_vocab_size_divisible_by: int = 64
46     gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None
47     extra_kwargs: Dict[str, Any] = field(default_factory=dict)
48     # pipeline_parallel_size: int
49     # data_parallel_size: int
50     # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']
51 
52     @property
53     def tensor_parallel_size(self):
54         return self._tensor_parallel_size
55 
56     @property
57     def sequence_parallel_size(self):
58         return self._sequence_parallel_size
59 
60     def __post_init__(self):
61         # turn on all optimization if all_optimization is set to True
62         if self.enable_all_optimization:
63             self._turn_on_all_optimization()
64 
65         if self.enable_sequence_parallelism:
66             self.sequence_parallelism_mode = (
67                 "split_gather" if self.sequence_parallelism_mode is None else self.sequence_parallelism_mode
68             )
69             assert (
70                 self.sequence_parallelism_mode in SUPPORT_SP_MODE
71             ), f"Sequence parallelism mode {self.sequence_parallelism_mode} is not in the supported list {SUPPORT_SP_MODE}"
72             if self.sequence_parallelism_mode in ["split_gather", "ring"]:
73                 assert (
74                     self.enable_tensor_parallelism
75                 ), f"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is True"
76             elif self.sequence_parallelism_mode in ["all_to_all"]:
77                 assert (
78                     not self.enable_tensor_parallelism
79                 ), f"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is False"
80                 if self.enable_sequence_overlap:
81                     self.enable_sequence_overlap = False
82                     warnings.warn(
83                         f"The enable_sequence_overlap flag will be ignored in sequence parallelism mode {self.sequence_parallelism_mode}"
84                     )
85         else:
86             if self.sequence_parallelism_mode:
87                 self.sequence_parallelism_mode = None
88                 warnings.warn(
89                     f"The sequence_parallelism_mode will be ignored when enable_sequence_parallelism is False"
90                 )
91             assert (
92                 not self.enable_sequence_overlap
93             ), f"enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True"
94 
95         # get the tensor parallel size
96         if not self.enable_tensor_parallelism:
97             self._tensor_parallel_size = 1
98         else:
99             self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)
100 
101         # get the sequence parallel size
102         if not self.enable_sequence_parallelism:
103             self._sequence_parallel_size = 1
104         else:
105             self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)
106 
107     def _turn_on_all_optimization(self):
108         """
109         Turn on all optimization.
110         """
111         # you can add all the optimization flag here
112         self.enable_fused_normalization = True
113         self.enable_flash_attention = True
114         self.enable_jit_fused = True
115         # This can cause non-in-place param sharding when used without ZeRO.
116         # It may also slow down training when seq len is small. Plz enable manually.
117         # self.enable_sequence_parallelism = True
118         # self.enable_sequence_overlap = True
119 
120     def _infer(self):
121         """
122         Set default params for inference.
123         """
124         # assert self.pipeline_stage_manager is None, "pipeline parallelism is not supported in inference for now"
125
```

Path: `colossalai/shardformer/shard/grad_ckpt_config.py`
Content:
```
1 from dataclasses import dataclass
2 from typing import List, Optional
3 
4 
5 @dataclass
6 class GradientCheckpointConfig:
7     gradient_checkpointing_ratio: float = 0.0
8 
9     def get_num_ckpt_layers(self, num_layers: int) -> int:
10         return int(self.gradient_checkpointing_ratio * num_layers)
11 
12 
13 @dataclass
14 class PipelineGradientCheckpointConfig(GradientCheckpointConfig):
15     r"""
16     The pipeline gradient config is designed to provide more flexibility for users to control gradient checkpoint in pipeline parallelism.
17     Combined with PipelineStageManager.set_distribution_config, user can fully control the distribution of layers and checkpointed layers in pipeline parallelism.
18     Refer to https://github.com/hpcaitech/ColossalAI/issues/5509 for more details.
19 
20     It provides the following features:
21         1. `gradient_checkpointing_ratio`: This is used to control gradient checkpointing more precisely, e.g., set 50% of the layers to use gradient checkpointing.
22         2. Customize # ckpt layers assigned to each stage. This takes precedence over `gradient_checkpointing_ratio`.
23 
24     """
25     """
26     Args:
27         gradient_checkpointing_ratio (Optional[float]): The ratio of gradient checkpointing. It can only be used in pipeline parallelism. Defaults to None.
28         num_stages (Optional[int]): Number of stages in the pipeline. Defaults to None. For sanity check.
29         num_model_chunks (Optional[int]): Number of model chunks (1F1B or Interleaved). Defaults to None. For sanity check.
30         num_model_layers (Optional[int]): Number of model layers. Defaults to None. For sanity check.
31         num_ckpt_layers_per_stage (Optional[List[int]]): Number of checkpointed layers for each stage. Defaults to None.
32 
33     Example 1:
34         num_stages = 8
35         num_layers = 80
36         num_model_chunks = 1
37         num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]
38         num_ckpt_layers_per_stage = [4, 4, 2, 2, 0, 0, 0, 0]
39 
40     Example 2:
41         num_stages = 4
42         num_layers = 80
43         num_model_chunks = 2
44         num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]
45         # device 0 holds num_layers_per_stage[0] and num_layers_per_stage[4] layers
46         ...
47 
48     """
49     num_stages: Optional[int] = None
50     num_model_chunks: Optional[int] = None
51     num_model_layers: Optional[int] = None
52     num_ckpt_layers_per_stage: Optional[List[int]] = None
53 
54     def __post_init__(self):
55         if self._enable_gradient_checkpointing_ratio:
56             if not (0 <= self.gradient_checkpointing_ratio <= 1):
57                 raise ValueError("gradient_checkpointing_ratio should be in 0% to 100%")
58 
59         if self._enable_customized_ckpt_layers_per_stage:
60             assert (
61                 self.num_stages is not None and self.num_model_chunks is not None and self.num_model_layers is not None
62             )
63             assert len(self.num_ckpt_layers_per_stage) == self.num_stages * self.num_model_chunks
64             assert all(
65                 [0 <= num_ckpt_layers < self.num_model_layers for num_ckpt_layers in self.num_ckpt_layers_per_stage]
66             )
67             self.gradient_checkpointing_ratio = sum(self.num_ckpt_layers_per_stage) / self.num_model_layers
68 
69     @property
70     def _enable_gradient_checkpointing_ratio(self) -> bool:
71         return self.gradient_checkpointing_ratio is not None
72 
73     @property
74     def _enable_customized_ckpt_layers_per_stage(self) -> bool:
75         return self.num_ckpt_layers_per_stage is not None
76 
77     def get_num_ckpt_layers(self, stage: int, num_layers: int, model_chunk_id: int = 0) -> int:
78         if not self._enable_gradient_checkpointing_ratio and not self._enable_customized_ckpt_layers_per_stage:
79             raise RuntimeError("No checkpointed layers information is provided")
80 
81         if self._enable_customized_ckpt_layers_per_stage:
82             assert stage <= self.num_stages and model_chunk_id <= self.num_model_chunks
83             num_ckpt_layers = self.num_ckpt_layers_per_stage[stage + model_chunk_id * self.num_stages]
84             assert num_ckpt_layers <= num_layers
85             return num_ckpt_layers
86         else:
87             return int(self.gradient_checkpointing_ratio * num_layers)
88
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/shardformer/shard/grad_ckpt_config.py b/colossalai/shardformer/shard/grad_ckpt_config.py
--- a/colossalai/shardformer/shard/grad_ckpt_config.py
+++ b/colossalai/shardformer/shard/grad_ckpt_config.py
@@ -22,6 +22,7 @@
         2. Customize # ckpt layers assigned to each stage. This takes precedence over `gradient_checkpointing_ratio`.
 
     """
+
     """
     Args:
         gradient_checkpointing_ratio (Optional[float]): The ratio of gradient checkpointing. It can only be used in pipeline parallelism. Defaults to None.
@@ -49,6 +50,7 @@
     num_stages: Optional[int] = None
     num_model_chunks: Optional[int] = None
     num_model_layers: Optional[int] = None
+    num_layers_per_stage: Optional[List[int]] = None
     num_ckpt_layers_per_stage: Optional[List[int]] = None
 
     def __post_init__(self):
@@ -70,6 +72,10 @@
     def _enable_gradient_checkpointing_ratio(self) -> bool:
         return self.gradient_checkpointing_ratio is not None
 
+    @property
+    def _customize_num_layers_per_stage(self) -> bool:
+        return self.num_layers_per_stage is not None and self.num_model_layers is not None
+
     @property
     def _enable_customized_ckpt_layers_per_stage(self) -> bool:
         return self.num_ckpt_layers_per_stage is not None
diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py
--- a/colossalai/shardformer/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -7,7 +7,7 @@
 
 from colossalai.pipeline.stage_manager import PipelineStageManager
 
-from .grad_ckpt_config import GradientCheckpointConfig
+from .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig
 
 __all__ = ["ShardConfig"]
 SUPPORT_SP_MODE = ["split_gather", "ring", "all_to_all"]
@@ -30,6 +30,7 @@
         gradient_checkpoint_config (Optional[GradientCheckpointConfig]): The gradient checkpoint config. Defaults to None.
         enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.
     """
+
     tensor_parallel_process_group: Optional[ProcessGroup] = None
     sequence_parallel_process_group: Optional[ProcessGroup] = None
     pipeline_stage_manager: Optional[PipelineStageManager] = None
@@ -104,6 +105,16 @@
         else:
             self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)
 
+        if (
+            self.pipeline_stage_manager is not None
+            and isinstance(self.gradient_checkpoint_config, PipelineGradientCheckpointConfig)
+            and self.gradient_checkpoint_config._customize_num_layers_per_stage
+        ):
+            self.pipeline_stage_manager.set_distribution_config(
+                self.gradient_checkpoint_config.num_model_layers,
+                self.gradient_checkpoint_config.num_layers_per_stage,
+            )
+
     def _turn_on_all_optimization(self):
         """
         Turn on all optimization.

{"golden_diff": "diff --git a/colossalai/shardformer/shard/grad_ckpt_config.py b/colossalai/shardformer/shard/grad_ckpt_config.py\n--- a/colossalai/shardformer/shard/grad_ckpt_config.py\n+++ b/colossalai/shardformer/shard/grad_ckpt_config.py\n@@ -22,6 +22,7 @@\n         2. Customize # ckpt layers assigned to each stage. This takes precedence over `gradient_checkpointing_ratio`.\n \n     \"\"\"\n+\n     \"\"\"\n     Args:\n         gradient_checkpointing_ratio (Optional[float]): The ratio of gradient checkpointing. It can only be used in pipeline parallelism. Defaults to None.\n@@ -49,6 +50,7 @@\n     num_stages: Optional[int] = None\n     num_model_chunks: Optional[int] = None\n     num_model_layers: Optional[int] = None\n+    num_layers_per_stage: Optional[List[int]] = None\n     num_ckpt_layers_per_stage: Optional[List[int]] = None\n \n     def __post_init__(self):\n@@ -70,6 +72,10 @@\n     def _enable_gradient_checkpointing_ratio(self) -> bool:\n         return self.gradient_checkpointing_ratio is not None\n \n+    @property\n+    def _customize_num_layers_per_stage(self) -> bool:\n+        return self.num_layers_per_stage is not None and self.num_model_layers is not None\n+\n     @property\n     def _enable_customized_ckpt_layers_per_stage(self) -> bool:\n         return self.num_ckpt_layers_per_stage is not None\ndiff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py\n--- a/colossalai/shardformer/shard/shard_config.py\n+++ b/colossalai/shardformer/shard/shard_config.py\n@@ -7,7 +7,7 @@\n \n from colossalai.pipeline.stage_manager import PipelineStageManager\n \n-from .grad_ckpt_config import GradientCheckpointConfig\n+from .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig\n \n __all__ = [\"ShardConfig\"]\n SUPPORT_SP_MODE = [\"split_gather\", \"ring\", \"all_to_all\"]\n@@ -30,6 +30,7 @@\n         gradient_checkpoint_config (Optional[GradientCheckpointConfig]): The gradient checkpoint config. Defaults to None.\n         enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.\n     \"\"\"\n+\n     tensor_parallel_process_group: Optional[ProcessGroup] = None\n     sequence_parallel_process_group: Optional[ProcessGroup] = None\n     pipeline_stage_manager: Optional[PipelineStageManager] = None\n@@ -104,6 +105,16 @@\n         else:\n             self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)\n \n+        if (\n+            self.pipeline_stage_manager is not None\n+            and isinstance(self.gradient_checkpoint_config, PipelineGradientCheckpointConfig)\n+            and self.gradient_checkpoint_config._customize_num_layers_per_stage\n+        ):\n+            self.pipeline_stage_manager.set_distribution_config(\n+                self.gradient_checkpoint_config.num_model_layers,\n+                self.gradient_checkpoint_config.num_layers_per_stage,\n+            )\n+\n     def _turn_on_all_optimization(self):\n         \"\"\"\n         Turn on all optimization.\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import warnings\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, Optional\n\nimport torch.distributed as dist\nfrom torch.distributed import ProcessGroup\n\nfrom colossalai.pipeline.stage_manager import PipelineStageManager\n\nfrom .grad_ckpt_config import GradientCheckpointConfig\n\n__all__ = [\"ShardConfig\"]\nSUPPORT_SP_MODE = [\"split_gather\", \"ring\", \"all_to_all\"]\n\n\n@dataclass\nclass ShardConfig:\n    r\"\"\"\n    The config for sharding the huggingface model\n\n    Args:\n        tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.\n        pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.\n        enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.\n        enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.\n        enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.\n        enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.\n        enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.\n        enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.\n        gradient_checkpoint_config (Optional[GradientCheckpointConfig]): The gradient checkpoint config. Defaults to None.\n        enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.\n    \"\"\"\n    tensor_parallel_process_group: Optional[ProcessGroup] = None\n    sequence_parallel_process_group: Optional[ProcessGroup] = None\n    pipeline_stage_manager: Optional[PipelineStageManager] = None\n    enable_tensor_parallelism: bool = True\n    enable_all_optimization: bool = False\n    enable_fused_normalization: bool = False\n    enable_flash_attention: bool = False\n    enable_jit_fused: bool = False\n    enable_sequence_parallelism: bool = False\n    sequence_parallelism_mode: str = None\n    enable_sequence_overlap: bool = False\n    parallel_output: bool = True\n    make_vocab_size_divisible_by: int = 64\n    gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None\n    extra_kwargs: Dict[str, Any] = field(default_factory=dict)\n    # pipeline_parallel_size: int\n    # data_parallel_size: int\n    # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']\n\n    @property\n    def tensor_parallel_size(self):\n        return self._tensor_parallel_size\n\n    @property\n    def sequence_parallel_size(self):\n        return self._sequence_parallel_size\n\n    def __post_init__(self):\n        # turn on all optimization if all_optimization is set to True\n        if self.enable_all_optimization:\n            self._turn_on_all_optimization()\n\n        if self.enable_sequence_parallelism:\n            self.sequence_parallelism_mode = (\n                \"split_gather\" if self.sequence_parallelism_mode is None else self.sequence_parallelism_mode\n            )\n            assert (\n                self.sequence_parallelism_mode in SUPPORT_SP_MODE\n            ), f\"Sequence parallelism mode {self.sequence_parallelism_mode} is not in the supported list {SUPPORT_SP_MODE}\"\n            if self.sequence_parallelism_mode in [\"split_gather\", \"ring\"]:\n                assert (\n                    self.enable_tensor_parallelism\n                ), f\"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is True\"\n            elif self.sequence_parallelism_mode in [\"all_to_all\"]:\n                assert (\n                    not self.enable_tensor_parallelism\n                ), f\"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is False\"\n                if self.enable_sequence_overlap:\n                    self.enable_sequence_overlap = False\n                    warnings.warn(\n                        f\"The enable_sequence_overlap flag will be ignored in sequence parallelism mode {self.sequence_parallelism_mode}\"\n                    )\n        else:\n            if self.sequence_parallelism_mode:\n                self.sequence_parallelism_mode = None\n                warnings.warn(\n                    f\"The sequence_parallelism_mode will be ignored when enable_sequence_parallelism is False\"\n                )\n            assert (\n                not self.enable_sequence_overlap\n            ), f\"enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True\"\n\n        # get the tensor parallel size\n        if not self.enable_tensor_parallelism:\n            self._tensor_parallel_size = 1\n        else:\n            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)\n\n        # get the sequence parallel size\n        if not self.enable_sequence_parallelism:\n            self._sequence_parallel_size = 1\n        else:\n            self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)\n\n    def _turn_on_all_optimization(self):\n        \"\"\"\n        Turn on all optimization.\n        \"\"\"\n        # you can add all the optimization flag here\n        self.enable_fused_normalization = True\n        self.enable_flash_attention = True\n        self.enable_jit_fused = True\n        # This can cause non-in-place param sharding when used without ZeRO.\n        # It may also slow down training when seq len is small. Plz enable manually.\n        # self.enable_sequence_parallelism = True\n        # self.enable_sequence_overlap = True\n\n    def _infer(self):\n        \"\"\"\n        Set default params for inference.\n        \"\"\"\n        # assert self.pipeline_stage_manager is None, \"pipeline parallelism is not supported in inference for now\"\n", "path": "colossalai/shardformer/shard/shard_config.py"}, {"content": "from dataclasses import dataclass\nfrom typing import List, Optional\n\n\n@dataclass\nclass GradientCheckpointConfig:\n    gradient_checkpointing_ratio: float = 0.0\n\n    def get_num_ckpt_layers(self, num_layers: int) -> int:\n        return int(self.gradient_checkpointing_ratio * num_layers)\n\n\n@dataclass\nclass PipelineGradientCheckpointConfig(GradientCheckpointConfig):\n    r\"\"\"\n    The pipeline gradient config is designed to provide more flexibility for users to control gradient checkpoint in pipeline parallelism.\n    Combined with PipelineStageManager.set_distribution_config, user can fully control the distribution of layers and checkpointed layers in pipeline parallelism.\n    Refer to https://github.com/hpcaitech/ColossalAI/issues/5509 for more details.\n\n    It provides the following features:\n        1. `gradient_checkpointing_ratio`: This is used to control gradient checkpointing more precisely, e.g., set 50% of the layers to use gradient checkpointing.\n        2. Customize # ckpt layers assigned to each stage. This takes precedence over `gradient_checkpointing_ratio`.\n\n    \"\"\"\n    \"\"\"\n    Args:\n        gradient_checkpointing_ratio (Optional[float]): The ratio of gradient checkpointing. It can only be used in pipeline parallelism. Defaults to None.\n        num_stages (Optional[int]): Number of stages in the pipeline. Defaults to None. For sanity check.\n        num_model_chunks (Optional[int]): Number of model chunks (1F1B or Interleaved). Defaults to None. For sanity check.\n        num_model_layers (Optional[int]): Number of model layers. Defaults to None. For sanity check.\n        num_ckpt_layers_per_stage (Optional[List[int]]): Number of checkpointed layers for each stage. Defaults to None.\n\n    Example 1:\n        num_stages = 8\n        num_layers = 80\n        num_model_chunks = 1\n        num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]\n        num_ckpt_layers_per_stage = [4, 4, 2, 2, 0, 0, 0, 0]\n\n    Example 2:\n        num_stages = 4\n        num_layers = 80\n        num_model_chunks = 2\n        num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]\n        # device 0 holds num_layers_per_stage[0] and num_layers_per_stage[4] layers\n        ...\n\n    \"\"\"\n    num_stages: Optional[int] = None\n    num_model_chunks: Optional[int] = None\n    num_model_layers: Optional[int] = None\n    num_ckpt_layers_per_stage: Optional[List[int]] = None\n\n    def __post_init__(self):\n        if self._enable_gradient_checkpointing_ratio:\n            if not (0 <= self.gradient_checkpointing_ratio <= 1):\n                raise ValueError(\"gradient_checkpointing_ratio should be in 0% to 100%\")\n\n        if self._enable_customized_ckpt_layers_per_stage:\n            assert (\n                self.num_stages is not None and self.num_model_chunks is not None and self.num_model_layers is not None\n            )\n            assert len(self.num_ckpt_layers_per_stage) == self.num_stages * self.num_model_chunks\n            assert all(\n                [0 <= num_ckpt_layers < self.num_model_layers for num_ckpt_layers in self.num_ckpt_layers_per_stage]\n            )\n            self.gradient_checkpointing_ratio = sum(self.num_ckpt_layers_per_stage) / self.num_model_layers\n\n    @property\n    def _enable_gradient_checkpointing_ratio(self) -> bool:\n        return self.gradient_checkpointing_ratio is not None\n\n    @property\n    def _enable_customized_ckpt_layers_per_stage(self) -> bool:\n        return self.num_ckpt_layers_per_stage is not None\n\n    def get_num_ckpt_layers(self, stage: int, num_layers: int, model_chunk_id: int = 0) -> int:\n        if not self._enable_gradient_checkpointing_ratio and not self._enable_customized_ckpt_layers_per_stage:\n            raise RuntimeError(\"No checkpointed layers information is provided\")\n\n        if self._enable_customized_ckpt_layers_per_stage:\n            assert stage <= self.num_stages and model_chunk_id <= self.num_model_chunks\n            num_ckpt_layers = self.num_ckpt_layers_per_stage[stage + model_chunk_id * self.num_stages]\n            assert num_ckpt_layers <= num_layers\n            return num_ckpt_layers\n        else:\n            return int(self.gradient_checkpointing_ratio * num_layers)\n", "path": "colossalai/shardformer/shard/grad_ckpt_config.py"}], "after_files": [{"content": "import warnings\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, Optional\n\nimport torch.distributed as dist\nfrom torch.distributed import ProcessGroup\n\nfrom colossalai.pipeline.stage_manager import PipelineStageManager\n\nfrom .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig\n\n__all__ = [\"ShardConfig\"]\nSUPPORT_SP_MODE = [\"split_gather\", \"ring\", \"all_to_all\"]\n\n\n@dataclass\nclass ShardConfig:\n    r\"\"\"\n    The config for sharding the huggingface model\n\n    Args:\n        tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.\n        pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.\n        enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.\n        enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.\n        enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.\n        enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.\n        enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.\n        enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.\n        gradient_checkpoint_config (Optional[GradientCheckpointConfig]): The gradient checkpoint config. Defaults to None.\n        enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.\n    \"\"\"\n\n    tensor_parallel_process_group: Optional[ProcessGroup] = None\n    sequence_parallel_process_group: Optional[ProcessGroup] = None\n    pipeline_stage_manager: Optional[PipelineStageManager] = None\n    enable_tensor_parallelism: bool = True\n    enable_all_optimization: bool = False\n    enable_fused_normalization: bool = False\n    enable_flash_attention: bool = False\n    enable_jit_fused: bool = False\n    enable_sequence_parallelism: bool = False\n    sequence_parallelism_mode: str = None\n    enable_sequence_overlap: bool = False\n    parallel_output: bool = True\n    make_vocab_size_divisible_by: int = 64\n    gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None\n    extra_kwargs: Dict[str, Any] = field(default_factory=dict)\n    # pipeline_parallel_size: int\n    # data_parallel_size: int\n    # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']\n\n    @property\n    def tensor_parallel_size(self):\n        return self._tensor_parallel_size\n\n    @property\n    def sequence_parallel_size(self):\n        return self._sequence_parallel_size\n\n    def __post_init__(self):\n        # turn on all optimization if all_optimization is set to True\n        if self.enable_all_optimization:\n            self._turn_on_all_optimization()\n\n        if self.enable_sequence_parallelism:\n            self.sequence_parallelism_mode = (\n                \"split_gather\" if self.sequence_parallelism_mode is None else self.sequence_parallelism_mode\n            )\n            assert (\n                self.sequence_parallelism_mode in SUPPORT_SP_MODE\n            ), f\"Sequence parallelism mode {self.sequence_parallelism_mode} is not in the supported list {SUPPORT_SP_MODE}\"\n            if self.sequence_parallelism_mode in [\"split_gather\", \"ring\"]:\n                assert (\n                    self.enable_tensor_parallelism\n                ), f\"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is True\"\n            elif self.sequence_parallelism_mode in [\"all_to_all\"]:\n                assert (\n                    not self.enable_tensor_parallelism\n                ), f\"sequence parallelism mode {self.sequence_parallelism_mode} can only be used when enable_tensor_parallelism is False\"\n                if self.enable_sequence_overlap:\n                    self.enable_sequence_overlap = False\n                    warnings.warn(\n                        f\"The enable_sequence_overlap flag will be ignored in sequence parallelism mode {self.sequence_parallelism_mode}\"\n                    )\n        else:\n            if self.sequence_parallelism_mode:\n                self.sequence_parallelism_mode = None\n                warnings.warn(\n                    f\"The sequence_parallelism_mode will be ignored when enable_sequence_parallelism is False\"\n                )\n            assert (\n                not self.enable_sequence_overlap\n            ), f\"enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True\"\n\n        # get the tensor parallel size\n        if not self.enable_tensor_parallelism:\n            self._tensor_parallel_size = 1\n        else:\n            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)\n\n        # get the sequence parallel size\n        if not self.enable_sequence_parallelism:\n            self._sequence_parallel_size = 1\n        else:\n            self._sequence_parallel_size = dist.get_world_size(self.sequence_parallel_process_group)\n\n        if (\n            self.pipeline_stage_manager is not None\n            and isinstance(self.gradient_checkpoint_config, PipelineGradientCheckpointConfig)\n            and self.gradient_checkpoint_config._customize_num_layers_per_stage\n        ):\n            self.pipeline_stage_manager.set_distribution_config(\n                self.gradient_checkpoint_config.num_model_layers,\n                self.gradient_checkpoint_config.num_layers_per_stage,\n            )\n\n    def _turn_on_all_optimization(self):\n        \"\"\"\n        Turn on all optimization.\n        \"\"\"\n        # you can add all the optimization flag here\n        self.enable_fused_normalization = True\n        self.enable_flash_attention = True\n        self.enable_jit_fused = True\n        # This can cause non-in-place param sharding when used without ZeRO.\n        # It may also slow down training when seq len is small. Plz enable manually.\n        # self.enable_sequence_parallelism = True\n        # self.enable_sequence_overlap = True\n\n    def _infer(self):\n        \"\"\"\n        Set default params for inference.\n        \"\"\"\n        # assert self.pipeline_stage_manager is None, \"pipeline parallelism is not supported in inference for now\"\n", "path": "colossalai/shardformer/shard/shard_config.py"}, {"content": "from dataclasses import dataclass\nfrom typing import List, Optional\n\n\n@dataclass\nclass GradientCheckpointConfig:\n    gradient_checkpointing_ratio: float = 0.0\n\n    def get_num_ckpt_layers(self, num_layers: int) -> int:\n        return int(self.gradient_checkpointing_ratio * num_layers)\n\n\n@dataclass\nclass PipelineGradientCheckpointConfig(GradientCheckpointConfig):\n    r\"\"\"\n    The pipeline gradient config is designed to provide more flexibility for users to control gradient checkpoint in pipeline parallelism.\n    Combined with PipelineStageManager.set_distribution_config, user can fully control the distribution of layers and checkpointed layers in pipeline parallelism.\n    Refer to https://github.com/hpcaitech/ColossalAI/issues/5509 for more details.\n\n    It provides the following features:\n        1. `gradient_checkpointing_ratio`: This is used to control gradient checkpointing more precisely, e.g., set 50% of the layers to use gradient checkpointing.\n        2. Customize # ckpt layers assigned to each stage. This takes precedence over `gradient_checkpointing_ratio`.\n\n    \"\"\"\n\n    \"\"\"\n    Args:\n        gradient_checkpointing_ratio (Optional[float]): The ratio of gradient checkpointing. It can only be used in pipeline parallelism. Defaults to None.\n        num_stages (Optional[int]): Number of stages in the pipeline. Defaults to None. For sanity check.\n        num_model_chunks (Optional[int]): Number of model chunks (1F1B or Interleaved). Defaults to None. For sanity check.\n        num_model_layers (Optional[int]): Number of model layers. Defaults to None. For sanity check.\n        num_ckpt_layers_per_stage (Optional[List[int]]): Number of checkpointed layers for each stage. Defaults to None.\n\n    Example 1:\n        num_stages = 8\n        num_layers = 80\n        num_model_chunks = 1\n        num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]\n        num_ckpt_layers_per_stage = [4, 4, 2, 2, 0, 0, 0, 0]\n\n    Example 2:\n        num_stages = 4\n        num_layers = 80\n        num_model_chunks = 2\n        num_layers_per_stage = [9, 9, 9, 10, 11, 10, 11, 11]\n        # device 0 holds num_layers_per_stage[0] and num_layers_per_stage[4] layers\n        ...\n\n    \"\"\"\n    num_stages: Optional[int] = None\n    num_model_chunks: Optional[int] = None\n    num_model_layers: Optional[int] = None\n    num_layers_per_stage: Optional[List[int]] = None\n    num_ckpt_layers_per_stage: Optional[List[int]] = None\n\n    def __post_init__(self):\n        if self._enable_gradient_checkpointing_ratio:\n            if not (0 <= self.gradient_checkpointing_ratio <= 1):\n                raise ValueError(\"gradient_checkpointing_ratio should be in 0% to 100%\")\n\n        if self._enable_customized_ckpt_layers_per_stage:\n            assert (\n                self.num_stages is not None and self.num_model_chunks is not None and self.num_model_layers is not None\n            )\n            assert len(self.num_ckpt_layers_per_stage) == self.num_stages * self.num_model_chunks\n            assert all(\n                [0 <= num_ckpt_layers < self.num_model_layers for num_ckpt_layers in self.num_ckpt_layers_per_stage]\n            )\n            self.gradient_checkpointing_ratio = sum(self.num_ckpt_layers_per_stage) / self.num_model_layers\n\n    @property\n    def _enable_gradient_checkpointing_ratio(self) -> bool:\n        return self.gradient_checkpointing_ratio is not None\n\n    @property\n    def _customize_num_layers_per_stage(self) -> bool:\n        return self.num_layers_per_stage is not None and self.num_model_layers is not None\n\n    @property\n    def _enable_customized_ckpt_layers_per_stage(self) -> bool:\n        return self.num_ckpt_layers_per_stage is not None\n\n    def get_num_ckpt_layers(self, stage: int, num_layers: int, model_chunk_id: int = 0) -> int:\n        if not self._enable_gradient_checkpointing_ratio and not self._enable_customized_ckpt_layers_per_stage:\n            raise RuntimeError(\"No checkpointed layers information is provided\")\n\n        if self._enable_customized_ckpt_layers_per_stage:\n            assert stage <= self.num_stages and model_chunk_id <= self.num_model_chunks\n            num_ckpt_layers = self.num_ckpt_layers_per_stage[stage + model_chunk_id * self.num_stages]\n            assert num_ckpt_layers <= num_layers\n            return num_ckpt_layers\n        else:\n            return int(self.gradient_checkpointing_ratio * num_layers)\n", "path": "colossalai/shardformer/shard/grad_ckpt_config.py"}]}