Add mtc_data_parallelism config for multi-tier checkpointing.

PiperOrigin-RevId: 812337661

Author: abhinavclemson

src/MaxText/configs/base.yml

@@ -317,6 +317,11 @@ enable_multi_tier_checkpointing: False
 # It should be a positive number when enabling multi-tier checkpointing.
 multi_tier_checkpointing_backup_interval_minutes: 0
 
+# Number of identical pipelines in job, should be equal to ICI data parallelism * DCN data parallelism.
+# It should be a positive number when enabling multi-tier checkpointing. If set to 0, it will be set to num of slices.
+mtc_data_parallelism: 0
+
+
 # Whether to enable emergency checkpoint. If True, `local_checkpoint_directory` and a non-zero `local_checkpoint_period` must also be specified.
 # Emergency checkpoint is an experimental Orbax feature that: periodically saves to persistent storage and, with a larger invertal, saves to a local directory.
 # During restore, if a local copy is available in any slice, it will be broadcast to other slices without having to fetch from persistent storage.

src/MaxText/max_utils.py

@@ -187,7 +187,10 @@ def maybe_initialize_jax_distributed_system(raw_keys):
         local_checkpoint_directory=raw_keys["local_checkpoint_directory"],
         backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"],
         run_name=raw_keys["run_name"],
-        jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"],
+        jax_initialization_timeout_seconds=raw_keys[
+            "jax_distributed_initialization_timeout"
+        ],
+        data_parallelism=raw_keys["mtc_data_parallelism"],
     )
     max_logging.log("Jax distributed system initialized for multi-tier checkpointing!")
   elif (raw_keys["enable_checkpointing"] and raw_keys["compile_topology_num_slices"] == -1) or raw_keys[