From 1469e211f3e8804690909428cde00d3062431250 Mon Sep 17 00:00:00 2001 From: "Chang (AIML) Liu" Date: Thu, 19 Feb 2026 14:35:27 -0800 Subject: [PATCH 1/5] [2/n] Add colocated python path to checkpoint manager. (#2180) * [2/n] Add colocated python path to checkpoint manager. * Fix threading and asyncio bugs. * resolve dex-ai's comments. * Minor, remove redundant logging. --- .../examples/colocated_python_benchmark.py | 155 ++---------- axlearn/common/array_serialization.py | 233 ++++++++++++++++-- 2 files changed, 234 insertions(+), 154 deletions(-) diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index d548bc25e..638dfa0a6 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -19,13 +19,8 @@ """ import argparse -import asyncio -import functools -import logging import os -import sys import time -from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from datetime import datetime from typing import Any, Dict, Optional, Sequence @@ -34,12 +29,11 @@ import jax.numpy as jnp import pathwaysutils # pytype: disable=import-error from jax._src.mesh import thread_resources -from jax.experimental import colocated_python, mesh_utils +from jax.experimental import mesh_utils from jax.experimental.array_serialization import serialization as array_serialization -from jax.experimental.array_serialization import tensorstore_impl from axlearn.common import utils -from axlearn.common.array_serialization import GlobalAsyncCheckpointManager, _async_deserialize +from axlearn.common.array_serialization import GlobalAsyncCheckpointManager from axlearn.common.checkpointer import parse_step_from_dir, read_index_file from axlearn.common.utils import TensorSpec, infer_mesh_shape @@ -62,78 +56,6 @@ def maybe_profile(enabled: bool, profile_dir: Optional[str]): return nullcontext() -def _colocated_deserialize( - shardings: Sequence[jax.sharding.NamedSharding], - tensorstore_specs: Sequence[Dict[str, Any]], - global_shapes: Sequence[tuple], - dtypes: Sequence[jnp.dtype], -): - concurrent_bytes = 34359738368 * 6 # 32GB * 6 - cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) - print(f"{cpu_devices=}") - - if len(cpu_devices) > 1: - print(f"TPU Mesh: {thread_resources.env.physical_mesh}") - cpu_mesh = colocated_python.colocated_cpu_devices(thread_resources.env.physical_mesh) - print(f"CPU Mesh: {cpu_mesh}") - cpu_shardings = [ - jax.sharding.NamedSharding(cpu_mesh, sharding.spec) for sharding in shardings - ] - else: - cpu_shardings = [ - jax.sharding.SingleDeviceSharding(cpu_devices[0]) for sharding in shardings - ] - - def output_spec_fn(): - return [ - jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=sharding) - for shape, dtype, sharding in zip(global_shapes, dtypes, cpu_shardings) - ] - - @colocated_python.colocated_python - def run_deserializer(): - # Object should be created once per process. - # pylint: disable=protected-access - # print("Print statement inside colocated") - logging.info("Logging statement inside colocated") - sys.stderr.write("Stdder statement in colocated") - start_colocated_time = time.perf_counter() - byte_limiter = tensorstore_impl._LimitInFlightBytes(concurrent_bytes) - h2d_limiter = tensorstore_impl._LimitInFlightBytes(concurrent_bytes) - thread_pool = ThreadPoolExecutor(1) - multi_thread_pool = ThreadPoolExecutor(2) - - future_arrays = jax.tree.map( - functools.partial( - _async_deserialize, - byte_limiter=byte_limiter, - h2d_limiter=h2d_limiter, - single_thread_pool=thread_pool, - multi_thread_pool=multi_thread_pool, - ), - cpu_shardings, - tensorstore_specs, - global_shapes, - dtypes, - ) - - async def gather_func(): - return await asyncio.gather(*future_arrays) - - result = asyncio.run(gather_func()) - logging.info("Deserialize took %.2f seconds", time.perf_counter() - start_colocated_time) - return result - - run_deserializer = run_deserializer.specialize( - devices=cpu_devices, - out_specs_fn=output_spec_fn, - ) - - # Try running in the current event loop if one exists, otherwise create new one - result = run_deserializer() - return result - - def create_mesh(mesh_shape=(1, 1, 1, 1, 1, 16, -1)): """Create a JAX mesh for distributed computation.""" inferred_mesh_shape = infer_mesh_shape(mesh_shape) @@ -266,16 +188,26 @@ def cleanup_loaded_arrays(loaded_arrays: list) -> None: print("Cleanup complete.") -def load_model_default( +def load_model( tensorstore_specs: Sequence[Dict[str, Any]], shardings: Sequence[jax.sharding.NamedSharding], global_shapes: Sequence[tuple], dtypes: Sequence[jnp.dtype], + use_colocated_python: bool = False, ): - """Load model using default method (direct to TPU).""" - print("Preloading checkpoint to TPU HBM...") - start_time = time.perf_counter() + """Load model from checkpoint. + Args: + tensorstore_specs: TensorStore specifications for each array. + shardings: Target shardings for the restored arrays. + global_shapes: Global shapes for each array. + dtypes: Data types for each array. + use_colocated_python: If True, load to CPU first then transfer to TPU. + If False, load directly to TPU. + + Returns: + List of restored JAX arrays. + """ manager = GlobalAsyncCheckpointManager() restored_values = manager.deserialize( shardings=shardings, @@ -283,47 +215,9 @@ def load_model_default( global_shapes=global_shapes, dtypes=dtypes, concurrent_gb=192, + use_colocated_python=use_colocated_python, ) - - preload_time = time.perf_counter() - start_time - print(f"Preload completed in {preload_time:.2f} seconds") - print(f"Preloaded {len(restored_values)} arrays") - - return restored_values - - -def load_model_colocated( - tensorstore_specs: Sequence[Dict[str, Any]], - shardings: Sequence[jax.sharding.NamedSharding], - global_shapes: Sequence[tuple], - dtypes: Sequence[jnp.dtype], -): - """Load model using colocated Python (CPU preload then transfer to TPU).""" - print("Preloading checkpoint to CPU memory...") - start_time = time.perf_counter() - - preloaded_values = _colocated_deserialize( - shardings=shardings, - tensorstore_specs=tensorstore_specs, - global_shapes=global_shapes, - dtypes=dtypes, - ) - # for x in preloaded_values: - # x.block_until_ready() - - preload_time = time.perf_counter() - start_time - print(f"Preload completed in {preload_time:.2f} seconds") - print(f"Preloaded {len(preloaded_values)} arrays") - - print("Transferring arrays to TPU...") - start_time = time.perf_counter() - - restored_values = [jax.device_put(x, s) for x, s in zip(preloaded_values, shardings)] - for x in restored_values: - x.block_until_ready() - - transfer_time = time.perf_counter() - start_time - print(f"Transfer completed in {transfer_time:.2f} seconds") + print(f"Loaded {len(restored_values)} arrays") return restored_values @@ -359,13 +253,6 @@ def main(): print(f"JAX devices: {jax.devices()}") - # Select loading function and profile prefix based on method - if args.method == "colocated": - loader_fn = load_model_colocated - else: # args.method == "default" - loader_fn = load_model_default - print(f"--- Running {args.method} benchmark ---") - # Validate checkpoint path if not args.ckpt_path.startswith("gs://"): raise ValueError(f"Only GCS paths (gs://) are supported, got: {args.ckpt_path}") @@ -385,6 +272,7 @@ def main(): state_spec = create_state_spec_from_checkpoint(args.ckpt_path) print(f"Found {len(jax.tree_util.tree_leaves(state_spec))} tensors in checkpoint") + print(f"--- Running {args.method} benchmark ---") loaded_values = None try: with create_mesh(): @@ -395,14 +283,15 @@ def main(): with maybe_profile(args.profile, profile_dir): start_time = time.perf_counter() - loaded_values = loader_fn( + loaded_values = load_model( tensorstore_specs=tensorstore_specs, shardings=shardings, global_shapes=global_shapes, dtypes=dtypes, + use_colocated_python=(args.method == "colocated"), ) print(f"✅ Successfully loaded model from {args.ckpt_path}") - print(f"Deserialize took {time.perf_counter() - start_time:.2f} seconds") + print(f"Total time took {time.perf_counter() - start_time:.2f} seconds") print(f" Total parameters: {sum(x.size for x in loaded_values):,}") finally: # Always clean up, even if benchmark fails diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index db78cfa27..63bd04d5c 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -35,6 +35,8 @@ from absl import logging from jax._src import array, typing from jax._src.layout import Format +from jax._src.mesh import thread_resources +from jax.experimental import colocated_python from jax.experimental.array_serialization import serialization from axlearn.common.utils import Tensor @@ -569,6 +571,154 @@ async def cb(index: array.Index, device: jax.Device): ) +def _create_cpu_shardings( + cpu_devices: list, + tpu_shardings: Sequence[Union[jax.sharding.Sharding, Format]], + tpu_mesh: jax.sharding.Mesh, +) -> list[jax.sharding.Sharding]: + """Create CPU shardings that mirror the structure of TPU shardings. + + Args: + cpu_devices: List of CPU devices. + tpu_shardings: Target TPU shardings to mirror. + tpu_mesh: TPU mesh for creating multi-device CPU mesh. + + Returns: + List of CPU shardings matching the TPU sharding structure. + """ + if len(cpu_devices) > 1: + logging.info("Creating CPU mesh from TPU mesh: %s", tpu_mesh) + cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) + logging.info("CPU Mesh: %s", cpu_mesh) + + return [ + ( + jax.sharding.NamedSharding(cpu_mesh, sharding.spec) + if isinstance(sharding, jax.sharding.NamedSharding) + else jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) + ) + for sharding in tpu_shardings + ] + else: + # Single device - use simple single device sharding + return [jax.sharding.SingleDeviceSharding(cpu_devices[0]) for _ in tpu_shardings] + + +async def _run_colocated_deserialize( + shardings: Sequence[Union[jax.sharding.Sharding, Format]], + tensorstore_specs: Sequence[dict[str, Any]], + global_shapes: Sequence[tuple], + dtypes: Sequence[typing.DTypeLike], + *, + concurrent_bytes: int, + tpu_mesh: jax.sharding.Mesh, +): + """Deserialize checkpoint to CPU using colocated python, then transfer to TPU. + + This approach offloads checkpoint loading from TPU to CPU, which can improve + performance for large checkpoints. + + Args: + shardings: Target TPU shardings for the restored arrays. + tensorstore_specs: TensorStore specifications for each array. + global_shapes: Global shapes for each array. + dtypes: Data types for each array. + concurrent_bytes: Maximum concurrent bytes for reading from storage. + tpu_mesh: TPU mesh for creating CPU mesh. Should be captured from the main thread + where the mesh context is active. + + Returns: + List of JAX arrays on TPU devices. + """ + cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) + logging.info("Colocated CPU devices: %s", cpu_devices) + + # Create CPU shardings matching the TPU sharding structure + cpu_shardings = _create_cpu_shardings(cpu_devices, shardings, tpu_mesh) + + def output_spec_fn(): + return [ + jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=cpu_sharding) + for shape, dtype, cpu_sharding in zip(global_shapes, dtypes, cpu_shardings) + ] + + @colocated_python.colocated_python + def colocated_deserializer(): + logging.info("Starting colocated deserialization") + start_time = time.perf_counter() + + # Create limiters and thread pools in the CPU worker context. + # Note: Use concurrent_bytes for h2d_limiter since H2D happens from CPU worker to TPU, + # not within the same host process, so premapped buffer size doesn't apply. + # pylint: disable=protected-access + byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + h2d_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + + # Use context managers to ensure thread pools are properly shut down + with ( + ThreadPoolExecutor(1) as single_thread_pool, + ThreadPoolExecutor(max_workers=int(os.cpu_count() * 0.8)) as multi_thread_pool, + ): + future_arrays = jax.tree.map( + functools.partial( + _async_deserialize, + byte_limiter=byte_limiter, + h2d_limiter=h2d_limiter, + single_thread_pool=single_thread_pool, + multi_thread_pool=multi_thread_pool, + ), + cpu_shardings, + tensorstore_specs, + global_shapes, + dtypes, + ) + + # Wrap in async function so gather() is called after event loop is created + async def gather_func(): + return await asyncio.gather(*future_arrays) + + result = asyncio.run(gather_func()) + logging.info( + "Colocated deserialization to CPU took %.2f seconds", + time.perf_counter() - start_time, + ) + return result + + # Specialize the colocated function with CPU devices and output specs + colocated_deserializer = colocated_deserializer.specialize( + devices=cpu_devices, + out_specs_fn=output_spec_fn, + ) + + # Run deserializer to load checkpoint to CPU + logging.info("Loading checkpoint to CPU memory...") + start_time = time.perf_counter() + cpu_arrays = colocated_deserializer() + # Block until all transfers complete + for arr in cpu_arrays: + arr.block_until_ready() + cpu_load_time = time.perf_counter() - start_time + logging.info("CPU load completed in %.2f seconds", cpu_load_time) + logging.info("Preloaded %d arrays.", len(cpu_arrays)) + + # Transfer from CPU to TPU + logging.info("Transferring arrays from CPU to TPU...") + start_time = time.perf_counter() + tpu_arrays = [jax.device_put(arr, sharding) for arr, sharding in zip(cpu_arrays, shardings)] + + # Block until all transfers complete + for arr in tpu_arrays: + arr.block_until_ready() + + transfer_time = time.perf_counter() - start_time + logging.info("CPU to TPU transfer completed in %.2f seconds", transfer_time) + logging.info( + "Total colocated deserialization time: %.2f seconds", cpu_load_time + transfer_time + ) + + return tpu_arrays + + # Reference: # https://github.com/google/orbax/blob/ebb3e6d75f9ccb52bf862f1740943a45b18f4dac/checkpoint/orbax/checkpoint/future.py#L49 class _ThreadRaisingException(threading.Thread): @@ -677,35 +827,76 @@ def deserialize( global_shapes: Optional[Sequence[array.Shape]] = None, dtypes: Optional[Sequence[typing.DTypeLike]] = None, concurrent_gb: int = 32, + use_colocated_python: bool = False, ): + """Deserialize arrays from TensorStore. + + Args: + shardings: Sharding specifications for each array. + tensorstore_specs: TensorStore specifications for each array. + global_shapes: Global shapes for each array. If None, uses shape from TensorStore. + dtypes: Data types for each array. If None, uses dtype from TensorStore. + concurrent_gb: Maximum concurrent GB for reading from storage. + use_colocated_python: If True, use colocated Python to load checkpoint to CPU first, + then transfer to TPU. + Only works on systems with Colocated Python support (e.g., Pathways). + + Returns: + List of deserialized JAX arrays. + """ self.wait_until_finished() start_time = time.perf_counter() concurrent_bytes = concurrent_gb * 10**9 - async def _run_deserializer(): - # Object should be created once per process. - # pylint: disable=protected-access - byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - h2d_limiter = serialization._LimitInFlightBytes(_get_premapped_buffer_size()) - - future_arrays = jax.tree.map( - functools.partial( - _async_deserialize, - byte_limiter=byte_limiter, - h2d_limiter=h2d_limiter, - single_thread_pool=self._single_thread_pool, - multi_thread_pool=self._multi_thread_pool, - ), - shardings, - tensorstore_specs, - [None] * len(tensorstore_specs) if global_shapes is None else global_shapes, - [None] * len(tensorstore_specs) if dtypes is None else dtypes, + # Prepare global shapes and dtypes + global_shapes_list = ( + [None] * len(tensorstore_specs) if global_shapes is None else list(global_shapes) + ) + dtypes_list = [None] * len(tensorstore_specs) if dtypes is None else list(dtypes) + + # Create the appropriate coroutine based on mode + if use_colocated_python: + # Capture mesh from main thread where context is active + # (mesh is thread-local and won't be available in background event loop) + tpu_mesh = thread_resources.env.physical_mesh + + # Use colocated Python path + coro = _run_colocated_deserialize( + shardings=shardings, + tensorstore_specs=tensorstore_specs, + global_shapes=global_shapes_list, + dtypes=dtypes_list, + concurrent_bytes=concurrent_bytes, + tpu_mesh=tpu_mesh, ) - return await asyncio.gather(*future_arrays) + else: + # Use default deserialization path (direct to TPU) + async def _run_deserializer(): + # Object should be created once per process. + # pylint: disable=protected-access + byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + h2d_limiter = serialization._LimitInFlightBytes(_get_premapped_buffer_size()) + + future_arrays = jax.tree.map( + functools.partial( + _async_deserialize, + byte_limiter=byte_limiter, + h2d_limiter=h2d_limiter, + single_thread_pool=self._single_thread_pool, + multi_thread_pool=self._multi_thread_pool, + ), + shardings, + tensorstore_specs, + global_shapes_list, + dtypes_list, + ) + return await asyncio.gather(*future_arrays) + + coro = _run_deserializer() - fut = asyncio.run_coroutine_threadsafe(_run_deserializer(), self._loop) - result = fut.result() + # Run the coroutine and get result + result = asyncio.run_coroutine_threadsafe(coro, self._loop).result() logging.info("deserialize took %.4f seconds.", time.perf_counter() - start_time) return result From 259dca2a39fad031183320d7788677c7c08f002b Mon Sep 17 00:00:00 2001 From: "Chang (AIML) Liu" Date: Wed, 25 Feb 2026 14:32:11 -0800 Subject: [PATCH 2/5] [3/n] Use limiter to control cpu to tpu transfer to avoid OOM. (#2189) * [3/n] Use limiter to control cpu to tpu transfer to avoid OOM. * Use explicit profile call to trigger the patch properly. * Add tensor statics in the benchmark script. * Add more logging for located python. * Add accurate compute for per device bytes. --- .../examples/colocated_python_benchmark.py | 72 ++++++-- axlearn/common/array_serialization.py | 156 +++++++++++++++++- 2 files changed, 209 insertions(+), 19 deletions(-) diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index 638dfa0a6..bb987c4b6 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -21,7 +21,7 @@ import argparse import os import time -from contextlib import nullcontext +from contextlib import contextmanager from datetime import datetime from typing import Any, Dict, Optional, Sequence @@ -38,22 +38,22 @@ from axlearn.common.utils import TensorSpec, infer_mesh_shape +@contextmanager def maybe_profile(enabled: bool, profile_dir: Optional[str]): - """Return JAX profiler context if enabled, otherwise a no-op context manager. + """JAX profiler context if enabled, otherwise a no-op. Args: enabled: Whether profiling is enabled. profile_dir: Directory to save profiling results. - - Returns: - Context manager for profiling or no-op. """ if enabled: assert profile_dir is not None, "profile_dir must be set when profiling is enabled" - return jax.profiler.trace(profile_dir) - else: - # Return a no-op context manager - return nullcontext() + jax.profiler.start_trace(profile_dir) + try: + yield + finally: + if enabled: + jax.profiler.stop_trace() def create_mesh(mesh_shape=(1, 1, 1, 1, 1, 16, -1)): @@ -132,6 +132,13 @@ def create_checkpoint_spec_from_state(ckpt_dir: str, state_spec: dict): if not mesh.shape: raise RuntimeError("Checkpoint restoration must take place within the context of a Mesh") + # Track sharding statistics + sharded_bytes = 0 + replicated_bytes = 0 + per_shard_bytes = 0 + num_sharded = 0 + num_replicated = 0 + # Process each tensor in the state spec for path, value in utils.flatten_items(state_spec, separator="/"): if isinstance(value, TensorSpec): @@ -153,6 +160,40 @@ def create_checkpoint_spec_from_state(ckpt_dir: str, state_spec: dict): dtypes.append(dtype) shardings.append(sharding) + # Compute tensor size in bytes downloaded from GCS. + # Checkpoints are stored as fp32 (4 bytes per element). + element_size = 4 + tensor_bytes = 1 + for d in value.shape: + tensor_bytes *= d + tensor_bytes *= element_size + + if partition_spec == jax.sharding.PartitionSpec(): + replicated_bytes += tensor_bytes + num_replicated += 1 + else: + sharded_bytes += tensor_bytes + num_sharded += 1 + # Compute per-shard size by dividing by the number of shards + num_shards = 1 + for axis in partition_spec: + if axis is not None: + if isinstance(axis, tuple): + for a in axis: + num_shards *= mesh.shape[a] + else: + num_shards *= mesh.shape[axis] + per_shard_bytes += tensor_bytes // num_shards + + num_devices = len(mesh.devices.flat) + print(f"Sharding stats ({num_devices} devices):") + print( + f" Sharded: {num_sharded} tensors, {sharded_bytes / 10**9:.2f} GB " + f"(per-shard total: {per_shard_bytes / 10**9:.2f} GB)" + ) + print(f" Replicated: {num_replicated} tensors, {replicated_bytes / 10**9:.2f} GB") + print(f" Per-device total: {(per_shard_bytes + replicated_bytes) / 10**9:.2f} GB") + return tensorstore_specs, shardings, global_shapes, dtypes @@ -194,6 +235,7 @@ def load_model( global_shapes: Sequence[tuple], dtypes: Sequence[jnp.dtype], use_colocated_python: bool = False, + transfer_concurrent_gb: int = 16, ): """Load model from checkpoint. @@ -204,6 +246,8 @@ def load_model( dtypes: Data types for each array. use_colocated_python: If True, load to CPU first then transfer to TPU. If False, load directly to TPU. + transfer_concurrent_gb: Maximum concurrent GB in flight during CPU to TPU transfer + when using colocated Python. Defaults to 16GB. Returns: List of restored JAX arrays. @@ -214,8 +258,9 @@ def load_model( tensorstore_specs=tensorstore_specs, global_shapes=global_shapes, dtypes=dtypes, - concurrent_gb=192, + concurrent_gb=400, use_colocated_python=use_colocated_python, + transfer_concurrent_gb=transfer_concurrent_gb, ) print(f"Loaded {len(restored_values)} arrays") @@ -240,6 +285,12 @@ def main(): action="store_true", help="Enable JAX profiler (adds overhead, disable for accurate benchmarking)", ) + parser.add_argument( + "--transfer_concurrent_gb", + type=int, + default=16, + help="Maximum concurrent GB during CPU to TPU transfer (colocated mode only). Default: 16", + ) args = parser.parse_args() # Disable persistent compilation cache for fair benchmarking @@ -289,6 +340,7 @@ def main(): global_shapes=global_shapes, dtypes=dtypes, use_colocated_python=(args.method == "colocated"), + transfer_concurrent_gb=args.transfer_concurrent_gb, ) print(f"✅ Successfully loaded model from {args.ckpt_path}") print(f"Total time took {time.perf_counter() - start_time:.2f} seconds") diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index 63bd04d5c..51126bdff 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -404,8 +404,17 @@ async def _run_serializer( raise e -def _blocking_device_put(out: Tensor, layout: Format) -> Tensor: - return jax.block_until_ready(jax.device_put(out, layout)) +def _blocking_device_put(tensor: Tensor, target: Union[Format, jax.sharding.Sharding]) -> Tensor: + """Device put and block until ready. + + Args: + tensor: Array to transfer. + target: Either a Format (with layout + sharding) or Sharding. + + Returns: + Transferred array. + """ + return jax.block_until_ready(jax.device_put(tensor, target)) async def _async_deserialize( @@ -473,7 +482,10 @@ async def _async_deserialize( # which only runs on GCP for now. context = serialization.TS_CONTEXT if os.getenv("ENABLE_GCS_GRPC", "false") == "true": + logging.debug("gcs_grpc enabled") tensorstore_spec, context = use_gcs_grpc(tensorstore_spec) + else: + logging.debug("gcs_grpc not enabled") t = await ts.open( tensorstore_spec, @@ -571,6 +583,112 @@ async def cb(index: array.Index, device: jax.Device): ) +async def _async_transfer_cpu_to_tpu( + tensor, + sharding: jax.sharding.NamedSharding, + global_limiter: serialization._LimitInFlightBytes, + single_thread_pool: ThreadPoolExecutor, + multi_thread_pool: ThreadPoolExecutor, +): + """Async worker to transfer a single array from CPU to TPU with rate limiting. + + Uses a global limiter to control concurrent transfer load. + Specifically designed for colocated Python CPU→TPU transfers. + + Args: + tensor: The CPU array to transfer to TPU. + sharding: The target TPU sharding for the array. + global_limiter: Global byte limiter for all transfers. + single_thread_pool: Thread pool with single worker for oversized transfers. + multi_thread_pool: Thread pool with multiple workers for normal transfers. + + Returns: + The transferred array on TPU. + """ + # Calculate bytes per device shard from the sharding spec. + shard_shape = sharding.shard_shape(tensor.shape) + bytes_per_device = math.prod(shard_shape) * tensor.dtype.itemsize + + # Check if we should skip limiting (similar to h2d_limiter pattern) + # pylint: disable-next=protected-access + max_capacity = global_limiter._max_bytes + + # Skip limiting if per-device bytes exceed capacity + if bytes_per_device > max_capacity: + logging.log_first_n( + logging.WARNING, + "Array shard size per device (%d bytes, %.2f GB) exceeded " + "limiter capacity (%d bytes, %.2f GB). Skipping rate limiting for this transfer.", + 5, + bytes_per_device, + bytes_per_device / (1024**3), + max_capacity, + max_capacity / (1024**3), + ) + # Perform transfer without limiting using single thread pool + # (serialize large transfers to avoid overwhelming memory) + loop = asyncio.get_event_loop() + result = await loop.run_in_executor( + single_thread_pool, _blocking_device_put, tensor, sharding + ) + return result + + # Acquire limiter once for the array (reserves bytes_per_device) + await global_limiter.wait_for_bytes(bytes_per_device) + + try: + # Perform the device_put using multi thread pool for parallelism + loop = asyncio.get_event_loop() + result = await loop.run_in_executor( + multi_thread_pool, _blocking_device_put, tensor, sharding + ) + return result + finally: + # Release bytes + await global_limiter.release_bytes(bytes_per_device) + + +async def _transfer_arrays_cpu_to_tpu( + arrays: Sequence, + shardings: Sequence[jax.sharding.NamedSharding], + concurrent_bytes: int, + single_thread_pool: ThreadPoolExecutor, + multi_thread_pool: ThreadPoolExecutor, +): + """Transfer multiple arrays from CPU to TPU with global rate limiting. + + Uses a single global limiter to control total concurrent transfer load. + Designed for colocated Python checkpoint loading workflow. + + Args: + arrays: List of CPU arrays to transfer. + shardings: List of target TPU shardings for each array. + concurrent_bytes: Maximum concurrent bytes in flight globally. + single_thread_pool: Thread pool with single worker for oversized transfers. + multi_thread_pool: Thread pool with multiple workers for normal transfers. + + Returns: + List of transferred arrays on TPU. + """ + # Create a single global limiter + # pylint: disable=protected-access + global_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + + logging.info( + "Created global transfer limiter with %.1fGB capacity", + concurrent_bytes / (1024**3), + ) + + # Create async tasks for all transfers + tasks = [ + _async_transfer_cpu_to_tpu( + array, sharding, global_limiter, single_thread_pool, multi_thread_pool + ) + for array, sharding in zip(arrays, shardings) + ] + return await asyncio.gather(*tasks) + + def _create_cpu_shardings( cpu_devices: list, tpu_shardings: Sequence[Union[jax.sharding.Sharding, Format]], @@ -612,6 +730,9 @@ async def _run_colocated_deserialize( *, concurrent_bytes: int, tpu_mesh: jax.sharding.Mesh, + transfer_concurrent_bytes: int, + single_thread_pool: ThreadPoolExecutor, + multi_thread_pool: ThreadPoolExecutor, ): """Deserialize checkpoint to CPU using colocated python, then transfer to TPU. @@ -626,6 +747,10 @@ async def _run_colocated_deserialize( concurrent_bytes: Maximum concurrent bytes for reading from storage. tpu_mesh: TPU mesh for creating CPU mesh. Should be captured from the main thread where the mesh context is active. + transfer_concurrent_bytes: Maximum concurrent bytes in flight during + CPU to TPU transfer. + single_thread_pool: Thread pool with single worker for oversized transfers. + multi_thread_pool: Thread pool with multiple workers for normal transfers. Returns: List of JAX arrays on TPU devices. @@ -655,9 +780,18 @@ def colocated_deserializer(): h2d_limiter = serialization._LimitInFlightBytes(concurrent_bytes) # Use context managers to ensure thread pools are properly shut down + num_workers = int(os.cpu_count() * 0.8) + logging.info( + "Colocated Python: os.cpu_count()=%s, multi_thread_pool workers=%d, " + "byte_limiter=%d GB, h2d_limiter=%d GB", + os.cpu_count(), + num_workers, + concurrent_bytes // 10**9, + concurrent_bytes // 10**9, + ) with ( ThreadPoolExecutor(1) as single_thread_pool, - ThreadPoolExecutor(max_workers=int(os.cpu_count() * 0.8)) as multi_thread_pool, + ThreadPoolExecutor(max_workers=num_workers) as multi_thread_pool, ): future_arrays = jax.tree.map( functools.partial( @@ -701,14 +835,12 @@ async def gather_func(): logging.info("CPU load completed in %.2f seconds", cpu_load_time) logging.info("Preloaded %d arrays.", len(cpu_arrays)) - # Transfer from CPU to TPU + # Transfer from CPU to TPU with global rate limiting logging.info("Transferring arrays from CPU to TPU...") start_time = time.perf_counter() - tpu_arrays = [jax.device_put(arr, sharding) for arr, sharding in zip(cpu_arrays, shardings)] - - # Block until all transfers complete - for arr in tpu_arrays: - arr.block_until_ready() + tpu_arrays = await _transfer_arrays_cpu_to_tpu( + cpu_arrays, shardings, transfer_concurrent_bytes, single_thread_pool, multi_thread_pool + ) transfer_time = time.perf_counter() - start_time logging.info("CPU to TPU transfer completed in %.2f seconds", transfer_time) @@ -828,6 +960,7 @@ def deserialize( dtypes: Optional[Sequence[typing.DTypeLike]] = None, concurrent_gb: int = 32, use_colocated_python: bool = False, + transfer_concurrent_gb: int = 16, ): """Deserialize arrays from TensorStore. @@ -840,6 +973,8 @@ def deserialize( use_colocated_python: If True, use colocated Python to load checkpoint to CPU first, then transfer to TPU. Only works on systems with Colocated Python support (e.g., Pathways). + transfer_concurrent_gb: Maximum concurrent GB in flight during + CPU to TPU transfer when using colocated Python. Defaults to 16GB. Returns: List of deserialized JAX arrays. @@ -869,6 +1004,9 @@ def deserialize( dtypes=dtypes_list, concurrent_bytes=concurrent_bytes, tpu_mesh=tpu_mesh, + transfer_concurrent_bytes=transfer_concurrent_gb * 10**9, + single_thread_pool=self._single_thread_pool, + multi_thread_pool=self._multi_thread_pool, ) else: # Use default deserialization path (direct to TPU) From bb5d51373c46bd94e1ed3db1fcec5d21a82b39bf Mon Sep 17 00:00:00 2001 From: "Chang (AIML) Liu" Date: Fri, 27 Feb 2026 10:31:44 -0800 Subject: [PATCH 3/5] [4/n] Use global limiter to control the model loading pipeline to further reduce memory usage (#2214) * [4/n] Use global limiter to control the full model loading pipeline to further reduce memory usage. * Fix the index array issue. Plus minor optimization of the thread pool. * Fix parallel execution of colocated python. * more accurate byte per device computation. * Enable colocated python by default when it's available. Controlled through ENV when it's available. Controlled through ENV.. * simplify limiter usage in load_and_transfer_one * Resolve devx-ai's comments. * fix global_shape issue * resolve Ethan's comment. * minor tensorspec open improvement * resolve Ethan's new comments. And minor improvement in case ts has no read chunk. --- .../examples/colocated_python_benchmark.py | 18 +- axlearn/common/array_serialization.py | 413 +++++++++--------- 2 files changed, 209 insertions(+), 222 deletions(-) diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index bb987c4b6..18a50bedb 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -234,8 +234,6 @@ def load_model( shardings: Sequence[jax.sharding.NamedSharding], global_shapes: Sequence[tuple], dtypes: Sequence[jnp.dtype], - use_colocated_python: bool = False, - transfer_concurrent_gb: int = 16, ): """Load model from checkpoint. @@ -244,10 +242,6 @@ def load_model( shardings: Target shardings for the restored arrays. global_shapes: Global shapes for each array. dtypes: Data types for each array. - use_colocated_python: If True, load to CPU first then transfer to TPU. - If False, load directly to TPU. - transfer_concurrent_gb: Maximum concurrent GB in flight during CPU to TPU transfer - when using colocated Python. Defaults to 16GB. Returns: List of restored JAX arrays. @@ -259,8 +253,6 @@ def load_model( global_shapes=global_shapes, dtypes=dtypes, concurrent_gb=400, - use_colocated_python=use_colocated_python, - transfer_concurrent_gb=transfer_concurrent_gb, ) print(f"Loaded {len(restored_values)} arrays") @@ -285,12 +277,6 @@ def main(): action="store_true", help="Enable JAX profiler (adds overhead, disable for accurate benchmarking)", ) - parser.add_argument( - "--transfer_concurrent_gb", - type=int, - default=16, - help="Maximum concurrent GB during CPU to TPU transfer (colocated mode only). Default: 16", - ) args = parser.parse_args() # Disable persistent compilation cache for fair benchmarking @@ -333,14 +319,14 @@ def main(): ) with maybe_profile(args.profile, profile_dir): + if args.method == "default": + os.environ["COLOCATED_PYTHON_DESERIALIZE"] = "0" start_time = time.perf_counter() loaded_values = load_model( tensorstore_specs=tensorstore_specs, shardings=shardings, global_shapes=global_shapes, dtypes=dtypes, - use_colocated_python=(args.method == "colocated"), - transfer_concurrent_gb=args.transfer_concurrent_gb, ) print(f"✅ Successfully loaded model from {args.ckpt_path}") print(f"Total time took {time.perf_counter() - start_time:.2f} seconds") diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index 51126bdff..bd6da7f50 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -213,6 +213,22 @@ def running_on_pathways(): return os.getenv("JAX_PLATFORMS") == "proxy" +def colocated_python_available() -> bool: + """Returns True if colocated Python CPU workers are available at runtime. + + Colocated Python requires: + 1. Running on Pathways (JAX_PLATFORMS == "proxy"). + 2. At least one colocated CPU device is present (i.e. the sidecar is running). + """ + if not running_on_pathways(): + return False + try: + cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) + return len(cpu_devices) > 0 + except Exception: # pylint: disable=broad-exception-caught + return False + + async def _slice_shard_and_copy_to_host(shard_infos: list[_ShardInfo]): """Slices each shard according to shard info and then copy the sliced result to host. @@ -423,10 +439,10 @@ async def _async_deserialize( global_shape: Optional[Sequence[int]], dtype: Optional[typing.DTypeLike], *, - h2d_limiter: serialization._LimitInFlightBytes, + h2d_limiter: Optional[serialization._LimitInFlightBytes], byte_limiter: serialization._LimitInFlightBytes, - single_thread_pool: ThreadPoolExecutor, - multi_thread_pool: ThreadPoolExecutor, + single_thread_pool: Optional[ThreadPoolExecutor], + multi_thread_pool: Optional[ThreadPoolExecutor], ): """Modified from https://github.com/jax-ml/jax/blob/e7ec418eba9ada336f755613948cbdf4a9e97d59/jax/experimental/array_serialization/serialization.py#L345 @@ -443,6 +459,9 @@ async def _async_deserialize( in flight H2D is imposed. Note that Pathways checkpoint loading does not require h2d limiter since the H2D doesn't happen in the head node, and each worker has preemapped a chunk of host memory that is larger than the total device memory. + h2d_limiter, single_thread_pool, and multi_thread_pool are all optional. When omitted (e.g. + in the colocated Python path where a global pipeline limiter already bounds concurrency), + H2D is submitted to the default asyncio thread pool without per-shard gating. 4. Let user pass in a multi_thread_pool thread pool for ckpt loading, instead of letting async io to create a default pool, to make it more configurable. @@ -542,10 +561,14 @@ async def cb(index: array.Index, device: jax.Device): dll, jax.sharding.SingleDeviceSharding(device, memory_kind=in_sharding.memory_kind) ) + if h2d_limiter is None: + # No per-shard H2D gating — caller (e.g. colocated Python path) relies on a global + # pipeline limiter to bound concurrency instead. + result = await loop.run_in_executor(None, _blocking_device_put, out, layout) # Jax >= 0.6.2 changes the behavior of _LimitInFlightBytes, where wait_for_bytes no longer # throws an exception if requested_bytes > max_bytes # pylint: disable-next=protected-access - if out_size > h2d_limiter._max_bytes: + elif out_size > h2d_limiter._max_bytes: logging.log_first_n( logging.WARNING, "Tensor shard for tensor %s (padded size %d bytes) exceeded " @@ -583,112 +606,6 @@ async def cb(index: array.Index, device: jax.Device): ) -async def _async_transfer_cpu_to_tpu( - tensor, - sharding: jax.sharding.NamedSharding, - global_limiter: serialization._LimitInFlightBytes, - single_thread_pool: ThreadPoolExecutor, - multi_thread_pool: ThreadPoolExecutor, -): - """Async worker to transfer a single array from CPU to TPU with rate limiting. - - Uses a global limiter to control concurrent transfer load. - Specifically designed for colocated Python CPU→TPU transfers. - - Args: - tensor: The CPU array to transfer to TPU. - sharding: The target TPU sharding for the array. - global_limiter: Global byte limiter for all transfers. - single_thread_pool: Thread pool with single worker for oversized transfers. - multi_thread_pool: Thread pool with multiple workers for normal transfers. - - Returns: - The transferred array on TPU. - """ - # Calculate bytes per device shard from the sharding spec. - shard_shape = sharding.shard_shape(tensor.shape) - bytes_per_device = math.prod(shard_shape) * tensor.dtype.itemsize - - # Check if we should skip limiting (similar to h2d_limiter pattern) - # pylint: disable-next=protected-access - max_capacity = global_limiter._max_bytes - - # Skip limiting if per-device bytes exceed capacity - if bytes_per_device > max_capacity: - logging.log_first_n( - logging.WARNING, - "Array shard size per device (%d bytes, %.2f GB) exceeded " - "limiter capacity (%d bytes, %.2f GB). Skipping rate limiting for this transfer.", - 5, - bytes_per_device, - bytes_per_device / (1024**3), - max_capacity, - max_capacity / (1024**3), - ) - # Perform transfer without limiting using single thread pool - # (serialize large transfers to avoid overwhelming memory) - loop = asyncio.get_event_loop() - result = await loop.run_in_executor( - single_thread_pool, _blocking_device_put, tensor, sharding - ) - return result - - # Acquire limiter once for the array (reserves bytes_per_device) - await global_limiter.wait_for_bytes(bytes_per_device) - - try: - # Perform the device_put using multi thread pool for parallelism - loop = asyncio.get_event_loop() - result = await loop.run_in_executor( - multi_thread_pool, _blocking_device_put, tensor, sharding - ) - return result - finally: - # Release bytes - await global_limiter.release_bytes(bytes_per_device) - - -async def _transfer_arrays_cpu_to_tpu( - arrays: Sequence, - shardings: Sequence[jax.sharding.NamedSharding], - concurrent_bytes: int, - single_thread_pool: ThreadPoolExecutor, - multi_thread_pool: ThreadPoolExecutor, -): - """Transfer multiple arrays from CPU to TPU with global rate limiting. - - Uses a single global limiter to control total concurrent transfer load. - Designed for colocated Python checkpoint loading workflow. - - Args: - arrays: List of CPU arrays to transfer. - shardings: List of target TPU shardings for each array. - concurrent_bytes: Maximum concurrent bytes in flight globally. - single_thread_pool: Thread pool with single worker for oversized transfers. - multi_thread_pool: Thread pool with multiple workers for normal transfers. - - Returns: - List of transferred arrays on TPU. - """ - # Create a single global limiter - # pylint: disable=protected-access - global_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - - logging.info( - "Created global transfer limiter with %.1fGB capacity", - concurrent_bytes / (1024**3), - ) - - # Create async tasks for all transfers - tasks = [ - _async_transfer_cpu_to_tpu( - array, sharding, global_limiter, single_thread_pool, multi_thread_pool - ) - for array, sharding in zip(arrays, shardings) - ] - return await asyncio.gather(*tasks) - - def _create_cpu_shardings( cpu_devices: list, tpu_shardings: Sequence[Union[jax.sharding.Sharding, Format]], @@ -722,6 +639,37 @@ def _create_cpu_shardings( return [jax.sharding.SingleDeviceSharding(cpu_devices[0]) for _ in tpu_shardings] +async def _effective_bytes_per_device( + spec: dict[str, Any], + shape: tuple, + dtype: typing.DTypeLike, + sharding: jax.sharding.Sharding, +) -> int: + """Estimates effective bytes read from storage per device, accounting for chunk overhead. + + TensorStore reads whole chunks even if only a partial chunk is needed, so actual bytes + read can exceed the raw tensor shard size. + """ + t = await ts.open(ts.Spec(spec), open=True) + shard_shape = sharding.shard_shape(shape) + raw_bytes = math.prod(shard_shape) * np.dtype(dtype).itemsize + + read_chunk = t.chunk_layout.read_chunk + if read_chunk is None or read_chunk.shape is None: + return raw_bytes + + chunk_shape = read_chunk.shape + overhead_ratio = 1.0 + for dim_idx in range(len(shape)): + cs = chunk_shape[dim_idx] + needed = shard_shape[dim_idx] + if needed > 0: + num_chunks = (needed + cs - 1) // cs + overhead_ratio *= (num_chunks * cs) / needed + + return int(raw_bytes * overhead_ratio) + + async def _run_colocated_deserialize( shardings: Sequence[Union[jax.sharding.Sharding, Format]], tensorstore_specs: Sequence[dict[str, Any]], @@ -730,27 +678,26 @@ async def _run_colocated_deserialize( *, concurrent_bytes: int, tpu_mesh: jax.sharding.Mesh, - transfer_concurrent_bytes: int, - single_thread_pool: ThreadPoolExecutor, + pipeline_concurrent_bytes: int, multi_thread_pool: ThreadPoolExecutor, ): - """Deserialize checkpoint to CPU using colocated python, then transfer to TPU. + """Deserialize checkpoint with pipelined load to CPU then transfer to TPU. - This approach offloads checkpoint loading from TPU to CPU, which can improve - performance for large checkpoints. + This approach uses colocated Python to load checkpoint data to CPU worker, then + transfers each array to TPU. A global limiter controls the entire pipeline (load + transfer) + to keep CPU memory usage bounded. Args: shardings: Target TPU shardings for the restored arrays. tensorstore_specs: TensorStore specifications for each array. global_shapes: Global shapes for each array. dtypes: Data types for each array. - concurrent_bytes: Maximum concurrent bytes for reading from storage. + concurrent_bytes: Maximum concurrent bytes for reading from storage (used by worker). tpu_mesh: TPU mesh for creating CPU mesh. Should be captured from the main thread where the mesh context is active. - transfer_concurrent_bytes: Maximum concurrent bytes in flight during - CPU to TPU transfer. - single_thread_pool: Thread pool with single worker for oversized transfers. - multi_thread_pool: Thread pool with multiple workers for normal transfers. + pipeline_concurrent_bytes: Maximum concurrent bytes for the entire pipeline. + Controls peak CPU memory usage. + multi_thread_pool: Thread pool for blocking operations (block_until_ready, TPU transfer). Returns: List of JAX arrays on TPU devices. @@ -761,92 +708,133 @@ async def _run_colocated_deserialize( # Create CPU shardings matching the TPU sharding structure cpu_shardings = _create_cpu_shardings(cpu_devices, shardings, tpu_mesh) - def output_spec_fn(): - return [ - jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=cpu_sharding) - for shape, dtype, cpu_sharding in zip(global_shapes, dtypes, cpu_shardings) - ] + # Pre-compute effective per-device bytes in parallel across all arrays. + effective_bytes = list( + await asyncio.gather( + *[ + _effective_bytes_per_device( + tensorstore_specs[i], + global_shapes[i], + dtypes[i], + cpu_sharding, + ) + for i, cpu_sharding in enumerate(cpu_shardings) + ] + ) + ) - @colocated_python.colocated_python - def colocated_deserializer(): - logging.info("Starting colocated deserialization") - start_time = time.perf_counter() + # Global limiter controls the entire pipeline (load + transfer) + # This limits total bytes in flight across all arrays + # pylint: disable=protected-access + global_limiter = serialization._LimitInFlightBytes(pipeline_concurrent_bytes) - # Create limiters and thread pools in the CPU worker context. - # Note: Use concurrent_bytes for h2d_limiter since H2D happens from CPU worker to TPU, - # not within the same host process, so premapped buffer size doesn't apply. + logging.info( + "Created global pipeline limiter with %.1fGB capacity", + pipeline_concurrent_bytes / (10**9), + ) + + # Capture all per-array data as lists in the closure so that load_to_cpu can be defined + # and specialized exactly once. colocated_python only supports jax.Array as inputs, so we + # pass a scalar index array and look up the per-array data inside the function. + # _get_specialized_func is cached on (FunctionInfo, Specialization): with a single function + # object the cache hits for all arrays that share the same output shape/dtype/sharding, + # reducing compilations from O(num_arrays) to O(num_unique_output_specs). + @colocated_python.colocated_python + def load_to_cpu(idx: jax.Array): + i = int(idx) # pylint: disable=protected-access byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - h2d_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - - # Use context managers to ensure thread pools are properly shut down - num_workers = int(os.cpu_count() * 0.8) - logging.info( - "Colocated Python: os.cpu_count()=%s, multi_thread_pool workers=%d, " - "byte_limiter=%d GB, h2d_limiter=%d GB", - os.cpu_count(), - num_workers, - concurrent_bytes // 10**9, - concurrent_bytes // 10**9, - ) - with ( - ThreadPoolExecutor(1) as single_thread_pool, - ThreadPoolExecutor(max_workers=num_workers) as multi_thread_pool, - ): - future_arrays = jax.tree.map( - functools.partial( - _async_deserialize, - byte_limiter=byte_limiter, - h2d_limiter=h2d_limiter, - single_thread_pool=single_thread_pool, - multi_thread_pool=multi_thread_pool, - ), - cpu_shardings, - tensorstore_specs, - global_shapes, - dtypes, + + return asyncio.run( + _async_deserialize( + cpu_shardings[i], + tensorstore_specs[i], + global_shapes[i], + dtypes[i], + byte_limiter=byte_limiter, + h2d_limiter=None, + single_thread_pool=None, + multi_thread_pool=None, ) + ) - # Wrap in async function so gather() is called after event loop is created - async def gather_func(): - return await asyncio.gather(*future_arrays) + # Pre-create replicated index arrays across all cpu_devices. colocated_python requires inputs + # to have one shard per device in the specialized device list. + cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) + idx_sharding = jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) + idx_arrays = [jax.device_put(jnp.array(i), idx_sharding) for i in range(len(shardings))] + + async def _load_and_transfer_one( + idx: int, + tpu_sharding: jax.sharding.Sharding, + dispatch_pool: ThreadPoolExecutor, + ): + """Load one array to CPU via colocated Python, then transfer to TPU.""" + # Specialize per array with explicit out_specs_fn to avoid stale cached output specs + # when arrays have different shapes. _get_specialized_func caches on + # (FunctionInfo, Specialization) so arrays sharing the same (shape, dtype, sharding) + # will reuse the compiled executable. + specialized = load_to_cpu.specialize( + devices=cpu_devices, + out_specs_fn=lambda _: jax.ShapeDtypeStruct( + shape=global_shapes[idx], dtype=dtypes[idx], sharding=cpu_shardings[idx] + ), + ) + loop = asyncio.get_running_loop() + # Each call dispatched from a separate thread enables concurrent sidecar execution. + cpu_array = await loop.run_in_executor(dispatch_pool, specialized, idx_arrays[idx]) + # block_until_ready is run in an executor so the event loop can continue + # dispatching other coroutines while waiting for this array to materialize. + await loop.run_in_executor(multi_thread_pool, cpu_array.block_until_ready) + return await loop.run_in_executor( + multi_thread_pool, + _blocking_device_put, + cpu_array, + tpu_sharding, + ) - result = asyncio.run(gather_func()) - logging.info( - "Colocated deserialization to CPU took %.2f seconds", - time.perf_counter() - start_time, + async def _load_and_transfer_one_rate_limited( + idx: int, + tpu_sharding: jax.sharding.Sharding, + dispatch_pool: ThreadPoolExecutor, + ): + """Wrapper that applies global limiter to the entire load+transfer operation.""" + bytes_per_device = effective_bytes[idx] + # pylint: disable-next=protected-access + max_capacity = global_limiter._max_bytes + + # Clamp to limiter capacity so oversized arrays don't deadlock wait_for_bytes. + if bytes_per_device > max_capacity: + logging.warning( + "Array shard size per device (%.2f GB) exceeded " + "limiter capacity (%.2f GB). Clamping reservation.", + bytes_per_device / (1024**3), + max_capacity / (1024**3), ) - return result + bytes_to_reserve = min(bytes_per_device, max_capacity) - # Specialize the colocated function with CPU devices and output specs - colocated_deserializer = colocated_deserializer.specialize( - devices=cpu_devices, - out_specs_fn=output_spec_fn, - ) + await global_limiter.wait_for_bytes(bytes_to_reserve) + try: + return await _load_and_transfer_one(idx, tpu_sharding, dispatch_pool) + finally: + await global_limiter.release_bytes(bytes_to_reserve) + + # Dedicated pool for colocated_python dispatch. colocated_python only executes calls + # concurrently when dispatched from different threads — a separate pool ensures dispatch + # threads are never starved by blocking work (block_until_ready, TPU transfer) on + # multi_thread_pool. + with ThreadPoolExecutor(max_workers=min(len(shardings), 256)) as dispatch_pool: + tasks = [ + _load_and_transfer_one_rate_limited(idx, tpu_sharding, dispatch_pool) + for idx, tpu_sharding in enumerate(shardings) + ] - # Run deserializer to load checkpoint to CPU - logging.info("Loading checkpoint to CPU memory...") - start_time = time.perf_counter() - cpu_arrays = colocated_deserializer() - # Block until all transfers complete - for arr in cpu_arrays: - arr.block_until_ready() - cpu_load_time = time.perf_counter() - start_time - logging.info("CPU load completed in %.2f seconds", cpu_load_time) - logging.info("Preloaded %d arrays.", len(cpu_arrays)) - - # Transfer from CPU to TPU with global rate limiting - logging.info("Transferring arrays from CPU to TPU...") - start_time = time.perf_counter() - tpu_arrays = await _transfer_arrays_cpu_to_tpu( - cpu_arrays, shardings, transfer_concurrent_bytes, single_thread_pool, multi_thread_pool - ) + logging.info("Starting pipelined colocated load and transfer...") + start_time = time.perf_counter() + tpu_arrays = await asyncio.gather(*tasks) - transfer_time = time.perf_counter() - start_time - logging.info("CPU to TPU transfer completed in %.2f seconds", transfer_time) - logging.info( - "Total colocated deserialization time: %.2f seconds", cpu_load_time + transfer_time - ) + total_time = time.perf_counter() - start_time + logging.info("Pipelined colocated deserialization completed in %.2f seconds", total_time) return tpu_arrays @@ -959,8 +947,6 @@ def deserialize( global_shapes: Optional[Sequence[array.Shape]] = None, dtypes: Optional[Sequence[typing.DTypeLike]] = None, concurrent_gb: int = 32, - use_colocated_python: bool = False, - transfer_concurrent_gb: int = 16, ): """Deserialize arrays from TensorStore. @@ -970,23 +956,31 @@ def deserialize( global_shapes: Global shapes for each array. If None, uses shape from TensorStore. dtypes: Data types for each array. If None, uses dtype from TensorStore. concurrent_gb: Maximum concurrent GB for reading from storage. - use_colocated_python: If True, use colocated Python to load checkpoint to CPU first, - then transfer to TPU. - Only works on systems with Colocated Python support (e.g., Pathways). - transfer_concurrent_gb: Maximum concurrent GB in flight during - CPU to TPU transfer when using colocated Python. Defaults to 16GB. Returns: List of deserialized JAX arrays. + + Environment variables: + COLOCATED_PYTHON_DESERIALIZE: Set to "0" or "false" to disable colocated Python + deserialization even when available. Defaults to enabled when available. + COLOCATED_PYTHON_PIPELINE_CONCURRENT_GB: Maximum concurrent GB in flight during + CPU to TPU transfer when using colocated Python. Defaults to 64. """ self.wait_until_finished() start_time = time.perf_counter() + use_colocated_python = colocated_python_available() and ( + os.getenv("COLOCATED_PYTHON_DESERIALIZE", "").lower() not in ("0", "false") + ) + logging.info("use_colocated_python=%s", use_colocated_python) + concurrent_bytes = concurrent_gb * 10**9 # Prepare global shapes and dtypes global_shapes_list = ( - [None] * len(tensorstore_specs) if global_shapes is None else list(global_shapes) + [None] * len(tensorstore_specs) + if global_shapes is None + else [tuple(s) for s in global_shapes] ) dtypes_list = [None] * len(tensorstore_specs) if dtypes is None else list(dtypes) @@ -995,17 +989,24 @@ def deserialize( # Capture mesh from main thread where context is active # (mesh is thread-local and won't be available in background event loop) tpu_mesh = thread_resources.env.physical_mesh + pipeline_concurrent_gb = int(os.getenv("COLOCATED_PYTHON_PIPELINE_CONCURRENT_GB", "64")) + + # Resolve any None shapes from TensorStore metadata before entering the colocated + # path, since out_specs_fn requires concrete shapes. + resolved_shapes = [ + tuple(ts.open(ts.Spec(spec), open=True).result().shape) if shape is None else shape + for spec, shape in zip(tensorstore_specs, global_shapes_list) + ] # Use colocated Python path coro = _run_colocated_deserialize( shardings=shardings, tensorstore_specs=tensorstore_specs, - global_shapes=global_shapes_list, + global_shapes=resolved_shapes, dtypes=dtypes_list, concurrent_bytes=concurrent_bytes, tpu_mesh=tpu_mesh, - transfer_concurrent_bytes=transfer_concurrent_gb * 10**9, - single_thread_pool=self._single_thread_pool, + pipeline_concurrent_bytes=pipeline_concurrent_gb * 10**9, multi_thread_pool=self._multi_thread_pool, ) else: From 01a11ed8dfd0ce04d21cfb8ba5592f940d958bff Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Tue, 3 Mar 2026 12:01:22 -0800 Subject: [PATCH 4/5] Try to fix and test colocated-python ckpt deserilization memory leak. --- .../examples/colocated_python_benchmark.py | 42 ++++++++++++------- axlearn/common/array_serialization.py | 7 +++- 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index 18a50bedb..28a052a52 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -277,6 +277,12 @@ def main(): action="store_true", help="Enable JAX profiler (adds overhead, disable for accurate benchmarking)", ) + parser.add_argument( + "--num_iters", + type=int, + default=1, + help="Number of times to repeat the load benchmark (default: 1)", + ) args = parser.parse_args() # Disable persistent compilation cache for fair benchmarking @@ -309,7 +315,8 @@ def main(): state_spec = create_state_spec_from_checkpoint(args.ckpt_path) print(f"Found {len(jax.tree_util.tree_leaves(state_spec))} tensors in checkpoint") - print(f"--- Running {args.method} benchmark ---") + num_iterations = args.num_iters + print(f"--- Running {args.method} benchmark ({num_iterations} iterations) ---") loaded_values = None try: with create_mesh(): @@ -318,21 +325,28 @@ def main(): args.ckpt_path, state_spec ) + if args.method == "default": + os.environ["COLOCATED_PYTHON_DESERIALIZE"] = "0" + + loaded_values = None with maybe_profile(args.profile, profile_dir): - if args.method == "default": - os.environ["COLOCATED_PYTHON_DESERIALIZE"] = "0" - start_time = time.perf_counter() - loaded_values = load_model( - tensorstore_specs=tensorstore_specs, - shardings=shardings, - global_shapes=global_shapes, - dtypes=dtypes, - ) - print(f"✅ Successfully loaded model from {args.ckpt_path}") - print(f"Total time took {time.perf_counter() - start_time:.2f} seconds") - print(f" Total parameters: {sum(x.size for x in loaded_values):,}") + for i in range(num_iterations): + if loaded_values is not None: + del loaded_values + print(f"\n--- Iteration {i + 1}/{num_iterations} ---") + start_time = time.perf_counter() + loaded_values = load_model( + tensorstore_specs=tensorstore_specs, + shardings=shardings, + global_shapes=global_shapes, + dtypes=dtypes, + ) + elapsed = time.perf_counter() - start_time + print(f"✅ Successfully loaded model from {args.ckpt_path}") + print(f"Total time took {elapsed:.2f} seconds") + print(f" Total parameters: {sum(x.size for x in loaded_values):,}") finally: - # Always clean up, even if benchmark fails + # Always clean up, even if benchmark fails. if loaded_values is not None: cleanup_loaded_arrays(loaded_values) diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index bd6da7f50..d541a0269 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -18,6 +18,7 @@ """ import asyncio import functools +import gc import math import os import threading @@ -786,12 +787,14 @@ async def _load_and_transfer_one( # block_until_ready is run in an executor so the event loop can continue # dispatching other coroutines while waiting for this array to materialize. await loop.run_in_executor(multi_thread_pool, cpu_array.block_until_ready) - return await loop.run_in_executor( + tpu_array = await loop.run_in_executor( multi_thread_pool, _blocking_device_put, cpu_array, tpu_sharding, ) + del cpu_array # Release CPU buffer promptly; don't wait for coroutine frame to exit. + return tpu_array async def _load_and_transfer_one_rate_limited( idx: int, @@ -836,6 +839,8 @@ async def _load_and_transfer_one_rate_limited( total_time = time.perf_counter() - start_time logging.info("Pipelined colocated deserialization completed in %.2f seconds", total_time) + gc.collect() + return tpu_arrays From 9e1965173bc88c3298924d368904e367dcdca73a Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Fri, 6 Mar 2026 10:50:18 -0800 Subject: [PATCH 5/5] snapshot of 03/05 night --- axlearn/common/array_serialization.py | 146 +++++++++++++++++++++++--- 1 file changed, 131 insertions(+), 15 deletions(-) diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index d541a0269..897e20c96 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -43,6 +43,88 @@ from axlearn.common.utils import Tensor +class _ColocatedResourceManager: + """Manages TensorStore context and event loop on colocated Python workers. + + Encapsulates all sidecar-side state for checkpoint deserialization. + Created by _colocated_deserialize_setup and stored as an attribute on the + colocated_python module so it can be accessed from other colocated functions + without cloudpickle frozen-globals issues. + + Must be created before each deserialization to ensure fresh resources, + avoiding accumulation of stale resources across checkpoint reloads. + """ + + def __init__(self): + """Initialize TensorStore context and event loop on colocated Python workers.""" + self.ts_context = ts.Context(serialization.TS_CONTEXT.spec) + + # Creates a dedicated event loop with a dedicated executor pool. + self.event_loop = asyncio.new_event_loop() + self.event_loop.set_default_executor(futures.ThreadPoolExecutor(max_workers=os.cpu_count())) + + # Runs the event loop in a background thread. + self.loop_thread = threading.Thread(target=self.event_loop.run_forever, daemon=True) + self.loop_thread.start() + logging.info("_ColocatedResourceManager initialized on sidecar.") + + def __del__(self): + """Release TensorStore context and event loop on colocated Python workers.""" + if hasattr(self, "event_loop") and self.event_loop is not None: + # Stop the dedicated event loop. + self.event_loop.call_soon_threadsafe(self.event_loop.stop) + # Exit the background thread and release resources of the event loop. + self.loop_thread.join() + self.event_loop.close() + logging.info("_ColocatedResourceManager destroyed on sidecar.") + + +@colocated_python.colocated_python +def _colocated_deserialize_setup(dummy_array: jax.Array) -> jax.Array: + """Initialize TensorStore context and event loop on colocated Python workers. + + Must be called before each deserialization to ensure fresh resources, + avoiding accumulation of stale resources across checkpoint reloads. + """ + colocated_python.resource_manager = _ColocatedResourceManager() + return dummy_array + + +@colocated_python.colocated_python +def _colocated_deserialize_teardown(dummy_array: jax.Array) -> jax.Array: + """Release TensorStore context and event loop on colocated Python workers. + + Must be called after each deserialization to release resources. + """ + import tracemalloc # pylint: disable=import-outside-toplevel + from collections import Counter # pylint: disable=import-outside-toplevel + + if hasattr(colocated_python, "resource_manager"): + colocated_python.resource_manager = None + + gc.collect() + + # Log object counts by type to identify what's accumulating. + type_counts = Counter(type(obj).__name__ for obj in gc.get_objects()) + logging.info("Top 20 object types: %s", type_counts.most_common(20)) + + # Log top memory allocations by source location. + if not tracemalloc.is_tracing(): + tracemalloc.start() + else: + snapshot = tracemalloc.take_snapshot() + top_stats = snapshot.statistics("lineno") + for stat in top_stats[:20]: + logging.info("tracemalloc: %s", stat) + + logging.info( + "Teardown complete. live objects: %d, gc garbage: %d", + len(gc.get_objects()), + len(gc.garbage), + ) + return dummy_array + + @dataclass class _ShardInfo: """Stores information for a maybe sliced jax.Shard. @@ -421,17 +503,24 @@ async def _run_serializer( raise e -def _blocking_device_put(tensor: Tensor, target: Union[Format, jax.sharding.Sharding]) -> Tensor: +def _blocking_device_put( + tensor: Tensor, target: Union[Format, jax.sharding.Sharding], *, verbose: bool = False +) -> Tensor: """Device put and block until ready. Args: tensor: Array to transfer. target: Either a Format (with layout + sharding) or Sharding. + verbose: If True, log execution time. Returns: Transferred array. """ - return jax.block_until_ready(jax.device_put(tensor, target)) + start = time.perf_counter() + result = jax.block_until_ready(jax.device_put(tensor, target)) + if verbose: + logging.info("device_put took %.3f seconds", time.perf_counter() - start) + return result async def _async_deserialize( @@ -444,6 +533,7 @@ async def _async_deserialize( byte_limiter: serialization._LimitInFlightBytes, single_thread_pool: Optional[ThreadPoolExecutor], multi_thread_pool: Optional[ThreadPoolExecutor], + ts_context: Optional[ts.Context] = None, ): """Modified from https://github.com/jax-ml/jax/blob/e7ec418eba9ada336f755613948cbdf4a9e97d59/jax/experimental/array_serialization/serialization.py#L345 @@ -500,7 +590,12 @@ async def _async_deserialize( # - On AWS (or other non-GCP environments) accessing GCS, gcs_grpc may hit auth/network # issues due to cross-cloud constraints. So we enable this optimization on Pathways # which only runs on GCP for now. - context = serialization.TS_CONTEXT + if ts_context is not None: + logging.info("Using provided TensorStore context for deserialization.") + context = ts_context + else: + logging.info("Using global TensorStore context for deserialization.") + context = serialization.TS_CONTEXT if os.getenv("ENABLE_GCS_GRPC", "false") == "true": logging.debug("gcs_grpc enabled") tensorstore_spec, context = use_gcs_grpc(tensorstore_spec) @@ -565,7 +660,9 @@ async def cb(index: array.Index, device: jax.Device): if h2d_limiter is None: # No per-shard H2D gating — caller (e.g. colocated Python path) relies on a global # pipeline limiter to bound concurrency instead. - result = await loop.run_in_executor(None, _blocking_device_put, out, layout) + result = await loop.run_in_executor( + None, functools.partial(_blocking_device_put, verbose=True), out, layout + ) # Jax >= 0.6.2 changes the behavior of _LimitInFlightBytes, where wait_for_bytes no longer # throws an exception if requested_bytes > max_bytes # pylint: disable-next=protected-access @@ -651,7 +748,7 @@ async def _effective_bytes_per_device( TensorStore reads whole chunks even if only a partial chunk is needed, so actual bytes read can exceed the raw tensor shard size. """ - t = await ts.open(ts.Spec(spec), open=True) + t = await ts.open(ts.Spec(spec), open=True, context=serialization.TS_CONTEXT) shard_shape = sharding.shard_shape(shape) raw_bytes = math.prod(shard_shape) * np.dtype(dtype).itemsize @@ -740,13 +837,21 @@ async def _run_colocated_deserialize( # _get_specialized_func is cached on (FunctionInfo, Specialization): with a single function # object the cache hits for all arrays that share the same output shape/dtype/sharding, # reducing compilations from O(num_arrays) to O(num_unique_output_specs). + cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) + replicated_sharding = jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) + + # Initialize fresh TensorStore context and event loop on the sidecar workers. + dummy_array = jax.device_put(jnp.array(0), replicated_sharding) + _colocated_deserialize_setup.specialize(devices=cpu_devices)(dummy_array) + @colocated_python.colocated_python def load_to_cpu(idx: jax.Array): i = int(idx) # pylint: disable=protected-access byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - return asyncio.run( + rm = colocated_python.resource_manager + return asyncio.run_coroutine_threadsafe( _async_deserialize( cpu_shardings[i], tensorstore_specs[i], @@ -756,14 +861,14 @@ def load_to_cpu(idx: jax.Array): h2d_limiter=None, single_thread_pool=None, multi_thread_pool=None, - ) - ) + ts_context=rm.ts_context, + ), + rm.event_loop, + ).result() # Pre-create replicated index arrays across all cpu_devices. colocated_python requires inputs # to have one shard per device in the specialized device list. - cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) - idx_sharding = jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) - idx_arrays = [jax.device_put(jnp.array(i), idx_sharding) for i in range(len(shardings))] + idx_arrays = [jax.device_put(jnp.array(i), replicated_sharding) for i in range(len(shardings))] async def _load_and_transfer_one( idx: int, @@ -793,7 +898,9 @@ async def _load_and_transfer_one( cpu_array, tpu_sharding, ) - del cpu_array # Release CPU buffer promptly; don't wait for coroutine frame to exit. + # Explicitly release the CPU buffer once H2D transfer is complete. + # This is for deterministic memory release within colocated-python sidecar container. + del cpu_array return tpu_array async def _load_and_transfer_one_rate_limited( @@ -839,7 +946,8 @@ async def _load_and_transfer_one_rate_limited( total_time = time.perf_counter() - start_time logging.info("Pipelined colocated deserialization completed in %.2f seconds", total_time) - gc.collect() + # Tear down the TensorStore context and event loop on colocated workers. + _colocated_deserialize_teardown.specialize(devices=cpu_devices)(dummy_array) return tpu_arrays @@ -999,7 +1107,15 @@ def deserialize( # Resolve any None shapes from TensorStore metadata before entering the colocated # path, since out_specs_fn requires concrete shapes. resolved_shapes = [ - tuple(ts.open(ts.Spec(spec), open=True).result().shape) if shape is None else shape + ( + tuple( + ts.open(ts.Spec(spec), open=True, context=serialization.TS_CONTEXT) + .result() + .shape + ) + if shape is None + else shape + ) for spec, shape in zip(tensorstore_specs, global_shapes_list) ] @@ -1178,4 +1294,4 @@ def serialize( self._tensorstore_spec_log_fn(tensorstore_specs) logging.info("D2H during save took %fs. Starting async commit.", time.time() - start_t) - self._start_async_commit(on_commit_callback) + self._start_async_commit(on_commit_callback) \ No newline at end of file