Fix PP issue

huggingface · Sep 24, 2024 · b0cc1d1 · b0cc1d1
1 parent 3748a06
commit b0cc1d1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/neuron/distributed/decoder_models.py b/optimum/neuron/distributed/decoder_models.py
@@ -29,6 +29,7 @@
     LlamaDecoderLayer,
     LlamaForQuestionAnswering,
     LlamaRMSNorm,
+    LlamaRotaryEmbedding,
     repeat_kv,
 )
 from transformers.models.mistral.modeling_mistral import (
@@ -554,7 +555,7 @@ class LlamaPipelineParallelismSpecs(PipelineParallelismSpecs):
         "LlamaForQuestionAnswering": ("input_ids", "attention_mask", "start_positions", "end_positions"),
     }
 
-    LEAF_MODULE_CLASSES_NAMES = [LlamaRMSNorm]
+    LEAF_MODULE_CLASSES_NAMES = [LlamaRMSNorm, LlamaRotaryEmbedding]
 
 
 class LlamaParallelizer(Parallelizer):