[Draft] show how to not use a global mesh.

[nouiz]

The TE/JAX back-end currently has a notion of global mesh.
But JAX doesn't have that.

I think I found a way to remove this notion of global mesh and I hope it would simplify the code.

@denera @mingxu1067 What do you think of that?
The impl() lambda can capture the mesh information instead of having a global mesh.

[mingxu1067]

This might be a good solution to avoid on getting mesh from a global variable.
Usually, all_reduce_sum_along_dp_fsdp and all_reduce_max_along_all_axes_except_PP are invoked in partition of custom_partitioning, and it is passed with the current used mesh, then we could pass that mesh to p* functions.

[ksivaman]

After a discussion with @nouiz, we will re-open a fresh PR.

[nouiz]

New version of that PR: #1112