Update Fuji mesh selector for TPU 7x

GitOrigin-RevId: a4bbda9

Author: andersensam

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-flash-fp8.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-flash.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-fp8.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-tiktoken-flash-fp8.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-tiktoken-flash.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-tiktoken-fp8.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-405B-v3-tiktoken.txt

@@ -144,18 +144,22 @@ mesh_rules[1][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[1][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[1]: -1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[2]: 1
-mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: -1
+mesh_rules[1][1].config_modifiers[0].mesh_shape[3]: 256
 mesh_rules[1][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[1][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.experiments.trainer_config_utils.V7xFlashConfigModifier'
 mesh_rules[1][1].config_modifiers[1].tpu_block_size: 2048
-mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[1][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.experiments.trainer_config_utils.SplashAttentionConfigModifier'
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_kv_dkv_compute: 256
+mesh_rules[1][1].config_modifiers[2].splash_block_q_dkv: 1024
+mesh_rules[1][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: None
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[1][1].config_modifiers[3].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_shape[0]: 1
 mesh_shape[1]: 1