Add support loading ref checkpoint with colocated Python and PathwaysColocatedPythonGuide.md

PiperOrigin-RevId: 912169322

Author: ChromeHearts

checkpoint/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Tensorstore non-atomic file I/O locking on local/POSIX-compatible filesystems
 to avoid unnecessary renames.
+- Add loading source checkpoint with Pathways colocated-Python in Benchmark
+suite.
 
 ## [0.11.38] - 2026-05-05
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/configs/llama-70b-v5p-64-pw-example.yaml

@@ -0,0 +1,23 @@
+# The name for the entire test suite run.
+suite_name: "Llama 3.1 70B v5p-64"
+num_repeats: 5
+
+mesh_configs:
+  - mesh_axes: ["data", "fsdp", "tensor"]
+    ici_parallelism: {"data": 1, "fsdp": 32, "tensor": 1}
+
+# The checkpoint configuration, shared across all generated benchmarks.
+checkpoint_configs:
+  - path: "gs://orbax-benchmarks/checkpoints/llama-3.1-70B-checkpoints/0/items"
+    sharding_config_path: "gs://orbax-benchmarks/sharding-configs/llama3.1-70b-v5p-64-data-1-fsdp-32-tensor-1/abstract_state.json"
+    load_with_colocated_python: true
+
+benchmarks:
+  - generator: "orbax.checkpoint._src.testing.benchmarks.pytree_checkpoint_benchmark.PyTreeCheckpointBenchmark"
+    options:
+      async_enabled: true
+      use_ocdbt: true
+      use_zarr3: true
+      use_replica_parallel: false
+      use_compression: true
+      use_colocated_python: true

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/checkpoint_generation.py

@@ -23,10 +23,12 @@
 import jax.numpy as jnp
 import numpy as np
 from orbax.checkpoint import checkpoint_utils
+from orbax.checkpoint import pathways
 from orbax.checkpoint._src.arrays import abstract_arrays
 from orbax.checkpoint._src.arrays import sharding as sharding_utils
 from orbax.checkpoint._src.checkpointers import checkpointer
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
+from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint._src.testing.benchmarks.core import configs
 from orbax.checkpoint._src.tree import utils as tree_utils
@@ -214,6 +216,16 @@ def load_checkpoint(config: configs.CheckpointConfig) -> Any:
   abstract_state = _get_abstract_state(config, use_ocdbt=use_ocdbt)
   restore_args = checkpoint_utils.construct_restore_args(abstract_state)
 
+  if multihost.is_pathways_backend():
+    checkpointing_impl = pathways.CheckpointingImpl.from_options(
+        use_colocated_python=config.load_with_colocated_python,
+    )
+    pathways.register_type_handlers(
+        checkpointing_impl=checkpointing_impl,
+        use_replica_parallel=False,
+        enable_replica_parallel_separate_folder=False,
+    )
+
   with checkpointer.Checkpointer(
       pytree_checkpoint_handler.PyTreeCheckpointHandler(use_ocdbt=use_ocdbt)
   ) as ckptr:

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/configs.py

@@ -81,12 +81,15 @@ class CheckpointConfig:
         1024], 'sharding': ['data', 'model']  # PartitionSpec }, 'step': 'int' }
       sharding_config_path: A path to a file containing sharding specifications,
         used alongside `path`. See above.
+      load_with_colocated_python: If True, the checkpoint will be loaded with
+        colocated Python.  Only effective in Pathways and `path` is specified.
   """
 
   path: str | None = None
   random_seed: int = 0
   spec: dict[str, Any] | None = None
   sharding_config_path: str | None = None
+  load_with_colocated_python: bool = False
 
   def __post_init__(self):
     if self.path is not None and self.spec is not None:

checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/PathwaysColocatedPythonGuide.md

@@ -0,0 +1,85 @@
+# Orbax Benchmark on Pathways Colocated Python Guide
+
+## Introduction
+
+This guide provides step-by-step instructions for running Orbax benchmarks
+on **Pathways with Colocated Python**. Before you start, make sure you are
+familiarize with [README.md](README.md) on how to run `xpk`. This setup allows
+running benchmarks where the Python code executes directly on the TPU worker
+nodes (colocated with the Pathways server), which is useful for reducing data
+transfer between Pathways Client and Workers.
+
+## Step-by-step Instructions
+
+### Create a Pathways cluster
+
+Use `launch_xpk.py --create_cluster` option or command below
+
+```shell
+xpk cluster create-pathways \
+  --cluster orbax-benchmark-v5p-64-pw \
+  --tpu-type=v5p-64 \
+  --num-slices=1 \
+  --spot \
+  --zone=us-east5-a
+```
+
+### Install latest XPK
+
+```shell
+pip install --upgrade xpk
+```
+
+### Build both benchmark & sidecar images.
+
+- Use `--jax-version 0.10.0` for Benchmark image build. This has to match
+  with `Dockerfile.sidecar`'s source image.
+- Use `--base-image python:3.12-slim` for the Benchmark image build, the python
+  version also has with `Dockerfile.sidecar`'s source image.
+- Use `--build-sidecar true` to build sidecar image as well.
+
+This is a full example of `build_image.sh` command.
+```shell
+export DATE_STR=(`date +"%Y%m%d-%H%M%Z"`)
+export IMG_TAG=orbax-benchmark-local-tpu-$DATE_STR
+
+bash build_image.sh  \
+  --project orbax-checkpoint \
+  --local-repo ./ \
+  --tag $IMG_TAG \
+  --device tpu \
+  --jax-version 0.10.0 \
+  --base-image python:3.12-slim \
+  --no-cache \
+  --build-benchmark true \
+  --build-sidecar true
+```
+
+### Run Benchmark
+
+Following is the full example to run a Orbax Benchmark on Pathways Colocated
+Python. Make sure to set the sidecar docker_image in `--pathways_sidecar_image`
+to the images you have just built above.
+
+```shell
+MODEL=llama-70b-v5p-64-pw
+OUTPUT_DIR=gs://orbax-benchmarks/${USER}/${DATE_STR}/$MODEL # don't end with slash
+CONFIG_DIR=$OUPTUT_DIR
+
+python3 launch_xpk.py \
+  --verbose \
+  --enable_pathways \
+  --cluster_name dnlng-v5p-64-pw \
+  --create_cluster=False \
+  --delete_cluster_on_completion=False \
+  --spot \
+  --tpu_type v5p-64 \
+  --num_slices 1 \
+  --zone us-east5-a \
+  --config_file=orbax/checkpoint/_src/testing/benchmarks/configs/llama-70b-v5p-64-pw-example.yaml \
+  --docker_image=gcr.io/orbax-checkpoint/orbax-benchmarks:$IMG_TAG \
+  --pathways_sidecar_image=gcr.io/orbax-checkpoint/orbax-benchmarks/sidecar:$IMG_TAG \
+  --output_directory=$OUTPUT_DIR \
+  --config_directory=$CONFIG_DIR
+
+```