add util function get_save_and_offload_names

Author: aireenmei

src/maxtext/layers/decoders.py

@@ -414,9 +414,9 @@ def get_remat_policy(self):
       elif cfg.remat_policy in ("qkv_proj_offloaded", "minimal_offloaded", "custom"):
         # minimal_offloaded offloads all except context. All three share a single
         # source of truth for their save/offload name lists (see
-        # maxtext_utils.get_offload_remat_names) so that offloading configured via
+        # maxtext_utils.get_save_and_offload_names) so that offloading configured via
         # `custom` resolves identically to the named presets.
-        save_names, offload_names = maxtext_utils.get_offload_remat_names(cfg)
+        save_names, offload_names = maxtext_utils.get_save_and_offload_names(cfg)
         policy = jax.checkpoint_policies.save_and_offload_only_these_names(
             names_which_can_be_saved=save_names,
             names_which_can_be_offloaded=offload_names,

src/maxtext/utils/maxtext_utils.py

@@ -195,18 +195,26 @@ def should_prevent_cse_in_remat(config):
   return True
 
 
-def get_offload_remat_names(config):
-  """Returns ``(save_names, offload_names)`` for offloading remat policies, else ``None``.
-
-  Single source of truth for which checkpointed tensors a remat policy saves on
-  device vs. offloads to pinned host. Shared by ``Decoder.get_remat_policy`` (which
-  builds the save-and-offload policy) and by models that must convert an offload
-  into a device-save at a scan boundary that cannot carry pinned-host residuals
-  (e.g. Gemma4's global-layer trip-count-one scan). Keeping both callers on this
-  helper ensures ``remat_policy=qkv_proj_offloaded`` and ``remat_policy=custom``
-  with the same tensors marked ``offload`` resolve to identical name sets.
-
-  Returns ``None`` for non-offloading policies.
+def get_save_and_offload_names(config) -> tuple[list[str], list[str]]:
+  """Returns the ``(save_names, offload_names)`` split for remat policies built via
+  ``jax.checkpoint_policies.save_and_offload_only_these_names``.
+
+  ``save_names`` are checkpointed tensors kept in device HBM; ``offload_names`` are moved to
+  pinned host. This is the single source of truth shared by ``Decoder.get_remat_policy`` (which
+  builds the save-and-offload policy) and by models that use custom ways to handle offload (
+  e.g. Gemma4's global-layer with scan). It also makes ``remat_policy=custom`` with tensors marked
+  ``offload`` resolve to the same name sets as the named presets.
+
+  Returns a ``(save_names, offload_names)`` tuple:
+    * ``custom``: ``(config.tensors_on_device, config.tensors_to_offload)`` -- the per-tensor
+      assignments. Either list may be empty: all tensors set to ``device`` gives an empty offload
+      list, all set to ``remat`` gives ``([], [])``.
+    * ``qkv_proj_offloaded`` / ``minimal_offloaded``: ``([], <hardcoded offload names>)`` -- presets
+      that only offload and save nothing on device.
+    * any other policy (``full``, ``minimal``, ``save_*``, ``none``, ...): ``([], [])`` -- these do
+      not use ``save_and_offload_only_these_names``, so they contribute no names to this split.
+      Note ``([], [])`` here means "no names for this split", not that the policy saves nothing
+      overall (e.g. ``save_out_proj`` still saves ``out_proj`` via ``save_only_these_names``).
   """
   if config.remat_policy == "qkv_proj_offloaded":
     return [], ["query_proj", "value_proj", "key_proj", "kv_proj"]
@@ -225,7 +233,7 @@ def get_offload_remat_names(config):
     ]
   if config.remat_policy == "custom":
     return list(config.tensors_on_device or []), list(config.tensors_to_offload or [])
-  return None
+  return [], []
 
 
 def load_compiled(config, partial_train, state, execution_devices):

tests/unit/maxtext_utils_test.py

@@ -1731,8 +1731,9 @@ def test_update_kv_caches_after_scan_invalid_type(self):
       maxtext_utils.update_kv_caches_after_scan(kv_caches_tuple, returned_kv_cache, scan_length=1, block_len=2)
 
 
-class TestGetOffloadRematNames(unittest.TestCase):
-  """Tests for maxtext_utils.get_offload_remat_names."""
+@pytest.mark.cpu_only
+class TestGetSaveAndOffloadNames(unittest.TestCase):
+  """Tests for maxtext_utils.get_save_and_offload_names (pure config logic, no device needed)."""
 
   @staticmethod
   def _cfg(remat_policy, tensors_on_device=None, tensors_to_offload=None):
@@ -1743,33 +1744,42 @@ def _cfg(remat_policy, tensors_on_device=None, tensors_to_offload=None):
     )
 
   def test_named_preset_matches_equivalent_custom(self):
-    """qkv_proj_offloaded resolves identically to custom with the same offloads."""
-    preset = maxtext_utils.get_offload_remat_names(self._cfg("qkv_proj_offloaded"))
-    custom = maxtext_utils.get_offload_remat_names(
-        self._cfg("custom", tensors_on_device=[], tensors_to_offload=["query_proj", "value_proj", "key_proj", "kv_proj"])
+    """qkv_proj_offloaded's offload names resolve identically to an equivalent custom config.
+
+    A real custom config keeps decoder_layer_input on device by default, so its full tuple
+    differs from the preset by that (benign, boundary) save entry -- assert only the offload halves.
+    """
+    _, preset_offload = maxtext_utils.get_save_and_offload_names(self._cfg("qkv_proj_offloaded"))
+    _, custom_offload = maxtext_utils.get_save_and_offload_names(
+        self._cfg(
+            "custom",
+            tensors_on_device=["decoder_layer_input"],
+            tensors_to_offload=["query_proj", "value_proj", "key_proj", "kv_proj"],
+        )
     )
-    self.assertEqual(preset, custom)
+    self.assertEqual(custom_offload, preset_offload)
 
   def test_kv_proj_retained_in_offload_presets(self):
     """Regression guard: kv_proj must stay in the offload presets (it is a real checkpoint name)."""
-    _, qkv_offload = maxtext_utils.get_offload_remat_names(self._cfg("qkv_proj_offloaded"))
-    _, minimal_offload = maxtext_utils.get_offload_remat_names(self._cfg("minimal_offloaded"))
+    _, qkv_offload = maxtext_utils.get_save_and_offload_names(self._cfg("qkv_proj_offloaded"))
+    _, minimal_offload = maxtext_utils.get_save_and_offload_names(self._cfg("minimal_offloaded"))
     self.assertIn("kv_proj", qkv_offload)
     self.assertIn("kv_proj", minimal_offload)
 
   def test_custom_reads_config_lists(self):
-    save, offload = maxtext_utils.get_offload_remat_names(
+    save, offload = maxtext_utils.get_save_and_offload_names(
         self._cfg("custom", tensors_on_device=["context"], tensors_to_offload=["out_proj"])
     )
     self.assertEqual(save, ["context"])
     self.assertEqual(offload, ["out_proj"])
 
   def test_custom_handles_none_lists(self):
-    self.assertEqual(maxtext_utils.get_offload_remat_names(self._cfg("custom")), ([], []))
+    self.assertEqual(maxtext_utils.get_save_and_offload_names(self._cfg("custom")), ([], []))
 
-  def test_non_offloading_policies_return_none(self):
+  def test_non_offloading_policies_return_empty(self):
+    """Policies that don't use the save/offload split contribute no names to it."""
     for policy in ("full", "minimal", "save_out_proj", "save_qkv_proj", "none"):
-      self.assertIsNone(maxtext_utils.get_offload_remat_names(self._cfg(policy)))
+      self.assertEqual(maxtext_utils.get_save_and_offload_names(self._cfg(policy)), ([], []))
 
 
 if __name__ == "__main__":