Update golden config for GB200

Author: andersensam

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-flash-fp8.txt

@@ -243,7 +243,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
 mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -254,8 +254,15 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
-mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
-mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
+mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-flash.txt

@@ -241,7 +241,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
 mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -252,6 +252,13 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-fp8.txt

@@ -243,7 +243,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
 mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -254,8 +254,15 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
-mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
-mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
+mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1.txt

@@ -241,7 +241,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
 mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -252,6 +252,13 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2-flash-fp8.txt

@@ -243,7 +243,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
 mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -254,8 +254,15 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
-mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
-mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
+mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2-flash.txt

@@ -241,7 +241,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
 mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -252,6 +252,13 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2-fp8.txt

@@ -243,7 +243,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
 mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -254,8 +254,15 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
-mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
-mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
+mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
+mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2.txt

@@ -241,7 +241,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
 mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
 mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -252,6 +252,13 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
+mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'