[PyTorch] set SP attr on bias param for reduction (#440)

Fix for sequence-parallel
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/module/layernorm_linear.py

@@ -737,6 +737,8 @@ class LayerNormLinear(TransformerEngineBaseModule):
                 self.register_parameter(
                     bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
                 )
+                if parallel_mode == "row":
+                    setattr(getattr(self, bname), "sequence_parallel", sequence_parallel)
             else:
                 setattr(self, bname, torch.Tensor().to(dtype=params_dtype, device=device))
 

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -1054,6 +1054,9 @@ class LayerNormMLP(TransformerEngineBaseModule):
             self.fc2_bias = Parameter(
                 torch.empty(hidden_size, device=device, dtype=params_dtype)
             )
+            # RPL
+            if self.set_parallel_mode:
+                setattr(self.fc2_bias, "sequence_parallel", sequence_parallel)
         else:
             self.fc2_bias = torch.Tensor().to(dtype=params_dtype, device=device)
 

transformer_engine/pytorch/module/linear.py

@@ -628,6 +628,8 @@ class Linear(TransformerEngineBaseModule):
                 self.register_parameter(
                     bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
                 )
+                if parallel_mode == "row":
+                    setattr(getattr(self, bname), "sequence_parallel", sequence_parallel)
             else:
                 setattr(self, bname, torch.Tensor().to(dtype=params_dtype, device=device))