Merge branch 'jax_0.8.0_py3.12_v7x' into jax_py3.12_nightly

Author: andersensam

axlearn/experiments/text/gpt/envy.py

@@ -35,9 +35,10 @@
     ScaleKey,
     ScaleQuery,
     TransformerLayer,
+    RematRegexSavePatterns,
 )
 from axlearn.common.base_layer import RematSpec
-from axlearn.common.config import TrainerConfigFn
+from axlearn.common.config import TrainerConfigFn, config_for_function
 from axlearn.common.embedding import TransformerTextEmbeddings
 from axlearn.common.layers import RMSNorm
 from axlearn.common.mixture_of_experts import TransformerFeedForwardMoE, get_outer_batch_from_mesh
@@ -65,6 +66,16 @@
     scaled_hidden_dim,
 )
 from axlearn.experiments.text.gpt.fuji import offload_attention_proj_policy
+from axlearn.experiments.trainer_config_utils import (
+    SplashAttentionConfigModifier,
+    V6eFlashConfigModifier,
+    V7xFlashConfigModifier,
+)
+from axlearn.common.utils import (
+    combine_remat_policies,
+    extended_checkpoint_policies,
+    save_and_offload_only_these_names_regex,
+)
 
 MODEL_SIZES = ("test", "Switch-Base", "Switch-Large", "Switch-XXL")
 
@@ -215,6 +226,37 @@ def get_trainer_kwargs(
                         ],
                     ),
                 ),
+                (
+                    "tpu-v7x-.*",
+                    ChainConfigModifier.default_config().set(
+                        config_modifiers=[
+                            MeshShapeModifier.default_config().set(
+                                mesh_shape=mesh_shape_from_axes(data=-1, expert=16, fsdp=16)
+                            ),
+                            # Use the updated block size for Splash Attention
+                            V7xFlashConfigModifier.default_config(),
+                            RematSpecModifier.default_config().set(
+                                remat_policies={
+                                    "model.decoder.transformer.layer": RematSpec(
+                                        prevent_cse=False,
+                                        policy=config_for_function(
+                                            save_and_offload_only_these_names_regex
+                                        ).set(
+                                            names_which_can_be_saved=None,
+                                            names_which_can_be_offloaded="|".join(
+                                                [
+                                                    RematRegexSavePatterns.INPUT.value,
+                                                ]
+                                            ),
+                                            offload_src="device",
+                                            offload_dst="pinned_host",
+                                        ),
+                                    ),
+                                }
+                            ),
+                        ],
+                    ),
+                )
             ),
         )
     elif model_size == "Switch-Large":