add sharding debug feature

Author: NuojCheng

src/MaxText/train_utils.py

@@ -219,6 +219,11 @@ def setup_train_loop(config, recorder, devices=None):
       # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage
       sharding.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance)
 
+    # print weights sharding info under debug sharding mode
+    if config.debug_sharding:
+      max_utils.print_non_trivial_mesh_axis(model.mesh)
+      maxtext_utils.print_state_mesh_shardings_params(state, state_mesh_shardings, model.mesh)
+
     if config.use_dpo:
       abstract_state, _, _ = maxtext_utils.get_abstract_state(model, tx, config, init_rng, mesh, is_training=True)
       max_logging.log(