Merge pull request #6322 from wangbluo/fix_falcon

fix falcon

Author: BurkeHulk

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))