Skip to content

Commit a6692f3

Browse files
committed
[model] fix: restore explicit gradient_accumulation_fusion=True overrides
The dynamic checker (can_enable_gradient_accumulation_fusion) is unreliable in the training environment — it returns False even when the fused_weight_gradient_mlp_cuda extension is available at runtime. Restore explicit True overrides in DSV2, DSV3, Kimi bridges and NT3 Nano recipe to guarantee the fusion is always enabled for these models. Signed-off-by: Yu Yao <yaoyu.094@gmail.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
1 parent 41520f3 commit a6692f3

4 files changed

Lines changed: 6 additions & 0 deletions

File tree

‎src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
6262
provider.moe_permute_fusion = True
6363

6464
provider.apply_rope_fusion = False
65+
provider.gradient_accumulation_fusion = True
6566
provider.bias_activation_fusion = True
6667
provider.bias_dropout_fusion = True
6768
provider.cross_entropy_fusion_impl = "te"

‎src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
6969
provider.moe_aux_loss_coeff = 0.0001
7070

7171
provider.apply_rope_fusion = False
72+
provider.gradient_accumulation_fusion = True
7273
provider.bias_activation_fusion = True
7374
provider.bias_dropout_fusion = True
7475
provider.cross_entropy_fusion_impl = "te"

‎src/megatron/bridge/models/kimi/kimi_bridge.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
6464
provider.moe_permute_fusion = True
6565

6666
provider.apply_rope_fusion = False
67+
provider.gradient_accumulation_fusion = True
6768
provider.bias_activation_fusion = True
6869
provider.bias_dropout_fusion = True
6970
provider.cross_entropy_fusion_impl = "te"

‎src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ def nemotron_3_nano_pretrain_config() -> ConfigContainer:
187187

188188
cfg.model.init_method_std = 0.0173
189189
cfg.model.apply_rope_fusion = False
190+
cfg.model.gradient_accumulation_fusion = True
190191
cfg.model.use_fused_weighted_squared_relu = True
191192

192193
return cfg
@@ -299,6 +300,7 @@ def nemotron_3_nano_sft_config() -> ConfigContainer:
299300
cfg.model.moe_grouped_gemm = True
300301
cfg.model.cross_entropy_loss_fusion = True
301302
cfg.model.cross_entropy_fusion_impl = "native"
303+
cfg.model.gradient_accumulation_fusion = True
302304

303305
# Memory saving (recompute & offloading)
304306
cfg.model.recompute_granularity = None
@@ -488,6 +490,7 @@ def nemotron_3_nano_peft_config(
488490
cfg.model.moe_grouped_gemm = True
489491
cfg.model.cross_entropy_loss_fusion = True
490492
cfg.model.cross_entropy_fusion_impl = "native"
493+
cfg.model.gradient_accumulation_fusion = True
491494

492495
# Memory saving
493496
cfg.model.recompute_granularity = None

0 commit comments

Comments
 (0)