@@ -246,7 +246,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[3]: 64
246246mesh_rules[7][1].config_modifiers[0].mesh_shape[4]: 1
247247mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
248248mesh_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)'
250250mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
251251mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
252252mesh_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
257257mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
258258mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
259259mesh_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'
267260mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
268261mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
269262mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
0 commit comments