feat: Use self.[distribute_layers|get_stage_index] to exploit custom layer distribution

Author: cwher

applications/ColossalMoE/colossal_moe/models/mixtral_policy.py

@@ -110,7 +110,7 @@ def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, poli
                 module = self.model.model
 
             layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)
-            stage_index = Policy.get_stage_index(layers_per_stage, stage_manager.stage)
+            stage_index = self.get_stage_index(layers_per_stage, stage_manager.stage)
             method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)}
             self.append_or_create_method_replacement(
                 description=method_replacement, policy=policy, target_key=model_cls

colossalai/shardformer/policies/base_policy.py

@@ -197,8 +197,7 @@ def get_shared_params(self) -> List[Dict[int, Tensor]]:
         """
         return []
 
-    @staticmethod
-    def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
+    def distribute_layers(self, num_layers: int, num_stages: int) -> List[int]:
         """Divide layers into stages"""
         quotient = num_layers // num_stages
         remainder = num_layers % num_stages
@@ -213,8 +212,8 @@ def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
                 layers_per_stage[i] += 1
         return layers_per_stage
 
-    @staticmethod
     def get_stage_index(
+        self,
         layers_per_stage: List[int],
         stage: int,
         num_model_chunks: int = 1,

colossalai/shardformer/policies/gpt2.py

@@ -175,7 +175,7 @@ def get_held_layers(self) -> List[nn.Module]:
             layers_per_stage = self.distribute_layers(
                 len(module.h), stage_manager.num_stages * stage_manager.num_model_chunks
             )
-            stage_indices = Policy.get_stage_index(
+            stage_indices = self.get_stage_index(
                 layers_per_stage,
                 stage_manager.stage,
                 num_model_chunks=stage_manager.num_model_chunks,
@@ -226,8 +226,8 @@ def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, poli
                 "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config)
             }
         else:
-            layers_per_stage = Policy.distribute_layers(len(module.h), stage_manager.num_stages)
-            stage_index = Policy.get_stage_index(layers_per_stage, stage_manager.stage)
+            layers_per_stage = self.distribute_layers(len(module.h), stage_manager.num_stages)
+            stage_index = self.get_stage_index(layers_per_stage, stage_manager.stage)
             method_replacement = {
                 "forward": partial(
                     new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config

examples/language/openmoe/model/openmoe_policy.py

@@ -98,11 +98,11 @@ def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, poli
                 module = self.model.model
 
             layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)
-            stage_index = Policy.get_stage_index(layers_per_stage, stage_manager.stage)
+            stage_index = self.get_stage_index(layers_per_stage, stage_manager.stage)
             method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)}
-            self.append_or_create_method_replacement(description=method_replacement,
-                                                     policy=policy,
-                                                     target_key=model_cls)
+            self.append_or_create_method_replacement(
+                description=method_replacement, policy=policy, target_key=model_cls
+            )
 
         return
 
@@ -126,12 +126,9 @@ def get_held_layers(self) -> List[Module]:
             held_layers.append(module.norm)
 
         return held_layers
-    
-    @staticmethod
-    def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
-        """Divide layers into stages
 
-        """
+    def distribute_layers(self, num_layers: int, num_stages: int) -> List[int]:
+        """Divide layers into stages"""
         if num_layers == 24 and num_stages == 4:
             return [7, 7, 7, 3]
         elif num_layers == 24 and num_stages == 2:
@@ -142,7 +139,7 @@ def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
             return [8, 4]
         else:
             print(f"num_layers: {num_layers}, num_stages: {num_stages} not optimized, use origin pp policy")
-            return Policy.distribute_layers(num_layers, num_stages)
+            return super().distribute_layers(num_layers, num_stages)
 
 
 class OpenMoeModelPolicy(OpenMoePolicy):