Add minimal remat policy for flash attention

AI-Hypercomputer · Apr 3, 2024 · 012e0a4 · 012e0a4
1 parent 5575702
commit 012e0a4
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'