Enable Memory Regulator for Orbax checkpoint saving.

PiperOrigin-RevId: 897298923

Author: JustinPan-goog

checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py

@@ -22,6 +22,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import dataclasses
 import functools
 import json
@@ -50,6 +51,7 @@
 from orbax.checkpoint._src.path import format_utils
 from orbax.checkpoint._src.path import types as path_types
 from orbax.checkpoint._src.serialization import limits
+from orbax.checkpoint._src.serialization import memory_regulator
 from orbax.checkpoint._src.serialization import ocdbt_utils
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 from orbax.checkpoint._src.serialization import type_handler_registry as type_handler_registry_lib
@@ -323,6 +325,17 @@ def _maybe_set_default_save_restore_args(v, leaf_args):
   )
 
 
+@contextlib.contextmanager
+def _memory_profiler_context():
+  """Context manager for memory_regulator profiler."""
+  memory_regulator.profiler_start()
+  try:
+    yield
+  finally:
+    # Explicitly stop the bg thread if an exception occurs
+    memory_regulator.profiler_end()
+
+
 
 
 def _format_bytes(bytes_value: Optional[int]) -> str:
@@ -347,7 +360,9 @@ def __init__(
       *,
       save_concurrent_bytes: Optional[int] = None,
       restore_concurrent_bytes: Optional[int] = None,
-      save_device_host_concurrent_bytes: Optional[int] = None,
+      save_device_host_concurrent_bytes: int | str | None = None,
+      max_save_device_host_concurrent_bytes: int | None = None,
+      fallback_host_limit_gb: int | None = None,
       use_ocdbt: bool = True,
       use_zarr3: bool = False,
       use_compression: bool = True,
@@ -377,7 +392,12 @@ def __init__(
       save_device_host_concurrent_bytes: max concurrent bytes allowed to be
         transferred from device to host memory at once when saving. When the
         limit is reached, arrays must be finished writing to the checkpoint
-        before a new array can start being transferred.
+        before a new array can start being transferred. Can be "auto".
+      max_save_device_host_concurrent_bytes: The maximum memory limit in bytes
+        allowed for regulation. Required if `save_device_host_concurrent_bytes`
+        is "auto".
+      fallback_host_limit_gb: Fallback physical machine size in GB to use if the
+        profiler fails to fetch the total memory dynamically.
       use_ocdbt: Whether to use OCDBT format for saving.
       use_zarr3: If True, use Zarr ver3 otherwise Zarr ver2.
       use_compression: If True, use zstd compression.
@@ -408,6 +428,30 @@ def __init__(
     self._save_concurrent_bytes = save_concurrent_bytes
     self._restore_concurrent_bytes = restore_concurrent_bytes
     self._save_device_host_concurrent_bytes = save_device_host_concurrent_bytes
+    self._max_save_device_host_concurrent_bytes = (
+        max_save_device_host_concurrent_bytes
+    )
+    self._fallback_host_limit_gib = None
+    if fallback_host_limit_gb is not None:
+      self._fallback_host_limit_gib = (
+          fallback_host_limit_gb * 10**9
+      ) / (1024**3)
+    if self._save_device_host_concurrent_bytes == 'auto':
+      if self._max_save_device_host_concurrent_bytes is None:
+        raise ValueError(
+            'max_save_device_host_concurrent_bytes must be provided if'
+            ' save_device_host_concurrent_bytes is "auto"'
+        )
+      max_memory_limit_gib = self._max_save_device_host_concurrent_bytes / (
+          1024**3
+      )
+      self._memory_regulator = memory_regulator.MemoryRegulator(
+          max_memory_limit_gib=max_memory_limit_gib,
+          fallback_host_limit_gib=self._fallback_host_limit_gib,
+      )
+      self._current_device_host_limit_bytes = int(
+          self._memory_regulator.min_memory_limit_gib * 1024**3
+      )
     self._use_ocdbt = use_ocdbt
     self._use_zarr3 = use_zarr3
     self._use_compression = use_compression
@@ -656,9 +700,38 @@ async def async_save(
 
     save_args = _fill_missing_save_or_restore_args(item, save_args, mode='save')
     byte_limiter = limits.get_byte_limiter(self._save_concurrent_bytes)
-    device_host_byte_limiter = limits.get_byte_limiter(
-        self._save_device_host_concurrent_bytes
-    )
+
+    device_host_concurrent_bytes = self._save_device_host_concurrent_bytes
+    if device_host_concurrent_bytes == 'auto':
+      peak_usage_gib = memory_regulator.profiler_peak_usage_gib()
+      blocking_time_sec = memory_regulator.get_prev_blocking_time_sec()
+      expected_surge_gib = memory_regulator.get_expected_surge_gib()
+
+      total_memory_gib = memory_regulator.get_total_memory_gib()
+      current_limit_gib = self._current_device_host_limit_bytes / (1024**3)
+      next_limit_gib = self._memory_regulator.get_next_memory_limit(
+          current_limit_gib=current_limit_gib,
+          peak_memory_usage_gib=peak_usage_gib,
+          blocking_time_sec=blocking_time_sec,
+          expected_surge_gib=expected_surge_gib,
+          total_memory_gib=total_memory_gib,
+      )
+      self._current_device_host_limit_bytes = int(next_limit_gib * 1024**3)
+      logging.info(
+          'MemoryRegulated: Updated device_host_concurrent_bytes to %s'
+          ' (peak=%f GiB)',
+          humanize.naturalsize(
+              self._current_device_host_limit_bytes, binary=True
+          ),
+          peak_usage_gib,
+      )
+      device_host_byte_limiter = limits.get_byte_limiter(
+          self._current_device_host_limit_bytes
+      )
+    else:
+      device_host_byte_limiter = limits.get_byte_limiter(
+          device_host_concurrent_bytes
+      )
     param_infos = self._get_param_infos(
         item,
         directory,
@@ -698,27 +771,34 @@ async def async_save(
         directory / PYTREE_METADATA_FILE
     )
     batch_requests_ready_time = time.time()
-    if partial_save:
-      serialize_ops, tree_memory_size, param_infos, save_args = (
-          await self._async_partial_save(
-              directory, item, batch_requests, param_infos, save_args
-          )
-      )
-    else:
-      tree_memory_size = 0
-      for request in batch_requests:
-        serialize_ops += [
-            _logging_serialize(
-                request.handler,
-                request.handler.serialize(
-                    request.values, request.infos, request.args
-                ),
+    with _memory_profiler_context():
+      if partial_save:
+        serialize_ops, tree_memory_size, param_infos, save_args = (
+            await self._async_partial_save(
+                directory, item, batch_requests, param_infos, save_args
             )
-        ]
-        write_size, _ = _get_batch_memory_size(request.handler, request.values)
-        tree_memory_size += write_size
-    # Await copy futures. Returns List[List[future.Future]].
-    commit_futures = await asyncio.gather(*serialize_ops)
+        )
+      else:
+        tree_memory_size = 0
+        for request in batch_requests:
+          serialize_ops += [
+              _logging_serialize(
+                  request.handler,
+                  request.handler.serialize(
+                      request.values, request.infos, request.args
+                  ),
+              )
+          ]
+          write_size, _ = _get_batch_memory_size(
+              request.handler, request.values
+          )
+          tree_memory_size += write_size
+      # Await copy futures. Returns List[List[future.Future]].
+      commit_futures = await asyncio.gather(*serialize_ops)
+    logging.info(
+        'MemoryRegulated: Peak usage: %f GiB',
+        memory_regulator.profiler_peak_usage_gib(),
+    )
     # Flatten to List[future.Future].
     commit_futures, _ = jax.tree.flatten(commit_futures)
 

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py

@@ -460,8 +460,10 @@ def _get_impl_save_args(
 
 
 def _concurrent_bytes(
-    concurrent_gb: int | None, *, use_default_if_none: bool = True
-) -> int | None:
+    concurrent_gb: int | str | None, *, use_default_if_none: bool = True
+) -> int | str | None:
+  if concurrent_gb == 'auto':
+    return 'auto'
   if concurrent_gb is None:
     if use_default_if_none:
       return DEFAULT_CONCURRENT_GB * 10**9
@@ -500,7 +502,9 @@ def __init__(
       *,
       save_concurrent_gb: Optional[int] = None,
       restore_concurrent_gb: Optional[int] = None,
-      save_device_host_concurrent_gb: Optional[int] = None,
+      save_device_host_concurrent_gb: int | str | None = None,
+      max_save_device_host_concurrent_gb: int | None = None,
+      fallback_host_limit_gb: int | None = None,
       use_ocdbt: bool = True,
       use_zarr3: bool = False,
       use_compression: bool = True,
@@ -544,6 +548,12 @@ def __init__(
         transferred from device to host. Note that asynchronous saves may not be
         truly asynchronous with this option enabled, as we have to block on some
         array writes before beginning others. Also see `is_prioritized_key_fn`.
+        Can be set to "auto" to enable Memory Regulator.
+      max_save_device_host_concurrent_gb: The maximum memory limit in GB allowed
+        for regulation. Required if `save_device_host_concurrent_gb` is set to
+        "auto".
+      fallback_host_limit_gb: Fallback physical machine size in GB to use if
+        the profiler fails to fetch the total memory dynamically.
       use_ocdbt: enables Tensorstore OCDBT driver. This option allows using a
         different checkpoint format which is faster to read and write, as well
         as more space efficient.
@@ -586,9 +596,22 @@ def __init__(
     self._type_handler_registry = type_handler_registry
     self._save_concurrent_bytes = _concurrent_bytes(save_concurrent_gb)
     self._restore_concurrent_bytes = _concurrent_bytes(restore_concurrent_gb)
+    if (
+        save_device_host_concurrent_gb == 'auto'
+        and max_save_device_host_concurrent_gb is None
+    ):
+      raise ValueError(
+          'max_save_device_host_concurrent_gb must be provided if'
+          ' save_device_host_concurrent_gb is "auto"'
+      )
     self._save_device_host_concurrent_bytes = _concurrent_bytes(
         save_device_host_concurrent_gb, use_default_if_none=False
     )
+    max_save_device_host_concurrent_bytes = (
+        None
+        if max_save_device_host_concurrent_gb is None
+        else int(max_save_device_host_concurrent_gb * 10**9)
+    )
     logging.info(
         'save_device_host_concurrent_bytes=%s',
         self._save_device_host_concurrent_bytes,
@@ -597,6 +620,8 @@ def __init__(
         save_concurrent_bytes=self._save_concurrent_bytes,
         restore_concurrent_bytes=self._restore_concurrent_bytes,
         save_device_host_concurrent_bytes=self._save_device_host_concurrent_bytes,
+        max_save_device_host_concurrent_bytes=max_save_device_host_concurrent_bytes,
+        fallback_host_limit_gb=fallback_host_limit_gb,
         use_ocdbt=use_ocdbt,
         use_zarr3=use_zarr3,
         use_compression=use_compression,

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py

@@ -1943,7 +1943,8 @@ def test_concurrent_gb_save(self, limit_bytes):
     with mock.patch.object(
         limits,
         'get_byte_limiter',
-        new=lambda _: byte_limiter,
+        autospec=True,
+        return_value=byte_limiter,
     ):
       handler.save(self.directory, args=PyTreeSaveArgs(tree))
     # Replicated shards are handled within the _write_array_shard function.
@@ -1959,6 +1960,25 @@ def test_concurrent_gb_save(self, limit_bytes):
           sleep_time,
       )
 
+  def test_concurrent_gb_auto_save(self):
+    handler = PyTreeCheckpointHandler(
+        save_device_host_concurrent_gb='auto',
+        max_save_device_host_concurrent_gb=80,
+        use_ocdbt=False,
+    )
+    tree = {'a': np.ones(10)}
+    # Verify it doesn't crash when saving with "auto".
+    with mock.patch.object(
+        base_pytree_checkpoint_handler.memory_regulator,
+        'get_total_memory_gib',
+        autospec=True,
+        return_value=250.0,
+    ):
+      handler.save(self.directory, args=PyTreeSaveArgs(tree))
+      # Verify we can also restore
+      restored = handler.restore(self.directory)
+    np.testing.assert_array_equal(restored['a'], tree['a'])
+
   @parameterized.parameters((5,), (9,))
   def test_concurrent_gb_restore(self, limit_bytes):
     # TODO(b/346811105): Enable for Pathways.
@@ -1975,7 +1995,8 @@ def test_concurrent_gb_restore(self, limit_bytes):
     with mock.patch.object(
         limits,
         'get_byte_limiter',
-        new=lambda _,: byte_limiter,
+        autospec=True,
+        return_value=byte_limiter,
     ):
       restored = handler.restore(
           self.directory,