add sharding debug feature

Author: NuojCheng

src/MaxText/configs/base.yml

@@ -832,6 +832,7 @@ enable_jax_profiler: False
 jax_profiler_port: 9999
 
 log_config: True # Prints the config (after defaults have been set by pyconfig logic)
+debug_sharding: False # Prints model weights sharding info
 
 # Checkpoint Structured logging
 enable_checkpoint_cloud_logger: False

src/MaxText/configs/types.py

@@ -243,6 +243,7 @@ class RunInfo(BaseModel):
       True,
       description="If True, prints the final configuration after initialization.",
   )
+  debug_sharding: bool = Field(False, description="If True, print model weight sharding details.")
   base_output_directory: PathStr = Field("", description="Base directory for all outputs, typically a GCS path.")
   sharding_strategy: None | Literal["experimental"] = Field(
       None,

src/MaxText/max_utils.py

@@ -989,6 +989,13 @@ def get_batch_seq_len_for_mode(config, model_mode):
   return batch_size, seq_len
 
 
+def print_non_trivial_mesh_axis(mesh):
+  """Print mesh axis if its axis size is larger than one."""
+  for mesh_axis, axis_size in mesh.shape.items():
+    if axis_size > 1:
+      print(f"{mesh_axis}: {axis_size}", flush=True)
+
+
 @contextmanager
 def maybe_get_transformer_engine_context(config):
   """Runs a transformer engine context engine manager for GPUs only."""

src/MaxText/maxtext_utils.py

@@ -1142,3 +1142,14 @@ def schedule(step):
     boundaries.append(warmup_steps + cos_steps + constant_zero_steps)
 
   return optax.join_schedules(pieces, boundaries)
+
+
+def print_state_mesh_shardings_params(state, state_sharding, mesh):
+  """Print state shardings."""
+  leaves_params, _ = jax.tree_util.tree_flatten_with_path(state.params)
+  leaves_sharding, _ = jax.tree_util.tree_flatten_with_path(state_sharding.params)
+  for (path, leaf_val), (_, leaf_sharding) in zip(leaves_params, leaves_sharding):
+    path_str = "/".join(str(p.key) for p in path)
+    shape = jax.typeof(leaf_val)
+    pspec = sharding.remove_size_one_mesh_axis(leaf_sharding.spec, mesh)
+    print(f"{path_str:.<80} {shape} {pspec}", flush=True)

src/MaxText/train_compile.py

@@ -29,7 +29,7 @@
 
 import jax
 from jax.experimental.topologies import get_topology_desc
-from jax.sharding import Mesh
+from jax.sharding import Mesh, AxisType
 from jax.experimental.serialize_executable import serialize
 
 from flax.linen import partitioning as nn_partitioning
@@ -41,7 +41,7 @@
 from MaxText import max_utils
 from MaxText import pyconfig
 from MaxText import sharding
-from MaxText.common_types import MODEL_MODE_TRAIN
+from MaxText.common_types import MODEL_MODE_TRAIN, ShardMode
 from MaxText.layers import models
 from MaxText.layers import quantizations
 from MaxText.utils import gcs_utils
@@ -77,8 +77,11 @@ def get_topology_mesh(config):
         num_slices=config.compile_topology_num_slices,
         wrap=target_hardware.wrap,
     ).devices
+  if config.shard_mode == ShardMode.EXPLICIT:
+    jax.config.update("jax_remove_size_one_mesh_axis_from_type", True)
   topology_device_mesh = maxtext_utils.create_device_mesh(config, topology_devices)
-  topology_mesh = Mesh(topology_device_mesh, config.mesh_axes)
+  mesh_axis_type = AxisType.Explicit if config.shard_mode == ShardMode.EXPLICIT else AxisType.Auto
+  topology_mesh = Mesh(topology_device_mesh, config.mesh_axes, axis_types=(mesh_axis_type,) * len(config.mesh_axes))
   return topology_mesh
 
 
@@ -221,6 +224,11 @@ def main(argv: Sequence[str]) -> None:
       )
   )
 
+  # print weights sharding info under debug sharding mode
+  if config.debug_sharding:
+    max_utils.print_non_trivial_mesh_axis(topology_mesh)
+    maxtext_utils.print_state_mesh_shardings_params(shaped_train_args[0], state_mesh_shardings, topology_mesh)
+
   # Compile
   print("Jitting and compiling train step...", flush=True)
   compiled = jit_and_compile(

src/MaxText/train_utils.py

@@ -19,6 +19,7 @@
 import jax
 from MaxText import checkpointing
 from MaxText import max_logging
+from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import sharding
 from MaxText import optimizers
@@ -213,6 +214,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(