tests passed

hpcaitech · Apr 30, 2024 · ff819a5 · ff819a5
1 parent 85ba7ab
commit ff819a5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py
@@ -381,7 +381,6 @@ def configure(
                 shard_config=self.shard_config,
                 dp_group=self.global_dp_group,
                 tp_group=self.tp_group,
-                sp_group=self.sp_group,  # TODO: add ep group. Modify shard_config to assign pg in policy
                 use_ddp=use_ddp,
                 ddp_config=self.ddp_config,
                 custom_policy=self.custom_policy,

diff --git a/colossalai/shardformer/policies/mixtral.py b/colossalai/shardformer/policies/mixtral.py
@@ -111,8 +111,8 @@ def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, poli
             else:
                 module = self.model.model
 
-            layers_per_stage = stage_manager.distribute_layers(len(module.layers))
-            stage_index = stage_manager.get_stage_index(layers_per_stage)
+            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)
             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
@@ -131,10 +131,10 @@ def get_held_layers(self) -> List[Module]:
         stage_manager = self.pipeline_stage_manager
 
         held_layers = []
-        layers_per_stage = stage_manager.distribute_layers(len(module.layers))
+        layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)
         if stage_manager.is_first_stage():
             held_layers.append(module.embed_tokens)
-        start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage)
+        start_idx, end_idx = self.get_stage_index(layers_per_stage, stage_manager.stage)
         held_layers.extend(module.layers[start_idx:end_idx])
         if stage_manager.is_last_stage():
             held_layers.append(module.norm)