Bugfix spec check for DP+FSDP configurations

tests/jax/distributed_test_base.py

@@ -17,6 +17,17 @@
 
 def generate_configs():
     configs = []
+    if is_devices_enough(8):
+        configs.append(
+            pytest.param(
+                8,
+                (2, 2, 2),
+                ("dp", "fsdp", "tpsp"),
+                MeshResource(dp_resource="dp", fsdp_resource="fsdp", tpsp_resource="tpsp"),
+                id="n8_dp2_fsdp2_tp2",
+            )
+        )
+
     if is_devices_enough(4):
         configs.append(
             pytest.param(

tests/jax/test_distributed_dense.py

@@ -48,6 +48,9 @@ def _get_sharding_for_gemm(mesh, mesh_resource, partition_layout="rowwise"):
     dp_axis = mesh_resource.dp_resource
     tp_axis = mesh_resource.tpsp_resource
 
+    if mesh_resource.fsdp_resource is not None:
+        dp_axis = (mesh_resource.dp_resource, mesh_resource.fsdp_resource)
+
     if partition_layout == "colwise":
         x_spec = PartitionSpec(dp_axis, None, None)
         weight_spec = PartitionSpec(None, tp_axis)

transformer_engine/jax/cpp_extensions/gemm.py

@@ -964,8 +964,13 @@ def _parse_operand_output_specs(
             # Non-contracting dims of RHS always needs to be gathered, i.e. for TP + activation_hidden
             # No batch-dim check needed as `rhs_non_cspecs` never contains batch-dim.
             # In `rhs_specs`, the batch dim appears only in Wgrad GEMM under `rhs_cspecs`.
+            # Flatten cspecs since a single element can be a tuple of mesh axes (e.g. ("data", "fsdp")).
+            flattened_lhs_non_cspecs = []
+            for spec in lhs_non_cspecs:
+                flattened_lhs_non_cspecs.extend(spec if isinstance(spec, tuple) else [spec])
             rhs_non_cspecs = tuple(
-                None if spec in lhs_non_cspecs else spec for spec in rhs_non_cspecs
+                None if spec in flattened_lhs_non_cspecs or spec in lhs_non_cspecs else spec
+                for spec in rhs_non_cspecs
             )
 
         else:
@@ -992,8 +997,13 @@ def _parse_operand_output_specs(
             # Non-contracting dims of LHS to be gathered along the SP axis.
             # Minor note: This causes MaxText TP (= Megatron TP + activation_hidden sharding) gathering x for
             # dW1 = x^T * dY1 which is unexpected. This is a known issue and no solution has found yet.
+            # Flatten cspecs since a single element can be a tuple of mesh axes (e.g. ("data", "fsdp")).
+            flattened_rhs_non_cspecs = []
+            for spec in rhs_non_cspecs:
+                flattened_rhs_non_cspecs.extend(spec if isinstance(spec, tuple) else [spec])
             lhs_non_cspecs = tuple(
-                None if spec in rhs_non_cspecs else spec for spec in lhs_non_cspecs
+                None if spec in flattened_rhs_non_cspecs or spec in rhs_non_cspecs else spec
+                for spec in lhs_non_cspecs
             )
 
         out_specs = lhs_non_cspecs + rhs_non_cspecs