fix: add optional args for and

Author: cwher

colossalai/pipeline/stage_manager.py

@@ -90,50 +90,66 @@ def set_distribution_config(self, num_model_layers: int, num_layers_per_stage: L
         self.num_model_layers = num_model_layers
         self.num_layers_per_stage = num_layers_per_stage
 
-    def distribute_layers(self, num_layers: int) -> List[int]:
+    def distribute_layers(
+        self, num_layers: int, num_stages: Optional[int] = None, num_model_chunks: Optional[int] = None
+    ) -> List[int]:
         """Divide layers into stages"""
+        num_stages = self.num_stages if num_stages is None else num_stages
+        num_model_chunks = (
+            (self.num_model_chunks if self.is_interleave else 1) if num_model_chunks is None else num_model_chunks
+        )
+
         if self.control_distribute_layers:
             assert num_layers == self.num_model_layers
             return self.num_layers_per_stage
 
         else:
-            num_model_chunk = self.num_model_chunks if self.is_interleave else 1
-            quotient = num_layers // (self.num_stages * num_model_chunk)
-            remainder = num_layers % (self.num_stages * num_model_chunk)
+            quotient = num_layers // (num_stages * num_model_chunks)
+            remainder = num_layers % (num_stages * num_model_chunks)
 
             # calculate the num_layers per stage
-            layers_per_stage = [quotient] * self.num_stages * num_model_chunk
+            layers_per_stage = [quotient] * num_stages * num_model_chunks
 
             # deal with the rest layers
             if remainder > 0:
-                start_position = (self.num_stages * num_model_chunk) // 2 - remainder // 2
+                start_position = (num_stages * num_model_chunks) // 2 - remainder // 2
                 for i in range(start_position, start_position + remainder):
                     layers_per_stage[i] += 1
             return layers_per_stage
 
     def get_stage_index(
         self,
         layers_per_stage: List[int],
+        stage: Optional[int] = None,
+        num_model_chunks: Optional[int] = None,
+        num_stages: Optional[int] = None,
     ) -> Union[Tuple[int, int], List[Tuple[int, int]]]:
         """
         Get the start index and end index of layers for each stage.
 
         Args:
             layers_per_stage (List[int]): number of layers for each stage
             stage (int): the stage index
+            num_stages (int): number of stages
+            num_model_chunks (int): number of model chunks
 
         Returns:
             - Tuple[int, int]: the start index and end index of this stage
             - List[Tuple[int, int]]: the start index and end index of this stage for each model chunk
 
         """
+        stage = self.stage if stage is None else stage
+        num_model_chunks = (
+            (self.num_model_chunks if self.is_interleave else 1) if num_model_chunks is None else num_model_chunks
+        )
+        num_stages = self.num_stages if num_stages is None else num_stages
+
         num_layers_per_stage_accumulated = np.insert(np.cumsum(layers_per_stage), 0, 0)
 
         stage_indices = []
-        num_model_chunks = self.num_model_chunks if self.is_interleave else 1
         for model_chunk in range(num_model_chunks):
-            start_idx = num_layers_per_stage_accumulated[self.stage + model_chunk * self.num_stages]
-            end_idx = num_layers_per_stage_accumulated[self.stage + model_chunk * self.num_stages + 1]
+            start_idx = num_layers_per_stage_accumulated[stage + model_chunk * num_stages]
+            end_idx = num_layers_per_stage_accumulated[stage + model_chunk * num_stages + 1]
             stage_indices.append([start_idx, end_idx])
 
         return stage_indices[0] if num_model_chunks == 1 else stage_indices

colossalai/shardformer/policies/t5.py

@@ -251,6 +251,8 @@ def distribute_t5_layers(
         Return the layer distribution as a list and the starting stage of decoder.
         If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers.
         """
+        stage_manager = self.pipeline_stage_manager
+        assert stage_manager is not None, "Pipeline stage manager is not set."
 
         # number of encoder layers must be a positive integer
         if num_encoder_layers <= 0:
@@ -262,7 +264,7 @@ def distribute_t5_layers(
 
         # in the case of T5EncoderModel, set decoder starting stage to num_stages since it doesn't exist
         if num_decoder_layers == 0:
-            return self.distribute_layers(num_encoder_layers, num_stages), num_stages
+            return stage_manager.distribute_layers(num_encoder_layers, num_stages), num_stages
 
         # the number of stages distributed between encoder and decoder is optimized in this way:
         # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages))
@@ -273,21 +275,26 @@ def objective(num_encoder_stages):
         num_encoder_stages = np.argmin([objective(i) for i in range(1, num_stages)]) + 1
         num_decoder_stages = num_stages - num_encoder_stages
 
-        encoder_distribution = self.distribute_layers(num_encoder_layers, num_encoder_stages)
-        decoder_distribution = self.distribute_layers(num_decoder_layers, num_decoder_stages)
+        encoder_distribution = stage_manager.distribute_layers(num_encoder_layers, num_encoder_stages)
+        decoder_distribution = stage_manager.distribute_layers(num_decoder_layers, num_decoder_stages)
         return encoder_distribution + decoder_distribution, num_encoder_stages
 
     def get_t5_stage_index(
         self, layers_per_stage: List[int], stage: int, decoder_starting_stage: int
-    ) -> Tuple[bool, int, int]:
+    ) -> Tuple[int, int]:
         """
         Input the distribution of layers among stages, the current stage and the first stage of decoder.
         Return the starting/ending idx of layers in encoder/decoder
         """
+        stage_manager = self.pipeline_stage_manager
+        assert stage_manager is not None, "Pipeline stage manager is not set."
+
         if stage < decoder_starting_stage:
-            return self.get_stage_index(layers_per_stage[:decoder_starting_stage], stage)
+            return stage_manager.get_stage_index(layers_per_stage[:decoder_starting_stage], stage)
         else:
-            return self.get_stage_index(layers_per_stage[decoder_starting_stage:], stage - decoder_starting_stage)
+            return stage_manager.get_stage_index(
+                layers_per_stage[decoder_starting_stage:], stage - decoder_starting_stage
+            )
 
     def get_held_layers(self) -> List[nn.Module]:
         """Get pipeline layers for current stage."""

colossalai/shardformer/policies/whisper.py

@@ -300,6 +300,8 @@ def distribute_whisper_layers(
         Return the layer distribution as a list and the starting stage of decoder.
         If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers.
         """
+        stage_manager = self.pipeline_stage_manager
+        assert stage_manager is not None, "pipeline_stage_manager is None"
 
         # number of encoder layers must be a positive integer
         if num_encoder_layers <= 0:
@@ -311,7 +313,7 @@ def distribute_whisper_layers(
 
         # in the case of whisperEncoderModel, set decoder starting stage to num_stages since it doesn't exist
         if num_decoder_layers == 0:
-            return self.distribute_layers(num_encoder_layers, num_stages), num_stages
+            return stage_manager.distribute_layers(num_encoder_layers, num_stages), num_stages
 
         # the number of stages distributed between encoder and decoder is optimized in this way:
         # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages))
@@ -322,21 +324,24 @@ def objective(num_encoder_stages):
         num_encoder_stages = np.argmin([objective(i) for i in range(1, num_stages)]) + 1
         num_decoder_stages = num_stages - num_encoder_stages
 
-        encoder_distribution = self.distribute_layers(num_encoder_layers, num_encoder_stages)
-        decoder_distribution = self.distribute_layers(num_decoder_layers, num_decoder_stages)
+        encoder_distribution = stage_manager.distribute_layers(num_encoder_layers, num_encoder_stages)
+        decoder_distribution = stage_manager.distribute_layers(num_decoder_layers, num_decoder_stages)
         return encoder_distribution + decoder_distribution, num_encoder_stages
 
     def get_whisper_stage_index(
         self, layers_per_stage: List[int], stage: int, decoder_starting_stage: int
-    ) -> Tuple[bool, int, int]:
+    ) -> Tuple[int, int]:
         """
         Input the distribution of layers among stages, the current stage and the first stage of decoder.
         Return the starting/ending idx of layers in encoder/decoder
         """
+        stage_manager = self.pipeline_stage_manager
+        assert stage_manager is not None, "pipeline_stage_manager is None"
+
         if stage < decoder_starting_stage:
-            return self.get_stage_index(layers_per_stage[:decoder_starting_stage], stage)
+            return stage_manager.get_stage_index(layers_per_stage[:decoder_starting_stage], stage)
         else:
-            return self.get_stage_index(
+            return stage_manager.get_stage_index(
                 layers_per_stage[decoder_starting_stage:],
                 stage - decoder_starting_stage,
             )