From 4a077e5dc3d8c6eca6595ff2722bcda9ff9d8dad Mon Sep 17 00:00:00 2001 From: wangbluo <2538539015@qq.com> Date: Thu, 22 May 2025 16:50:40 +0800 Subject: [PATCH] fix falcon --- colossalai/shardformer/policies/falcon.py | 1 + 1 file changed, 1 insertion(+) diff --git a/colossalai/shardformer/policies/falcon.py b/colossalai/shardformer/policies/falcon.py index 68a548aee869..362f33176a00 100644 --- a/colossalai/shardformer/policies/falcon.py +++ b/colossalai/shardformer/policies/falcon.py @@ -246,6 +246,7 @@ def get_held_layers(self) -> List[Module]: module = self.model.transformer stage_manager = self.pipeline_stage_manager held_layers = [] + held_layers.append(module.rotary_emb) if stage_manager.is_interleave: assert stage_manager.num_model_chunks is not None layers_per_stage = stage_manager.distribute_layers(len(module.h))