[PyTorch] Set flags in norm modules for Mcore sequence-parallel suppo…

…rt (NVIDIA#1528) Set flag in norm modules for Mcore sequence-parallel support Signed-off-by: Tim Moon <[email protected]>
phu0ngng · Mar 1, 2025 · 4b523d2 · 4b523d2
1 parent d3efaeb
commit 4b523d2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py
@@ -104,6 +104,9 @@ def __init__(
 
         # Flag for sequence parallelism (custom Megatron-LM integration)
         self.sequence_parallel: Optional[bool] = sequence_parallel
+        if sequence_parallel is not None:
+            self.weight.sequence_parallel = sequence_parallel
+            self.bias.sequence_parallel = sequence_parallel
 
     def reset_layer_norm_parameters(self) -> None:
         """Init LN params"""

diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py
@@ -108,6 +108,8 @@ def __init__(
 
         # Flag for sequence parallelism (custom Megatron-LM integration)
         self.sequence_parallel: Optional[bool] = sequence_parallel
+        if sequence_parallel is not None:
+            self.weight.sequence_parallel = sequence_parallel
 
     def reset_rms_norm_parameters(self) -> None:
         """Deprecated"""