Add 405B full config

Author: andersensam

axlearn/experiments/text/gpt/fuji.py

@@ -68,7 +68,7 @@
 from axlearn.experiments.trainer_config_utils import (
     SplashAttentionConfigModifier,
     V6eFlashConfigModifier,
-    V7xFlashConfigModifier
+    V7xFlashConfigModifier,
 )
 
 MODEL_SIZES = ("test", "1B", "3B", "7B", "8B", "70B", "405B")
@@ -731,7 +731,7 @@ def get_trainer_kwargs(
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(fsdp=-1)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=-1)
                             ),
                             # Ensure we set the default tpu_block_size=2048 on v7x
                             # With the 70B model, MaxText also modifies block_q
@@ -919,7 +919,40 @@ def get_trainer_kwargs(
                                             save_and_offload_only_these_names_regex
                                         ).set(
                                             names_which_can_be_saved=None,
-                                            names_which_can_be_offloaded=None,
+                                            names_which_can_be_offloaded=(
+                                                RematRegexSavePatterns.INPUT.value
+                                            ),
+                                            offload_src="device",
+                                            offload_dst="pinned_host",
+                                        ),
+                                    ),
+                                }
+                            ),
+                        ],
+                    ),
+                ),
+                (
+                    "tpu-v7x-.*",
+                    ChainConfigModifier.default_config().set(
+                        config_modifiers=[
+                            MeshShapeModifier.default_config().set(
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=-1)
+                            ),
+                            # 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",
                                         ),