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 under explicit shard mode
 
 # Checkpoint Structured logging
 enable_checkpoint_cloud_logger: False

src/MaxText/configs/types.py

@@ -243,6 +243,9 @@ 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 when using explicit shard mode."
+  )
   base_output_directory: PathStr = Field("", description="Base directory for all outputs, typically a GCS path.")
   sharding_strategy: None | Literal["experimental"] = Field(
       None,
@@ -2052,6 +2055,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         )
       if self.quantization:
         raise ValueError("Quantization is not supported with 'explicit' sharding.")
+    else:
+      if self.debug_sharding:
+        raise ValueError("Debug sharding function only works with explicit shard mode. Please set shard_mode=explicit.")
     if (
         self.per_device_batch_size > 0
         and (self.per_device_batch_size * self.max_target_length) % self.num_vocab_tiling != 0

src/MaxText/max_utils.py

@@ -989,6 +989,21 @@ 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)
+
+
+def print_params_sharding_shapes(tree):
+  """Print flatten tree and their leaf shapes."""
+  leaves, _ = jax.tree_util.tree_flatten_with_path(tree)
+  for path, leaf in leaves:
+    path_str = "/".join(str(p.key) for p in path)
+    print(f"{path_str:.<80} {jax.typeof(leaf)}", flush=True)
+
+
 @contextmanager
 def maybe_get_transformer_engine_context(config):
   """Runs a transformer engine context engine manager for GPUs only."""

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)
+    max_utils.print_params_sharding_shapes(state.params)
+
   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)
+    max_utils.print_params_sharding_shapes(shaped_train_args[0].params)
+
   # Serialize and save the compiled object
   if config.compiled_trainstep_file != "":
     print("Saving compiled object...")