Fix RS out dims

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Author: timmoon10

transformer_engine/pytorch/module/layernorm_linear.py

@@ -323,7 +323,9 @@ def forward(
         # Output buffer for Userbuffers reduce-scatter
         reduce_scatter_out = None
         if ub_overlap_rs_fprop:
-            out_shape = [reduce(multiply_op, inp_shape[:-1]) // tp_world_size, out_features]
+            out_shape = list(inp_shape)
+            out_shape[0] //= tp_world_size
+            out_shape[-1] = out_features
             reduce_scatter_out = torch.empty(
                 out_shape, dtype=activation_dtype, device=ln_out_total.device
             )

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -457,8 +457,8 @@ def forward(
         if ub_overlap_rs:
             ub_obj_fc2out = get_ub("fc2_fprop")
             dim_size = list(act_out.size())
-            dim_size[0] = dim_size[0] // tp_world_size
-            dim_size[1] = fc2_weight.size(0)
+            dim_size[0] //= tp_world_size
+            dim_size[-1] = fc2_weight.size(0)
             reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device)
 
         # ------------------------------------------------------

transformer_engine/pytorch/module/linear.py

@@ -278,7 +278,9 @@ def forward(
         # Output buffer for Userbuffers reduce-scatter
         reduce_scatter_out = None
         if ub_overlap_rs_fprop:
-            out_shape = [reduce(multiply_op, inp.shape[:-1]) // tp_world_size, out_features]
+            out_shape = list(inp.shape)
+            out_shape[0] //= tp_world_size
+            out_shape[-1] = out_features
             reduce_scatter_out = torch.empty(
                 out_shape, dtype=activation_dtype, device=inputmat_total.device
             )