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.py

@@ -240,7 +240,6 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
     _loss_fn = dpo_loss_fn
 
   params = state.params
-
   if config.gradient_accumulation_steps > 1:
     loss, aux, raw_grads = gradient_accumulation_loss_and_grad(
         _loss_fn,
@@ -411,6 +410,11 @@ def train_loop(config, recorder, state=None):
       compiled_stats = compiled.memory_analysis()
       max_utils.print_compiled_memory_stats(compiled_stats)
 
+  # print weights sharding info under debug sharding mode
+  if config.debug_sharding:
+    max_utils.print_non_trivial_mesh_axis(mesh)
+    maxtext_utils.print_state_mesh_shardings_params(state, state_mesh_shardings, mesh)
+
   start_step = get_first_step(state)  # this is the start_step for training
   prof = profiler.Profiler(config, offset_step=start_step)
   metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule)

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
 
 
@@ -236,6 +239,11 @@ def main(argv: Sequence[str]) -> None:
   )
   print("Jitting and compilation complete!", flush=True)
 
+  # 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)
+
   # Serialize and save the compiled object
   if config.compiled_trainstep_file != "":
     print("Saving compiled object...")