Simplify JAX distributed initialization in benchmarks.

PiperOrigin-RevId: 895163261

Author: Nikhil Bansal

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/checkpoint_manager_benchmark.yaml

@@ -4,8 +4,8 @@ mesh_configs:
   - mesh_axes: ["data", "model"]
     ici_parallelism: {"data": 2, "model": 2}
   - mesh_axes: ["data", "model"]
-    ici_parallelism: {"data": 1, "model": 1}
-    dcn_parallelism: {"data": 4, "model": 1}
+    ici_parallelism: {"data": 16, "model": 1}
+    dcn_parallelism: {"data": 2, "model": 1}
 
 checkpoint_config:
   spec:
@@ -16,4 +16,4 @@ benchmarks:
     options:
       save_interval_steps: [1]
       max_to_keep: [1]
-      train_steps: [1]
+      train_steps: [100]

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/pytree_checkpoint_benchmark.yaml

@@ -8,7 +8,7 @@ mesh_configs:
     process_is_granule: True
   - mesh_axes: ["data", "stage", "fsdp", "fsdp_transpose", "sequence", "tensor", "expert", "autoregressive"]
     # DCN: Across processes assuming 4 processes.
-    dcn_parallelism: {"tensor": 1, "data": 1}
+    dcn_parallelism: {"tensor": 1, "data": 2}
     ici_parallelism: {"tensor": 16, "data": 1}
 
 # The checkpoint configuration, shared across all generated benchmarks.
@@ -37,4 +37,4 @@ benchmarks:
       #   - use_ocdbt
       #   - use_zarr3
       use_ocdbt: [True]
-      use_zarr3: [True, False]
+      use_zarr3: [False]

checkpoint/orbax/checkpoint/_src/testing/benchmarks/pytree_checkpoint_benchmark.py

@@ -166,26 +166,28 @@ def test_fn(
       checkpointer = ocp.Checkpointer(handler)
     metrics_to_measure = _metrics_to_measure(options)
 
-    with metrics.measure("save", metrics_to_measure):
-      checkpointer.save(save_path, args=ocp.args.PyTreeSave(pytree))
-
-    if options.async_enabled:
-      with metrics.measure("wait_until_finished", metrics_to_measure):
-        assert hasattr(checkpointer, "wait_until_finished")
-        checkpointer.wait_until_finished()
-
-    context.pytree = self._clear_pytree(context.pytree)
-
-    with metrics.measure("restore", metrics_to_measure):
-      restored_pytree = checkpointer.restore(
-          save_path,
-          args=ocp.args.PyTreeRestore(
-              item=pytree,
-              restore_args=ocp.checkpoint_utils.construct_restore_args(pytree),
-          ),
-      )
-
-    self._clear_pytree(restored_pytree)
+    for i in range(300):
+      save_path = context.path / "pytree" / str(i)
+      with metrics.measure("save", metrics_to_measure):
+        checkpointer.save(
+            save_path, args=ocp.args.PyTreeSave(pytree)
+        )
+
+      if options.async_enabled:
+        with metrics.measure("wait_until_finished", metrics_to_measure):
+          assert hasattr(checkpointer, "wait_until_finished")
+          checkpointer.wait_until_finished()
+
+      with metrics.measure("restore", metrics_to_measure):
+        _ = checkpointer.restore(
+            save_path,
+            args=ocp.args.PyTreeRestore(
+                item=pytree,
+                restore_args=ocp.checkpoint_utils.construct_restore_args(
+                    pytree
+                ),
+            ),
+        )
 
     checkpointer.close()
     return benchmarks_core.TestResult(metrics=metrics)

checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py

@@ -30,6 +30,7 @@
 from etils import epath
 import jax
 from orbax.checkpoint._src.testing.benchmarks.core import config_parsing
+import pathwaysutils
 
 
 # Core Flags
@@ -63,29 +64,30 @@ def _init_jax_distributed():
   """Initializes JAX distributed system if not managed by XManager."""
 
   try:
-    jax_platforms = os.environ.get('JAX_PLATFORMS')
-    jax_coordinator_address = os.environ.get('JAX_COORDINATOR_ADDRESS')
-    jax_process_id = os.environ.get('JAX_PROCESS_ID')
-    jax_num_processes = os.environ.get('JAX_NUM_PROCESSES')
-    jax_coordinator_port = os.environ.get('JAX_COORDINATOR_PORT')
-    logging.info('JAX_PLATFORMS: %s', jax_platforms)
-    logging.info(
-        'JAX_COORDINATOR_ADDRESS: %s',
-        jax_coordinator_address,
-    )
-    logging.info('JAX_PROCESS_ID: %s', jax_process_id)
-    logging.info('JAX_NUM_PROCESSES: %s', jax_num_processes)
-    logging.info('JAX_COORDINATOR_PORT: %s', jax_coordinator_port)
-    if jax_num_processes is not None:
-      jax_num_processes = int(jax_num_processes)
-    if jax_process_id is not None:
-      jax_process_id = int(jax_process_id)
-    jax.distributed.initialize(
-        coordinator_address=jax_coordinator_address,
-        num_processes=jax_num_processes,
-        process_id=jax_process_id,
-        initialization_timeout=600,
-    )
+    # jax_platforms = os.environ.get('JAX_PLATFORMS')
+    # jax_coordinator_address = os.environ.get('JAX_COORDINATOR_ADDRESS')
+    # jax_process_id = os.environ.get('JAX_PROCESS_ID')
+    # jax_num_processes = os.environ.get('JAX_NUM_PROCESSES')
+    # jax_coordinator_port = os.environ.get('JAX_COORDINATOR_PORT')
+    # logging.info('JAX_PLATFORMS: %s', jax_platforms)
+    # logging.info(
+    #     'JAX_COORDINATOR_ADDRESS: %s',
+    #     jax_coordinator_address,
+    # )
+    # logging.info('JAX_PROCESS_ID: %s', jax_process_id)
+    # logging.info('JAX_NUM_PROCESSES: %s', jax_num_processes)
+    # logging.info('JAX_COORDINATOR_PORT: %s', jax_coordinator_port)
+    # if jax_num_processes is not None:
+    #   jax_num_processes = int(jax_num_processes)
+    # if jax_process_id is not None:
+    #   jax_process_id = int(jax_process_id)
+    # jax.distributed.initialize(
+    #     coordinator_address=jax_coordinator_address,
+    #     num_processes=jax_num_processes,
+    #     process_id=jax_process_id,
+    #     initialization_timeout=600,
+    # )
+    pathwaysutils.initialize()
     logging.info('JAX distributed system initialized.')
   except Exception as e:  # pylint: disable=broad-exception-caught
     logging.warning(
@@ -96,8 +98,8 @@ def _init_jax_distributed():
         exc_info=True,
     )
 
-  logging.info('Default JAX backend: %s', jax.default_backend())
-  logging.info('Available devices: %s', jax.devices())
+  # logging.info('Default JAX backend: %s', jax.default_backend())
+  # logging.info('Available devices: %s', jax.devices())
   logging.info('JAX process index: %d', jax.process_index())
   logging.info('JAX process count: %d', jax.process_count())
   logging.info('JAX device count: %d', jax.device_count())