Save O_PROJ on Fuji 70B-v2 for neuron

apoorvtintin · apoorvtintin · commit 5966cc597af2 · 2025-05-19T18:25:28.000Z
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-flash.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2-flash.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v2.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-flash.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-tiktoken-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-tiktoken-flash.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-tiktoken.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3-tiktoken.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v3.txt
@@ -234,7 +234,7 @@ mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modif
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: True
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: None
-mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*linear1_[01]'
+mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj|.*o_proj|.*linear1_[01]'
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: None
 mesh_rules[7][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: None
 mesh_rules[7][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ModuleConfigModifier'
diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
@@ -810,6 +810,7 @@ def get_trainer_kwargs(
                                             names_which_can_be_saved="|".join(
                                                 [
                                                     RematRegexSavePatterns.QKV_PROJ.value,
+                                                    RematRegexSavePatterns.O_PROJ.value,
                                                     RematRegexSavePatterns.LINEAR1_X.value,
                                                 ]
                                             ),

Original file line number	Diff line number	Diff line change
`@@ -810,6 +810,7 @@ def get_trainer_kwargs(`
`810`	`810`	`names_which_can_be_saved="\|".join(`
`811`	`811`	`[`
`812`	`812`	`RematRegexSavePatterns.QKV_PROJ.value,`
	`813`	`+ RematRegexSavePatterns.O_PROJ.value,`
`813`	`814`	`RematRegexSavePatterns.LINEAR1_X.value,`
`814`	`815`	`]`
`815`	`816`	`),`