add util function get_offload_remat_names

Author: aireenmei

src/maxtext/layers/decoders.py

@@ -411,36 +411,15 @@ def get_remat_policy(self):
             "kv_proj",
             "qkv_proj",
         )
-      elif cfg.remat_policy == "qkv_proj_offloaded":
+      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
+        # `custom` resolves identically to the named presets.
+        save_names, offload_names = maxtext_utils.get_offload_remat_names(cfg)
         policy = jax.checkpoint_policies.save_and_offload_only_these_names(
-            names_which_can_be_saved=[],
-            names_which_can_be_offloaded=["query_proj", "value_proj", "key_proj", "kv_proj"],
-            offload_src="device",
-            offload_dst="pinned_host",
-        )
-      elif cfg.remat_policy == "minimal_offloaded":
-        # offload all except context
-        policy = jax.checkpoint_policies.save_and_offload_only_these_names(
-            names_which_can_be_saved=[],
-            names_which_can_be_offloaded=[
-                "query_proj",
-                "value_proj",
-                "key_proj",
-                "kv_proj",
-                "qkv_proj",
-                "out_proj",
-                "mlpwi_0",
-                "mlpwi_1",
-                "mlpwi",
-                "mlpwo",
-            ],
-            offload_src="device",
-            offload_dst="pinned_host",
-        )
-      elif cfg.remat_policy == "custom":
-        policy = jax.checkpoint_policies.save_and_offload_only_these_names(
-            names_which_can_be_saved=cfg.tensors_on_device,
-            names_which_can_be_offloaded=cfg.tensors_to_offload,
+            names_which_can_be_saved=save_names,
+            names_which_can_be_offloaded=offload_names,
             offload_src="device",
             offload_dst="pinned_host",
         )

src/maxtext/utils/maxtext_utils.py

@@ -195,6 +195,39 @@ 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.
+  """
+  if config.remat_policy == "qkv_proj_offloaded":
+    return [], ["query_proj", "value_proj", "key_proj", "kv_proj"]
+  if config.remat_policy == "minimal_offloaded":
+    return [], [
+        "query_proj",
+        "value_proj",
+        "key_proj",
+        "kv_proj",
+        "qkv_proj",
+        "out_proj",
+        "mlpwi_0",
+        "mlpwi_1",
+        "mlpwi",
+        "mlpwo",
+    ]
+  if config.remat_policy == "custom":
+    return list(config.tensors_on_device or []), list(config.tensors_to_offload or [])
+  return None
+
+
 def load_compiled(config, partial_train, state, execution_devices):
   """# Loading a serialized compiled train step function."""
 

tests/unit/maxtext_utils_test.py

@@ -17,6 +17,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass, field
 import functools
+from types import SimpleNamespace
 from typing import Any, Sequence
 import unittest
 from unittest.mock import MagicMock, Mock, patch
@@ -1730,5 +1731,46 @@ 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."""
+
+  @staticmethod
+  def _cfg(remat_policy, tensors_on_device=None, tensors_to_offload=None):
+    return SimpleNamespace(
+        remat_policy=remat_policy,
+        tensors_on_device=tensors_on_device,
+        tensors_to_offload=tensors_to_offload,
+    )
+
+  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"])
+    )
+    self.assertEqual(preset, custom)
+
+  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"))
+    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(
+        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")), ([], []))
+
+  def test_non_offloading_policies_return_none(self):
+    for policy in ("full", "minimal", "save_out_proj", "save_qkv_proj", "none"):
+      self.assertIsNone(maxtext_utils.get_offload_remat_names(self._cfg(policy)))
+
+
 if __name__ == "__main__":
   unittest.main()