Skip to content

Commit 68f4ff6

Browse files
judyzhaoxinwucamiloCienet
authored andcommitted
chore: revert accidental changes to test data folder during rebase
1 parent 65fcab9 commit 68f4ff6

35 files changed

Lines changed: 1987 additions & 2337 deletions

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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -246,7 +246,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
246246
mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
247247
mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
249-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
249+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
250250
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
251251
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
252252
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -257,13 +257,6 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
257257
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
258258
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
259259
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
260-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
261-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
262-
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'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
264-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
267260
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
268261
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
269262
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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
249249
mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
250250
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
251-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
251+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
252252
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
253253
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
254254
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -259,15 +259,8 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
259259
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
260260
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
261261
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
262-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
264-
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'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
267-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
268-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
269-
mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
270-
mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
262+
mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
263+
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
271264
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
272265
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
273266
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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -246,7 +246,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
246246
mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
247247
mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
249-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
249+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
250250
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
251251
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
252252
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -257,13 +257,6 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
257257
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
258258
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
259259
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
260-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
261-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
262-
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'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
264-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
267260
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
268261
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
269262
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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
249249
mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
250250
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
251-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
251+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
252252
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
253253
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
254254
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -259,15 +259,8 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
259259
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
260260
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
261261
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
262-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
264-
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'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
267-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
268-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
269-
mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
270-
mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
262+
mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
263+
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
271264
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
272265
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
273266
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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -246,7 +246,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
246246
mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
247247
mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
249-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
249+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
250250
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
251251
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
252252
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -257,13 +257,6 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
257257
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
258258
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
259259
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
260-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
261-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
262-
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'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
264-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
267260
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
268261
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
269262
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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
249249
mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
250250
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
251-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
251+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
252252
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
253253
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
254254
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -259,15 +259,8 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
259259
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
260260
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
261261
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
262-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
264-
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'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
267-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
268-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
269-
mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
270-
mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
262+
mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
263+
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
271264
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
272265
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
273266
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-v3-flash-fp8.txt

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
249249
mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
250250
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
251-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
251+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
252252
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
253253
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
254254
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -259,15 +259,8 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
259259
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
260260
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
261261
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
262-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
264-
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'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
267-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
268-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
269-
mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
270-
mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
262+
mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
263+
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
271264
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
272265
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
273266
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-v3-flash.txt

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -246,7 +246,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
246246
mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
247247
mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248
mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
249-
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
249+
mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
250250
mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
251251
mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
252252
mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -257,13 +257,6 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
257257
mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
258258
mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
259259
mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
260-
mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
261-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
262-
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'
263-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
264-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
265-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
266-
mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
267260
mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
268261
mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
269262
mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'

0 commit comments

Comments
 (0)