Add sequence parallel support to Rope positional embedding (#6178)

* support sp Signed-off-by: Yi Dong <[email protected]> * use parallel_state Signed-off-by: Yi Dong <[email protected]> --------- Signed-off-by: Yi Dong <[email protected]>
NVIDIA · Mar 14, 2023 · 88b7a06 · 88b7a06
1 parent edb0812
commit 88b7a06
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/language_model.py b/nemo/collections/nlp/modules/common/megatron/language_model.py
@@ -27,7 +27,7 @@
 )
 
 try:
-    from apex.transformer import tensor_parallel
+    from apex.transformer import parallel_state, tensor_parallel
     from apex.transformer.enums import AttnMaskType
 
     HAVE_APEX = True
@@ -491,6 +491,7 @@ def __init__(
         self.output_layer_init_method = output_layer_init_method
         self.position_embedding_type = position_embedding_type
         self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+        self.sequence_parallel = sequence_parallel
 
         if kv_channels is None:
 
@@ -664,17 +665,25 @@ def forward(
         else:
             pass
 
-        # encoder_input: [s, b, h]
-
         # enc_attn_mask: [1, 1, s, s]
 
         if self.position_embedding_type == 'rope':
             if inference_max_sequence_len is not None:
                 rotary_pos_emb = self.rotary_pos_emb(inference_max_sequence_len)
             elif self.encoder.input_tensor is not None:
-                rotary_pos_emb = self.rotary_pos_emb(self.encoder.input_tensor.size(0))
+                if self.sequence_parallel:
+                    rotary_pos_emb = self.rotary_pos_emb(
+                        self.encoder.input_tensor.size(0) * parallel_state.get_tensor_model_parallel_world_size()
+                    )
+                else:
+                    rotary_pos_emb = self.rotary_pos_emb(self.encoder.input_tensor.size(0))
             else:
-                rotary_pos_emb = self.rotary_pos_emb(encoder_input.size(0))
+                if self.sequence_parallel:
+                    rotary_pos_emb = self.rotary_pos_emb(
+                        encoder_input.size(0) * parallel_state.get_tensor_model_parallel_world_size()
+                    )
+                else:
+                    rotary_pos_emb = self.rotary_pos_emb(encoder_input.size(0))
         else:
             rotary_pos_emb = None
         # encoder.