fix(recipes): default cross entropy fusion to native

scripts/performance/utils/overrides.py

@@ -67,7 +67,7 @@ def _set_common_perf_overrides(recipe: ConfigContainer) -> ConfigContainer:
     if hasattr(recipe.model, "apply_rope_fusion"):
         recipe.model.apply_rope_fusion = True
     if hasattr(recipe.model, "cross_entropy_fusion_impl"):
-        recipe.model.cross_entropy_fusion_impl = "te"
+        recipe.model.cross_entropy_fusion_impl = "native"
 
     # TODO: This needs to be adjusted when overlapping HybridEP with computation or
     # the number of SMs for HybridEP is reduced.

src/megatron/bridge/diffusion/recipes/nemotron_labs_diffusion/ar_to_dlm.py

@@ -229,7 +229,7 @@ def _nemotron_labs_diffusion_common(
         model_cfg.recompute_num_layers = 1
 
     model_cfg.cross_entropy_loss_fusion = True
-    model_cfg.cross_entropy_fusion_impl = "te"
+    model_cfg.cross_entropy_fusion_impl = "native"
 
     if lr_decay_style == "WSD":
         opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing_dllm(

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

@@ -65,7 +65,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.gradient_accumulation_fusion = True
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
-        provider.cross_entropy_fusion_impl = "te"
+        provider.cross_entropy_fusion_impl = "native"
         provider.cross_entropy_loss_fusion = True
         provider.masked_softmax_fusion = True
         provider.persist_layer_norm = True

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

@@ -79,7 +79,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.gradient_accumulation_fusion = True
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
-        provider.cross_entropy_fusion_impl = "te"
+        provider.cross_entropy_fusion_impl = "native"
         provider.cross_entropy_loss_fusion = True
         provider.masked_softmax_fusion = True
         provider.persist_layer_norm = True

src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py

@@ -441,7 +441,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.share_embeddings_and_output_weights = bool(hf_config.tie_word_embeddings)
         provider.gradient_accumulation_fusion = True
         provider.bias_dropout_fusion = True
-        provider.cross_entropy_fusion_impl = "te"
+        provider.cross_entropy_fusion_impl = "native"
         provider.cross_entropy_loss_fusion = True
         provider.masked_softmax_fusion = True
         provider.persist_layer_norm = True

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

@@ -67,7 +67,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.gradient_accumulation_fusion = True
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
-        provider.cross_entropy_fusion_impl = "te"
+        provider.cross_entropy_fusion_impl = "native"
         provider.cross_entropy_loss_fusion = True
         provider.masked_softmax_fusion = True
         provider.persist_layer_norm = True

src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py

@@ -105,7 +105,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> KimiK25VLModelProvide
         provider.apply_rope_fusion = False
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
-        provider.cross_entropy_fusion_impl = "te"
+        provider.cross_entropy_fusion_impl = "native"
         provider.cross_entropy_loss_fusion = True
         provider.masked_softmax_fusion = True
         provider.persist_layer_norm = True

src/megatron/bridge/models/sarvam/sarvam_provider.py

@@ -107,7 +107,7 @@ class SarvamMoEModelProvider(GPTModelProvider):
     attention_softmax_in_fp32: bool = True
     persist_layer_norm: bool = True
 
-    cross_entropy_fusion_impl: str = "te"
+    cross_entropy_fusion_impl: str = "native"
     cp_comm_type: str = "p2p"
     recompute_granularity: str = "selective"
     recompute_modules: List[str] = field(default_factory=lambda: ["layernorm", "shared_experts", "mlp", "moe_act"])
@@ -192,7 +192,7 @@ class SarvamMLAModelProvider(MLATransformerConfig, GPTModelProvider):
     attention_softmax_in_fp32: bool = True
     persist_layer_norm: bool = True
 
-    cross_entropy_fusion_impl: str = "te"
+    cross_entropy_fusion_impl: str = "native"
     cp_comm_type: str = "p2p"
     recompute_granularity: str = "selective"
     recompute_modules: List[str] = field(default_factory=lambda: ["moe"])

src/megatron/bridge/recipes/deepseek/deepseek_v2.py

@@ -84,7 +84,7 @@ def deepseek_v2_lite_pretrain_config() -> ConfigContainer:
         True  # MoE-specific: Use grouped GEMM for experts (default from DeepSeekModelProvider)
     )
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"  # Default from DeepSeekModelProvider
+    cfg.model.cross_entropy_fusion_impl = "native"  # Default from DeepSeekModelProvider
 
     # Memory saving (recompute & offloading) - ENABLED for V2-Lite
     cfg.model.recompute_granularity = "full"
@@ -202,7 +202,7 @@ def deepseek_v2_pretrain_config() -> ConfigContainer:
         True  # MoE-specific: Use grouped GEMM for experts (default from DeepSeekModelProvider)
     )
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"  # Default from DeepSeekModelProvider
+    cfg.model.cross_entropy_fusion_impl = "native"  # Default from DeepSeekModelProvider
 
     # Memory saving (recompute & offloading) - ENABLED for V2
     cfg.model.recompute_granularity = "full"

src/megatron/bridge/recipes/deepseek/deepseek_v3.py

@@ -170,7 +170,7 @@ def deepseek_v3_pretrain_config() -> ConfigContainer:
     cfg.model.moe_permute_fusion = True  # MoE-specific: Fuse permute operations
     cfg.model.moe_grouped_gemm = True  # MoE-specific: Use grouped GEMM for experts
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"  # Default from DeepSeekModelProvider
+    cfg.model.cross_entropy_fusion_impl = "native"  # Default from DeepSeekModelProvider
 
     # Memory saving (recompute & offloading) — no recompute by default.
     # Setting granularity="selective" with modules=None would cause MCore's
@@ -326,7 +326,7 @@ def deepseek_v3_pretrain_config_32nodes() -> ConfigContainer:
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving - FULL recompute for 32 nodes (memory efficiency)
     cfg.model.recompute_granularity = "full"

src/megatron/bridge/recipes/deepseek/deepseek_v4.py

@@ -121,7 +121,7 @@ def deepseek_v4_flash_pretrain_config() -> ConfigContainer:
     cfg.model.moe_aux_loss_coeff = 0.0
     cfg.model.moe_router_force_load_balancing = False
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = "selective"
     cfg.model.recompute_modules = ["moe_act", "mhc"]

src/megatron/bridge/recipes/kimi/kimi_k2.py

@@ -146,7 +146,7 @@ def kimi_k2_pretrain_config(optimizer_type: str = "muon") -> ConfigContainer:
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading) - already set in model provider
     # cfg.model.recompute_granularity = "selective"

src/megatron/bridge/recipes/kimi_vl/kimi_k25_vl.py

@@ -135,7 +135,7 @@ def kimi_k25_vl_sft_config() -> ConfigContainer:
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving
     cfg.model.fine_grained_activation_offloading = False

src/megatron/bridge/recipes/llama/llama3.py

@@ -95,7 +95,7 @@ def llama32_1b_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -182,7 +182,7 @@ def llama32_3b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -268,7 +268,7 @@ def llama3_8b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -349,7 +349,7 @@ def llama3_8b_16k_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -430,7 +430,7 @@ def llama3_8b_64k_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -511,7 +511,7 @@ def llama3_8b_128k_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -614,7 +614,7 @@ def llama3_8b_low_precision_pretrain_config(mixed_precision_recipe: str) -> Conf
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -701,7 +701,7 @@ def llama3_70b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -804,7 +804,7 @@ def llama3_70b_16k_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -893,7 +893,7 @@ def llama3_70b_64k_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -987,7 +987,7 @@ def llama31_8b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -1070,7 +1070,7 @@ def llama31_70b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None
@@ -1165,7 +1165,7 @@ def llama31_405b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     cfg.model.recompute_granularity = None
     cfg.model.recompute_modules = None

src/megatron/bridge/recipes/moonlight/moonlight_16b.py

@@ -125,7 +125,7 @@ def moonlight_16b_pretrain_config() -> ConfigContainer:
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = "selective"
@@ -307,7 +307,7 @@ def moonlight_16b_sft_config() -> ConfigContainer:
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     # recompute_granularity already set in model provider
@@ -520,7 +520,7 @@ def moonlight_16b_peft_config(
     cfg.model.moe_permute_fusion = True
     cfg.model.moe_grouped_gemm = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving
     cfg.model.fine_grained_activation_offloading = False

src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py

@@ -86,7 +86,7 @@ def nemotron_3_super_pretrain_config() -> ConfigContainer:
 
     # Kernel Selections
     cfg.model.attention_backend = "fused"
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
     cfg.model.use_te_rng_tracker = True
 
     # MTP Settings (HF config has num_nextn_predict_layers=1 for the shared block;

src/megatron/bridge/recipes/qwen/qwen3.py

@@ -67,7 +67,7 @@ def qwen3_600m_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -148,7 +148,7 @@ def qwen3_1p7b_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -229,7 +229,7 @@ def qwen3_4b_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -310,7 +310,7 @@ def qwen3_8b_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -391,7 +391,7 @@ def qwen3_14b_pretrain_config() -> ConfigContainer:
     # Kernel selections
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading)
     cfg.model.recompute_granularity = None
@@ -464,7 +464,7 @@ def qwen3_32b_pretrain_config() -> ConfigContainer:
 
     cfg.model.attention_backend = None
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.cross_entropy_fusion_impl = "native"
 
     # Memory saving (recompute & offloading) - ENABLED for 32B
     cfg.model.recompute_granularity = "full"