Merge pull request #573 from google:yooh/minimal_flash

PiperOrigin-RevId: 621911981
AI-Hypercomputer · Apr 4, 2024 · 264c823 · 264c823
2 parents de49d83 + 012e0a4
commit 264c823
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/MaxText/layers/models.py b/MaxText/layers/models.py
@@ -234,6 +234,11 @@ def __call__(self,
           offload_src="device", offload_dst="pinned_host")
       elif cfg.remat_policy == 'minimal_offloaded':
         policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(offload_src="device", offload_dst="pinned_host")
+      elif cfg.remat_policy == 'minimal_flash':
+        policy = jax.checkpoint_policies.save_from_both_policies(
+          jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims,
+          jax.checkpoint_policies.save_only_these_names('context',),
+        )
       else:
         assert (
             cfg.remat_policy == 'full'