fix: remove GradCkptCollection

Author: cwher

colossalai/booster/plugin/hybrid_parallel_plugin.py

@@ -26,7 +26,7 @@
 from colossalai.interface import AMPModelMixin, ModelWrapper, OptimizerWrapper
 from colossalai.pipeline.schedule import InterleavedSchedule, OneForwardOneBackwardSchedule
 from colossalai.pipeline.stage_manager import PipelineStageManager
-from colossalai.shardformer import GradCkptCollection, ShardConfig, ShardFormer
+from colossalai.shardformer import GradientCheckpointConfig, ShardConfig, ShardFormer
 from colossalai.shardformer.layer.utils import SeqParallelUtils
 from colossalai.shardformer.policies.base_policy import Policy
 from colossalai.tensor.d_tensor.api import is_distributed_tensor
@@ -930,7 +930,7 @@ class HybridParallelPlugin(PipelinePluginBase):
         custom_policy (Policy, optional): Custom policy for Shardformer. Defaults to None.
         pp_style (str, optional): The style for pipeline parallelism. Defaults to '1f1b'.
         num_model_chunks (int, optional): The number of model chunks for interleaved pipeline parallelism. Defaults to 1.
-        gradient_ckpt_collection (GradCkptCollection, optional): The configuration for gradient checkpointing. Defaults to None.
+        gradient_checkpoint_config (GradientCheckpointConfig, optional): Configuration for gradient checkpointing. Defaults to None.
         enable_metadata_cache (bool, optional): Whether to enable metadata cache for pipeline parallelism. Defaults to True.
     """
 
@@ -970,7 +970,7 @@ def __init__(
         custom_policy: Policy = None,
         pp_style: str = "1f1b",
         num_model_chunks: int = 1,
-        gradient_ckpt_collection: Optional[GradCkptCollection] = None,
+        gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None,
         enable_metadata_cache: bool = True,
     ) -> None:
         super().__init__()
@@ -1045,7 +1045,7 @@ def __init__(
             enable_sequence_parallelism=enable_sequence_parallelism,
             enable_sequence_overlap=enable_sequence_overlap,
             parallel_output=parallel_output,
-            gradient_ckpt_collection=gradient_ckpt_collection,
+            gradient_checkpoint_config=gradient_checkpoint_config,
         )
         self.amp_config = dict(
             initial_scale=initial_scale,

colossalai/shardformer/__init__.py

@@ -1 +1 @@
-from .shard import GradCkptCollection, ModelSharder, PipelineGradCkptConfig, ShardConfig, ShardFormer
+from .shard import GradientCheckpointConfig, ModelSharder, PipelineGradientCheckpointConfig, ShardConfig, ShardFormer

colossalai/shardformer/shard/__init__.py

@@ -1,6 +1,6 @@
-from .grad_ckpt_config import GradCkptCollection, PipelineGradCkptConfig
+from .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig
 from .shard_config import ShardConfig
 from .sharder import ModelSharder
 from .shardformer import ShardFormer
 
-__all__ = ["ShardConfig", "ModelSharder", "ShardFormer", "PipelineGradCkptConfig", "GradCkptCollection"]
+__all__ = ["ShardConfig", "ModelSharder", "ShardFormer", "PipelineGradientCheckpointConfig", "GradientCheckpointConfig"]

colossalai/shardformer/shard/grad_ckpt_config.py

@@ -1,43 +1,22 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import List, Optional
 
 
 @dataclass
-class GradCkptConfig:
+class GradientCheckpointConfig:
     # TODO: for future use
     _dummy_value: Optional[float] = None
 
-    def __post_init__(self):
-        raise NotImplementedError()
-
     @property
     def control_gradient_checkpointing(self) -> bool:
-        raise NotImplementedError()
+        return False
 
     def get_num_ckpt_layers(self, *args, **kwargs) -> int:
         raise NotImplementedError()
 
 
 @dataclass
-class GradCkptCollection:
-    gradient_ckpt_configs: List[GradCkptConfig] = field(default_factory=list)
-
-    def __post_init__(self):
-        assert all([isinstance(config, GradCkptConfig) for config in self.gradient_ckpt_configs])
-
-    @property
-    def control_gradient_checkpointing(self) -> bool:
-        return any([config.control_gradient_checkpointing for config in self.gradient_ckpt_configs])
-
-    def get_num_ckpt_layers(self, *args, **kwargs) -> int:
-        for config in self.gradient_ckpt_configs:
-            if config.control_gradient_checkpointing:
-                return config.get_num_ckpt_layers(*args, **kwargs)
-        raise RuntimeError("No checkpointed layers information is provided")
-
-
-@dataclass
-class PipelineGradCkptConfig(GradCkptConfig):
+class PipelineGradientCheckpointConfig(GradientCheckpointConfig):
     r"""
     The pipeline gradient config is designed to provide more flexibility for users to control gradient checkpoint in pipeline parallelism.
     Combined with PipelineStageManager.set_distribution_config, user can fully control the distribution of layers and checkpointed layers in pipeline parallelism.

colossalai/shardformer/shard/shard_config.py

@@ -6,7 +6,7 @@
 
 from colossalai.pipeline.stage_manager import PipelineStageManager
 
-from .grad_ckpt_config import GradCkptCollection
+from .grad_ckpt_config import GradientCheckpointConfig
 
 __all__ = ["ShardConfig"]
 
@@ -25,7 +25,7 @@ class ShardConfig:
         enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.
         enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.
         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.
-        gradient_ckpt_collection (Optional[GradCkptCollection]): The gradient checkpointing configs. Defaults to None.
+        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
@@ -38,7 +38,7 @@ class ShardConfig:
     enable_sequence_parallelism: bool = False
     enable_sequence_overlap: bool = False
     parallel_output: bool = True
-    gradient_ckpt_collection: Optional[GradCkptCollection] = None
+    gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None
     extra_kwargs: Dict[str, Any] = field(default_factory=dict)
     # TODO padding vocab
     # make_vocab_size_divisible_by: int = 128

tests/test_shardformer/test_model/test_shard_llama.py

@@ -5,7 +5,7 @@
 
 import colossalai
 from colossalai.logging import disable_existing_loggers
-from colossalai.shardformer import GradCkptCollection, PipelineGradCkptConfig
+from colossalai.shardformer import PipelineGradientCheckpointConfig
 from colossalai.shardformer.layer.utils import Randomizer
 from colossalai.tensor.d_tensor.api import clear_layout_converter
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
@@ -107,7 +107,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
             "precision": "fp16",
             "initial_scale": 1,
             "enable_gradient_checkpointing": True,
-            "gradient_ckpt_collection": GradCkptCollection([PipelineGradCkptConfig(gradient_checkpointing_ratio=0.5)]),
+            "gradient_checkpoint_config": PipelineGradientCheckpointConfig(gradient_checkpointing_ratio=0.5),
         },
         {
             "tp_size": 1,
@@ -116,12 +116,8 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
             "use_lazy_init": False,
             "precision": "fp32",
             "enable_gradient_checkpointing": True,
-            "gradient_ckpt_collection": GradCkptCollection(
-                [
-                    PipelineGradCkptConfig(
-                        num_stages=2, num_model_chunks=1, num_model_layers=8, num_ckpt_layers_per_stage=[4, 0]
-                    )
-                ]
+            "gradient_checkpoint_config": PipelineGradientCheckpointConfig(
+                num_stages=2, num_model_chunks=1, num_model_layers=8, num_ckpt_layers_per_stage=[4, 0]
             ),
         },
         {
@@ -205,15 +201,11 @@ def run_llama_test(test_config):
             "zero_stage": 1,
             "initial_scale": 1,
             "enable_gradient_checkpointing": True,
-            "gradient_ckpt_collection": GradCkptCollection(
-                [
-                    PipelineGradCkptConfig(
-                        num_stages=2,
-                        num_model_chunks=2,
-                        num_model_layers=8,
-                        num_ckpt_layers_per_stage=[0, 1, 2, 2],
-                    )
-                ]
+            "gradient_checkpoint_config": PipelineGradientCheckpointConfig(
+                num_stages=2,
+                num_model_chunks=2,
+                num_model_layers=8,
+                num_ckpt_layers_per_stage=[0, 1, 2, 2],
             ),
         },
     ],