Add escape hatch for jax_init_info

PiperOrigin-RevId: 874188030

Author: Orbax Authors

checkpoint/CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - `uvloop` dependency for improved event loop performance
+- Add option to disable for multi-tier checkpointing process id initialization
 
 ### Removed
 

checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/initialization.py

@@ -70,6 +70,34 @@ def _create_replicator_file(
   os.rename(temp_file, replicator_file)
 
 
+def _initialize_jax_from_mtc(
+    local_checkpoint_directory: epath.Path,
+    jax_initialization_timeout_seconds: int = 900,
+) -> str:
+  """Initialize jax with jax_init_info."""
+  local_checkpoint_directory = epath.Path(local_checkpoint_directory)
+  process_id, coordinator_address = _retrieve_jax_init_info(
+      local_checkpoint_directory
+  )
+  if not process_id or not coordinator_address:
+    raise ValueError(
+        'Data is missing from the JAX init info file: Current values:'
+        f' process_id: {process_id}, coordinator_address: {coordinator_address}'
+    )
+  logging.info(
+      'Using process_id %s and coordinator_address %s to initialize JAX'
+      ' distributed runtime...',
+      process_id,
+      coordinator_address,
+  )
+  jax.distributed.initialize(
+      process_id=int(process_id),
+      coordinator_address=coordinator_address,
+      initialization_timeout=jax_initialization_timeout_seconds,
+  )
+  return process_id
+
+
 def initialize_multi_tier_checkpointing(
     local_checkpoint_directory: epath.Path,
     *,
@@ -78,6 +106,7 @@ def initialize_multi_tier_checkpointing(
     run_name: Optional[str] = None,
     data_parallelism: Optional[int] = None,
     jax_initialization_timeout_seconds: int = 900,
+    use_mtc_process_ids: bool = True,
 ):
   """Initializes multi-tier checkpointing.
 
@@ -91,27 +120,18 @@ def initialize_multi_tier_checkpointing(
       equal to ICI data parallelism * DCN data parallelism. If not provided, it
       will be inferred from the number of slices.
     jax_initialization_timeout_seconds: The timeout for JAX initialization.
+    use_mtc_process_ids: Use the MTC rank server to calculate process ids.
   """
-  local_checkpoint_directory = epath.Path(local_checkpoint_directory)
-  process_id, coordinator_address = _retrieve_jax_init_info(
-      local_checkpoint_directory
-  )
-  if not process_id or not coordinator_address:
-    raise ValueError(
-        'Data is missing from the JAX init info file: Current values:'
-        f' process_id: {process_id}, coordinator_address: {coordinator_address}'
+  if use_mtc_process_ids:
+    process_id = _initialize_jax_from_mtc(
+        local_checkpoint_directory, jax_initialization_timeout_seconds
     )
-  logging.info(
-      'Using process_id %s and coordinator_address %s to initialize JAX'
-      ' distributed runtime...',
-      process_id,
-      coordinator_address,
-  )
-  jax.distributed.initialize(
-      process_id=int(process_id),
-      coordinator_address=coordinator_address,
-      initialization_timeout=jax_initialization_timeout_seconds,
-  )
+  else:
+    process_id = None
+    jax.distributed.initialize(
+        initialization_timeout=jax_initialization_timeout_seconds,
+    )
+
   multihost.initialize_runtime_to_distributed_ids()
   multihost.initialize_distributed_to_device_ids()
   _wait_for_replicator_file_to_disappear(local_checkpoint_directory)
@@ -127,14 +147,24 @@ def initialize_multi_tier_checkpointing(
   process_index_to_node_rank = (
       multihost.runtime_to_distributed_ids()
   )
-  logging.info(
-      'Mapping of IDs: jax-init-info.txt=%s, NodeRank=%s, ProcessIndex=%s,'
-      ' ProcessIndex->NodeRank=%s',
-      process_id,
-      node_rank,
-      my_process_index,
-      process_index_to_node_rank,
-  )
+  if use_mtc_process_ids:
+    logging.info(
+        'Mapping of IDs: jax-init-info.txt=%s, NodeRank=%s, ProcessIndex=%s,'
+        ' ProcessIndex->NodeRank=%s',
+        process_id,
+        node_rank,
+        my_process_index,
+        process_index_to_node_rank,
+    )
+  else:
+    logging.info(
+        'Mapping of IDs (jax-init-info not used): NodeRank=%s, ProcessIndex=%s,'
+        ' ProcessIndex->NodeRank=%s',
+        node_rank,
+        my_process_index,
+        process_index_to_node_rank,
+    )
+
   my_in_pipeline_index = my_process_index % nodes_per_slice
   peer_ranks = []
   for i in range(num_slices):

checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/initialization_test.py

@@ -288,6 +288,61 @@ def test_initialize_multi_tier_checkpointing_run_name_not_set(
       mock_initialize_distributed_to_device_ids.assert_called_once()
       self.assertEqual(mock_wait_for_replicator_file_to_disappear.call_count, 1)
 
+  @mock.patch.object(
+      initialization, "_wait_for_replicator_file_to_disappear", autospec=True
+  )
+  @mock.patch.object(initialization, "_create_replicator_file", autospec=True)
+  @mock.patch.object(jax.distributed, "initialize", autospec=True)
+  @mock.patch.object(
+      multihost, "initialize_runtime_to_distributed_ids", autospec=True
+  )
+  @mock.patch.object(
+      multihost, "initialize_distributed_to_device_ids", autospec=True
+  )
+  @mock.patch.object(multihost, "runtime_to_distributed_ids", autospec=True)
+  def test_initialize_multi_tier_checkpointing_skip_init_info(
+      self,
+      mock_runtime_to_distributed_ids,
+      mock_initialize_distributed_to_device_ids,
+      mock_initialize_runtime_to_distributed_ids,
+      mock_jax_distributed_initialize,
+      mock_create_replicator_file,
+      mock_wait_for_replicator_file_to_disappear,
+  ):
+    mock_runtime_to_distributed_ids.return_value = [0, 1]
+    mock_jax_distributed_initialize.return_value = None
+    mock_initialize_runtime_to_distributed_ids.return_value = [None, None]
+    mock_initialize_distributed_to_device_ids.return_value = None
+    mock_create_replicator_file.return_value = [None, None]
+    mock_wait_for_replicator_file_to_disappear.return_value = False
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      epath.Path(tmp_dir).mkdir(parents=True, exist_ok=True)
+      replicator_file = epath.Path(tmp_dir) / initialization._REPLICATOR_FILE
+      replicator_file.write_text("replicator.yaml")
+      self.assertTrue(replicator_file.exists())
+
+      restore_dir = epath.Path(tmp_dir) / "test-run-s1-n0-w0.restore"
+      restore_dir.write_text("restore_dir")
+      self.assertTrue(restore_dir.exists())
+
+      initialization.initialize_multi_tier_checkpointing(
+          epath.Path(tmp_dir),
+          num_slices=1,
+          run_name="test-run",
+          data_parallelism=1,
+          use_mtc_process_ids=False,
+      )
+      mock_jax_distributed_initialize.assert_called_once_with(
+          initialization_timeout=900,
+      )
+      mock_initialize_runtime_to_distributed_ids.assert_called_once()
+      mock_initialize_distributed_to_device_ids.assert_called_once()
+      self.assertEqual(mock_wait_for_replicator_file_to_disappear.call_count, 2)
+      mock_create_replicator_file.assert_called_once()
+      expected_restore_dir = epath.Path(tmp_dir) / "1"
+      self.assertTrue(expected_restore_dir.exists())
+
 
 if __name__ == "__main__":
   absltest.main()