Update remat rules

Author: RissyRan

MaxText/configs/base.yml

@@ -240,7 +240,7 @@ set_remat_policy_on_pipeline_iterations: True
 set_remat_policy_on_layers_per_stage: False
 
 
-# Choose 'remat_policy' between 'minimal', 'save_dot_with_context_except_mlp', 'save_dot_except_mlpwi', 'save_dot_except_mlp',
+# Choose 'remat_policy' between 'minimal_with_context', 'minimal', 'save_dot_with_context_except_mlp', 'save_dot_except_mlpwi', 'save_dot_except_mlp',
 # 'save_qkv_proj', 'qkv_proj_offloaded', 'custom', 'minimal_offloaded', 'save_out_proj' and 'full'.
 # These options offer a trade-off between speed (fastest to slowest) and HBM usage (highest to lowest)
 remat_policy: 'full'

MaxText/configs/models/gpu/llama2_7b.yml

@@ -8,7 +8,7 @@ max_target_length: 4096
 model_name: "llama2-7b"
 enable_checkpointing: False
 attention: "cudnn_flash_te"
-remat_policy: "minimal_flash"
+remat_policy: "minimal_with_context"
 use_iota_embed: True
 scan_layers: False
 dataset_type: "synthetic"

MaxText/configs/models/gpu/llama3_8b.yml

@@ -23,7 +23,7 @@ steps: 30
 per_device_batch_size: 12
 max_target_length: 8192
 attention: "cudnn_flash_te"
-remat_policy: "minimal_flash"
+remat_policy: "minimal_with_context"
 use_iota_embed: True
 dataset_type: "synthetic"
 reuse_example_batch: 1

MaxText/configs/models/gpu/mixtral_8x7b.yml

@@ -23,7 +23,7 @@ steps: 30
 per_device_batch_size: 12
 max_target_length: 4096
 attention: "cudnn_flash_te"
-remat_policy: "minimal_flash"
+remat_policy: "minimal_with_context"
 use_iota_embed: True
 dataset_type: "synthetic"
 reuse_example_batch: 1

MaxText/layers/decoders.py

@@ -264,8 +264,29 @@ def get_remat_policy(self):
     policy = None
     cfg = self.config
     if cfg.remat_policy != "none":
-      if cfg.remat_policy == "minimal":
-        policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
+      if cfg.remat_policy == "minimal_with_context":
+        policy = jax.checkpoint_policies.save_only_these_names(
+            "query_proj",
+            "value_proj",
+            "key_proj",
+            "qkv_proj",
+            "context",
+            "out_proj",
+            "mlpwi_0",
+            "mlpwi_1",
+            "mlpwo",
+        )
+      elif cfg.remat_policy == "minimal":
+        policy = jax.checkpoint_policies.save_only_these_names(
+            "query_proj",
+            "value_proj",
+            "key_proj",
+            "qkv_proj",
+            "out_proj",
+            "mlpwi_0",
+            "mlpwi_1",
+            "mlpwo",
+        )
       elif cfg.remat_policy == "save_dot_with_context_except_mlp":
         policy = jax.checkpoint_policies.save_only_these_names(
             "query_proj",
@@ -307,21 +328,28 @@ def get_remat_policy(self):
             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")
+        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",
+                "qkv_proj",
+                "out_proj",
+                "mlpwi_0",
+                "mlpwi_1",
+                "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,
             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",
-            ),
-        )
       elif cfg.remat_policy == "save_out_proj":
         policy = jax.checkpoint_policies.save_only_these_names(
             "out_proj",
@@ -422,9 +450,7 @@ def get_norm_layer(self, num_features: int):
     else:
       raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}")
 
-  def scan_decoder_layers(
-      self, cfg, decoder_layer, length, metadata_axis_name, mesh, in_axes_tuple, model_mode, **kwargs
-  ):
+  def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, mesh, in_axes_tuple, model_mode, **kwargs):
     """scan decoder layers, calls `flax.linen.transforms.scan`"""
     initializing = self.is_mutable_collection("params")
     params_spec = cfg.param_scan_axis if initializing else ScanIn(cfg.param_scan_axis)
@@ -744,9 +770,7 @@ def __call__(
           # Iterate over the two layer groups (dense and MoE) and apply layer transformation
           for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes):
             for index in range(num_layers):
-              y = layer(
-                  config=cfg, mesh=mesh, name=f"{layer_prefix}_{index}", quant=self.quant, model_mode=self.model_mode
-              )(
+              y = layer(config=cfg, mesh=mesh, name=f"{layer_prefix}_{index}", quant=self.quant, model_mode=self.model_mode)(
                   y,
                   decoder_segment_ids,
                   decoder_positions,