Merge pull request #1533 from AI-Hypercomputer:lance-tp-fix2

PiperOrigin-RevId: 750766633

Author: maxtext authors

MaxText/configs/base.yml

@@ -63,11 +63,12 @@ checkpoint_storage_use_zarr3: True
 checkpoint_storage_concurrent_gb: 96
 ############################### END CHECKPOINTING ##################################
 
-
+############################### BEGIN TESTING ##################################
 reuse_example_batch: 0 # for testing TPU performance, this options repeated uses the same batch.
-
-
 metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written.
+disalbe_key_validation: False # for testing, if true, skip the key validation.
+############################### END TESTING ##################################
+
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 gcs_metrics: False
 
@@ -330,7 +331,7 @@ logical_axis_rules: [
                     ]
 # Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
 data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']]
-
+input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
 # sharding tolerance: float between 0.0 and 1.0 representing the allowed percentage of non-sharded parameters.
 sharding_tolerance: 0.02
 

MaxText/input_pipeline/_grain_data_processing.py

@@ -196,7 +196,7 @@ def make_grain_train_iterator(
           tokenize=config.tokenize_train_data,
           grain_worker_count=config.grain_worker_count,
       )
-    return multihost_dataloading.MultiHostDataLoadIterator(train_dataloader, global_mesh)
+    return multihost_dataloading.MultiHostDataLoadIterator(train_dataloader, global_mesh, config)
   else:
     get_ds_fn = functools.partial(
         get_datasets,
@@ -262,7 +262,7 @@ def make_grain_eval_iterator(
           tokenize=config.tokenize_eval_data,
           grain_worker_count=config.grain_worker_count_eval,
       )
-    return multihost_dataloading.MultiHostDataLoadIterator(eval_dataloader, global_mesh)
+    return multihost_dataloading.MultiHostDataLoadIterator(eval_dataloader, global_mesh, config)
   else:
     get_ds_fn = functools.partial(
         get_datasets,

MaxText/input_pipeline/_hf_data_processing.py

@@ -185,10 +185,7 @@ def preprocessing_pipeline(
       read_options=grain.ReadOptions(num_threads=num_threads, prefetch_buffer_size=128),
   )
 
-  multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(dataloader, global_mesh)
-
-  # Return multi-host jax.Array prep iterator
-  return multihost_gen
+  return dataloader
 
 
 def make_hf_train_iterator(
@@ -205,7 +202,7 @@ def make_hf_train_iterator(
       streaming=True,
       token=config.hf_access_token,
   )
-  train_iter = preprocessing_pipeline(
+  train_data_loader = preprocessing_pipeline(
       dataloading_host_index=process_indices_train.index(jax.process_index()),
       dataloading_host_count=len(process_indices_train),
       global_mesh=global_mesh,
@@ -226,6 +223,7 @@ def make_hf_train_iterator(
       use_sft=config.use_sft,
       sft_train_on_completion_only=config.sft_train_on_completion_only,
   )
+  train_iter = multihost_dataloading.MultiHostDataLoadIterator(train_data_loader, global_mesh, config)
   return train_iter
 
 
@@ -247,7 +245,7 @@ def make_hf_eval_iterator(
     eval_generate_padding_example = True
   else:
     eval_generate_padding_example = False
-  eval_iter = preprocessing_pipeline(
+  eval_data_loader = preprocessing_pipeline(
       dataloading_host_index=process_indices_eval.index(jax.process_index()),
       dataloading_host_count=len(process_indices_eval),
       global_mesh=global_mesh,
@@ -268,4 +266,5 @@ def make_hf_eval_iterator(
       use_sft=config.use_sft,
       sft_train_on_completion_only=config.sft_train_on_completion_only,
   )
+  eval_iter = multihost_dataloading.MultiHostDataLoadIterator(eval_data_loader, global_mesh, config)
   return eval_iter

MaxText/input_pipeline/_tfds_data_processing.py

@@ -197,7 +197,7 @@ def make_tfds_train_iterator(
         use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
-    return multihost_dataloading.MultiHostDataLoadIterator(train_dataloader, global_mesh)
+    return multihost_dataloading.MultiHostDataLoadIterator(train_dataloader, global_mesh, config)
   else:
     get_ds_fn = functools.partial(
         get_datasets,
@@ -261,7 +261,7 @@ def make_tfds_eval_iterator(
         use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
-    return multihost_dataloading.MultiHostDataLoadIterator(eval_dataloader, global_mesh)
+    return multihost_dataloading.MultiHostDataLoadIterator(eval_dataloader, global_mesh, config)
   else:
     get_ds_fn = functools.partial(
         get_datasets,

MaxText/input_pipeline/_tfds_data_processing_c4_mlperf.py

@@ -330,7 +330,7 @@ def make_c4_mlperf_train_iterator(
       shuffle_buffer_size=128,
       data_shuffle_seed=config.data_shuffle_seed,
   )
-  train_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(train_ds, global_mesh)
+  train_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(train_ds, global_mesh, config)
   return train_multihost_gen
 
 
@@ -360,7 +360,7 @@ def make_c4_mlperf_eval_iterator(
       max_target_length=config.max_target_length,
   )
 
-  eval_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(eval_ds, global_mesh)
+  eval_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(eval_ds, global_mesh, config)
 
   # Return multi-host jax.Array prep iterator
   return eval_multihost_gen

MaxText/input_pipeline/input_pipeline_interface.py

@@ -26,6 +26,7 @@
 from MaxText.input_pipeline._grain_data_processing import make_grain_train_iterator, make_grain_eval_iterator
 from MaxText.input_pipeline._tfds_data_processing_c4_mlperf import make_c4_mlperf_train_iterator, make_c4_mlperf_eval_iterator
 from MaxText.input_pipeline._hf_data_processing import make_hf_train_iterator, make_hf_eval_iterator
+from MaxText import maxtext_utils
 from MaxText import multihost_dataloading
 
 
@@ -35,8 +36,9 @@ class SyntheticDataIterator:
   def __init__(self, config, mesh):
     self.mesh = mesh
     self.config = config
-    data_pspec = P(*config.data_sharding)
-    data_pspec_shardings = jax.tree_util.tree_map(lambda p: jax.sharding.NamedSharding(mesh, p), data_pspec)
+    data_pspec_shardings = maxtext_utils.get_input_data_sharding(
+        mesh, config.input_data_sharding_logical_axes, config.logical_axis_rules
+    )
     self.data_generator = jax.jit(
         SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings, static_argnums=0
     )
@@ -82,7 +84,7 @@ class BadSyntheticDataIterator:
   def __init__(self, config, mesh):
     self.mesh = mesh
     dataset = BadSyntheticDataIterator.get_bad_synthetic_data(config)
-    self.data_generator = multihost_dataloading.MultiHostDataLoadIterator(dataset, self.mesh)
+    self.data_generator = multihost_dataloading.MultiHostDataLoadIterator(dataset, self.mesh, config)
 
   def __iter__(self):
     return self.data_generator
@@ -118,8 +120,7 @@ def get_process_loading_real_data(
     data_sharding, global_batch_size_to_load, global_batch_size_to_train_on, max_target_length, mesh
 ):
   """Get list of processes loading data from GCS when expansion_factor_real_data != -1"""
-  sharding = jax.sharding.NamedSharding(mesh, P(*data_sharding))
-  devices_indices_map = sharding.devices_indices_map((global_batch_size_to_load, max_target_length))
+  devices_indices_map = data_sharding.devices_indices_map((global_batch_size_to_load, max_target_length))
   batch_cutoff = global_batch_size_to_train_on
   process_loading_real_data = set()
   for p, indices in devices_indices_map.items():
@@ -149,16 +150,19 @@ def create_data_iterator(config, mesh):
   if config.dataset_type == "synthetic":
     return SyntheticDataIterator(config, mesh), None
 
+  input_data_sharding = maxtext_utils.get_input_data_sharding(
+      mesh, config.input_data_sharding_logical_axes, config.logical_axis_rules
+  )
   process_indices_train = get_process_loading_real_data(
-      config.data_sharding,
+      input_data_sharding,
       config.global_batch_size_to_load,
       config.global_batch_size_to_train_on,
       config.max_target_length,
       mesh,
   )
   if config.eval_interval > 0:
     process_indices_eval = get_process_loading_real_data(
-        config.data_sharding,
+        input_data_sharding,
         config.global_batch_size_to_load_eval,
         config.global_batch_size_to_eval_on,
         config.max_target_length,

MaxText/max_utils.py

@@ -373,21 +373,27 @@ def get_coordinator_ip_address():
   return coordinator_ip_address
 
 
-def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_type):
+def get_unspecified_mesh_axes_value(parallelism_vals, target_product, parallelism_type):
   """Evaluates unspecified DCN/ICI parallelism values"""
-  if -1 in parallelism_vals:
-    assert (
-        parallelism_vals.count(-1) == 1
-    ), f"Found unspecified values (-1) for more than one {parallelism_type}\
-      parallelism axis. At most one axis can be unspecified."
+  assert (
+      parallelism_vals.count(-1) == 1
+  ), f"Found unspecified values (-1) for more than one {parallelism_type}\
+    parallelism axis. At most one axis can be unspecified."
+
+  determined_val = target_product / np.prod(parallelism_vals) * -1
+
+  assert (
+      determined_val >= 1 and determined_val.is_integer
+  ), f"Unspecified value unable to be determined with the given\
+    {parallelism_type} parallelism values"
 
-    determined_val = target_product / np.prod(parallelism_vals) * -1
+  return int(determined_val)
 
-    assert (
-        determined_val >= 1 and determined_val.is_integer
-    ), f"Unspecified value unable to be determined with the given\
-      {parallelism_type} parallelism values"
 
+def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_type):
+  """Evaluates unspecified DCN/ICI parallelism values"""
+  if -1 in parallelism_vals:
+    determined_val = get_unspecified_mesh_axes_value(parallelism_vals, target_product, parallelism_type)
     parallelism_vals[parallelism_vals.index(-1)] = int(determined_val)
 
   target_type = "slices" if parallelism_type == "DCN" else "devices per slice"
@@ -780,6 +786,35 @@ def reorder_causal_load_balanced(batch, cp_size):
   }
 
 
+def shard_reorder_causal_load_balanced(batch, cp_size):
+  """Shard the output of the reordered sequence."""
+  reordered = reorder_causal_load_balanced(batch, cp_size)
+  for _, v in batch.items():
+    if isinstance(v, jax.Array):
+      reordered = jax.lax.with_sharding_constraint(reordered, v.sharding)
+      break
+  return reordered
+
+
 def get_reorder_callable(cp_size):
   """Creates a callable that can be used with map() to reorder batches."""
-  return functools.partial(reorder_causal_load_balanced, cp_size=cp_size)
+  return functools.partial(shard_reorder_causal_load_balanced, cp_size=cp_size)
+
+
+def compute_axis_product(axis_spec, mesh_dict):
+  """Computes the product of the axis specified in axis_spec."""
+  if isinstance(axis_spec, str):
+    axis_spec = (axis_spec,)
+  elif axis_spec is None:
+    return 1
+  product = 1
+  for dim_name in axis_spec:
+    if dim_name in mesh_dict:
+      product *= mesh_dict[dim_name]
+  return product
+
+
+def construct_parallelism_name(mesh_axis: str, prefix: str) -> str:
+  if mesh_axis == "stage":
+    return f"{prefix}_pipeline_parallelism"
+  return f"{prefix}_{mesh_axis}_parallelism"

MaxText/maxtext_utils.py

@@ -23,6 +23,7 @@
 from MaxText import max_utils
 from jax.sharding import PartitionSpec as P
 from jax.experimental.serialize_executable import deserialize_and_load
+from flax import linen as nn
 
 import pickle
 import functools
@@ -32,6 +33,7 @@
 from flax.linen import partitioning as nn_partitioning
 
 from MaxText import max_logging
+import ml_collections
 import numpy as np
 import jax.numpy as jnp
 from MaxText import checkpointing
@@ -50,12 +52,16 @@
 NUM_IMAGE_CHANNELS = 3
 
 
+def get_input_data_sharding(mesh, input_data_sharding_logical_axes, logical_axis_rules):
+  data_pspec = P(*input_data_sharding_logical_axes)
+  return nn.logical_to_mesh_sharding(data_pspec, mesh, logical_axis_rules)
+
+
 def get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config):
   """Get the shardings (both state and data) for train_step"""
   functional_train = get_functional_train_step(train_step, model, config, state_mesh_shardings)
   functional_train.__name__ = "train_step"
-  data_pspec = P(*config.data_sharding)
-  data_sharding = jax.tree_util.tree_map(lambda p: jax.sharding.NamedSharding(mesh, p), data_pspec)
+  data_sharding = get_input_data_sharding(mesh, config.input_data_sharding_logical_axes, config.logical_axis_rules)
   in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
   out_shardings = (state_mesh_shardings, None)  # State, metrics
   static_argnums = ()  # We partial out the static argnums of model and config
@@ -71,8 +77,7 @@ def get_functional_eval_with_signature(eval_step, mesh, state_mesh_shardings, mo
   """Get the shardings (both state and data) for eval_step"""
   functional_eval = get_functional_eval_step(eval_step, model, config)
   functional_eval.__name__ = "eval_step"
-  data_pspec = P(*config.data_sharding)
-  data_sharding = jax.tree_util.tree_map(lambda p: jax.sharding.NamedSharding(mesh, p), data_pspec)
+  data_sharding = get_input_data_sharding(mesh, config.input_data_sharding_logical_axes, config.logical_axis_rules)
   in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
   out_shardings = None  # metrics
   static_argnums = ()  # We partial out the static argnums of model, config
@@ -690,14 +695,36 @@ def add_config_to_summary_writer(config, summary_writer):
       max_utils.add_text_to_summary_writer(key, str(value), summary_writer)
 
 
-def create_device_mesh(config, devices=None):
-  """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
+def get_ici_parallelism(config, devices=None):
+  """Get the ICI parallelism for the model."""
   if devices is None:
     devices = jax.devices()
   num_devices = len(devices)
   num_slices = 1 if config.inference_benchmark_test else config.num_slices
   num_devices_per_slice = num_devices // num_slices
 
+  return max_utils.fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
+
+
+def get_dcn_parallelism(config):
+  """Get the DCN parallelism for the model."""
+  num_slices = 1 if config.inference_benchmark_test else config.num_slices
+  return max_utils.fill_unspecified_mesh_axes(config.dcn_parallelism.copy(), num_slices, "DCN")
+
+
+def get_slices_and_devices(config, devices=None):
+  if devices is None:
+    devices = jax.devices()
+  num_devices = len(devices)
+  num_slices = 1 if config.inference_benchmark_test else config.num_slices
+  num_devices_per_slice = num_devices // num_slices
+  return num_slices, num_devices_per_slice
+
+
+def create_device_mesh(config, devices=None):
+  """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
+  num_slices, num_devices_per_slice = get_slices_and_devices(config, devices)
+  num_devices = num_devices_per_slice * num_slices
   multi_slice_env = num_slices > 1
 
   # Find possible unspecified parallelisms