[fix] fix llama modeling policy;

2025-09-08 20:40:34 +00:00 · 2024-10-25 10:17:06 +00:00
parent cc0dfddcbc
commit 03fa79a55c
2 changed files with 3 additions and 1 deletions
--- a/colossalai/shardformer/policies/llama.py
+++ b/colossalai/shardformer/policies/llama.py
@@ -96,7 +96,8 @@ class LlamaPolicy(Policy):
                target_key=attn_cls,
            )

-        if self.pipeline_stage_manager is not None:
+        # if self.pipeline_stage_manager is not None:
+        if self.pipeline_stage_manager is None:
            self.append_or_create_method_replacement(
                description={
                    "forward": partial(