[JAX] Fix mesh resource requirement when no mesh (#2307)

* Fix mesh resource requirement when no mesh

Signed-off-by: Jeremy Berchtold <[email protected]>

* do not require meshresource if all axes are manual axes

Signed-off-by: Jeremy Berchtold <[email protected]>

* remove abstract_mesh is None check

Signed-off-by: Jeremy Berchtold <[email protected]>

---------

Signed-off-by: Jeremy Berchtold <[email protected]>

Author: jberchtold-nvidia

transformer_engine/jax/sharding.py

@@ -75,6 +75,16 @@ def get_sharding_map_logic_axis_to_mesh_axis():
     """
     Generate a dict to map logical axes to mesh axes.
     """
+    mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh
+    if mesh is None or mesh.empty:
+        # If no mesh is defined, return an empty dict and do not require a MeshResource context to be present
+        return {}
+
+    abstract_mesh = get_abstract_mesh()
+    if sorted(abstract_mesh.manual_axes) == sorted(mesh.axis_names):
+        # If all mesh axes are manual axes, return an empty dict and do not require a MeshResource context to be present
+        return {}
+
     gsr = global_mesh_resource()
 
     is_tpsp_enabled = gsr.tpsp_resource is not None and get_mesh_axis_size(gsr.tpsp_resource) > 1