diff --git a/scripts/performance/utils/overrides.py b/scripts/performance/utils/overrides.py index 903c63d026..b14323609e 100644 --- a/scripts/performance/utils/overrides.py +++ b/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. diff --git a/src/megatron/bridge/diffusion/recipes/nemotron_labs_diffusion/ar_to_dlm.py b/src/megatron/bridge/diffusion/recipes/nemotron_labs_diffusion/ar_to_dlm.py index 7d99ab8c89..6247779334 100644 --- a/src/megatron/bridge/diffusion/recipes/nemotron_labs_diffusion/ar_to_dlm.py +++ b/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( diff --git a/src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py index 0a900ca287..8f09f54cf7 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py +++ b/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 diff --git a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py index 1333aee9d2..ed61cd7f5e 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py +++ b/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 diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 9fbd657f65..8011d7efbb 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/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 diff --git a/src/megatron/bridge/models/kimi/kimi_bridge.py b/src/megatron/bridge/models/kimi/kimi_bridge.py index 058a026f16..c019bbccea 100755 --- a/src/megatron/bridge/models/kimi/kimi_bridge.py +++ b/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 diff --git a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py index 6087fbbba6..cdff57e555 100644 --- a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +++ b/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 diff --git a/src/megatron/bridge/models/sarvam/sarvam_provider.py b/src/megatron/bridge/models/sarvam/sarvam_provider.py index 0057e6e001..c0986914ac 100644 --- a/src/megatron/bridge/models/sarvam/sarvam_provider.py +++ b/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"]) diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v2.py b/src/megatron/bridge/recipes/deepseek/deepseek_v2.py index b00fe7c5b5..6af056fc09 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v2.py +++ b/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" diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v3.py b/src/megatron/bridge/recipes/deepseek/deepseek_v3.py index 8520d992a2..e112a2e098 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v3.py +++ b/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" diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index a8a04963ed..18956d43c0 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/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"] diff --git a/src/megatron/bridge/recipes/kimi/kimi_k2.py b/src/megatron/bridge/recipes/kimi/kimi_k2.py index 62fe3a7784..8aa2397598 100644 --- a/src/megatron/bridge/recipes/kimi/kimi_k2.py +++ b/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" diff --git a/src/megatron/bridge/recipes/kimi_vl/kimi_k25_vl.py b/src/megatron/bridge/recipes/kimi_vl/kimi_k25_vl.py index 69ded93599..744a3d19c0 100644 --- a/src/megatron/bridge/recipes/kimi_vl/kimi_k25_vl.py +++ b/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 diff --git a/src/megatron/bridge/recipes/llama/llama3.py b/src/megatron/bridge/recipes/llama/llama3.py index 3ca001a5ad..03d623e4a3 100644 --- a/src/megatron/bridge/recipes/llama/llama3.py +++ b/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 diff --git a/src/megatron/bridge/recipes/moonlight/moonlight_16b.py b/src/megatron/bridge/recipes/moonlight/moonlight_16b.py index 6a1530769f..cddbd10a12 100644 --- a/src/megatron/bridge/recipes/moonlight/moonlight_16b.py +++ b/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 diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py index 74cb6957c0..8098ed414f 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py +++ b/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; diff --git a/src/megatron/bridge/recipes/qwen/qwen3.py b/src/megatron/bridge/recipes/qwen/qwen3.py index fb5a49ea06..befed2a929 100644 --- a/src/megatron/bridge/recipes/qwen/qwen3.py +++ b/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" diff --git a/src/megatron/bridge/recipes/qwen/qwen3_moe.py b/src/megatron/bridge/recipes/qwen/qwen3_moe.py index 58f1fd2ca6..7a978abab1 100644 --- a/src/megatron/bridge/recipes/qwen/qwen3_moe.py +++ b/src/megatron/bridge/recipes/qwen/qwen3_moe.py @@ -77,7 +77,7 @@ def qwen3_30b_a3b_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) - ENABLED for 30B MoE cfg.model.recompute_granularity = "full" @@ -191,7 +191,7 @@ def qwen3_235b_a22b_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) # Enable if needed for memory optimization @@ -321,7 +321,7 @@ def qwen3_30b_a3b_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) cfg.model.recompute_granularity = None @@ -454,7 +454,7 @@ def qwen3_235b_a22b_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) cfg.model.recompute_granularity = None @@ -601,7 +601,7 @@ def qwen3_30b_a3b_peft_config(peft_scheme: str | PEFT = "lora") -> ConfigContain 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 = None @@ -748,7 +748,7 @@ def qwen3_235b_a22b_peft_config(peft_scheme: str | PEFT = "lora") -> ConfigConta 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 = None diff --git a/src/megatron/bridge/recipes/stepfun/step35.py b/src/megatron/bridge/recipes/stepfun/step35.py index d4b17d87d6..95133ff7db 100644 --- a/src/megatron/bridge/recipes/stepfun/step35.py +++ b/src/megatron/bridge/recipes/stepfun/step35.py @@ -72,7 +72,7 @@ def step35_196b_a11b_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 = "full" diff --git a/tests/unit_tests/recipes/gemma/test_gemma2_recipes.py b/tests/unit_tests/recipes/gemma/test_gemma2_recipes.py index 754dc15260..6d784b5257 100644 --- a/tests/unit_tests/recipes/gemma/test_gemma2_recipes.py +++ b/tests/unit_tests/recipes/gemma/test_gemma2_recipes.py @@ -61,7 +61,7 @@ def _safe_overrides_for(name: str) -> dict: class _FakeModelCfg: def __init__(self): - self.cross_entropy_fusion_impl = "te" + self.cross_entropy_fusion_impl = "native" self.vocab_size = 256000 self.context_parallel_size = 1 diff --git a/tests/unit_tests/recipes/kimi/test_kimi_k2.py b/tests/unit_tests/recipes/kimi/test_kimi_k2.py index 7778ff1478..c21a2fca99 100644 --- a/tests/unit_tests/recipes/kimi/test_kimi_k2.py +++ b/tests/unit_tests/recipes/kimi/test_kimi_k2.py @@ -225,7 +225,7 @@ def test_pretrain_config_kernel_selections(self): assert cfg.model.attention_backend is None assert cfg.model.cross_entropy_loss_fusion is True - assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cross_entropy_fusion_impl == "native" def test_pretrain_config_comm_overlap(self): """Test communication overlap configuration.""" diff --git a/tests/unit_tests/recipes/kimi_vl/test_kimi_k25_vl.py b/tests/unit_tests/recipes/kimi_vl/test_kimi_k25_vl.py index 54017acc48..9e5c4c3776 100644 --- a/tests/unit_tests/recipes/kimi_vl/test_kimi_k25_vl.py +++ b/tests/unit_tests/recipes/kimi_vl/test_kimi_k25_vl.py @@ -267,12 +267,12 @@ def test_sft_config_transformer_engine_and_cuda_graph(self): assert cfg.model.cuda_graph_warmup_steps == 3 def test_sft_config_kernel_selections(self): - """Default attention backend is None; cross-entropy fusion uses TE.""" + """Default attention backend is None; cross-entropy fusion uses native.""" cfg = kimi_k25_vl_sft_config() assert cfg.model.attention_backend is None assert cfg.model.cross_entropy_loss_fusion is True - assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cross_entropy_fusion_impl == "native" def test_sft_config_comm_overlap(self): """Comm overlap is off (TP overlap, wgrad delay, MoE EP overlap).""" diff --git a/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py b/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py index 23737a0454..b8f9e4b7d2 100644 --- a/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py +++ b/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py @@ -66,7 +66,7 @@ class _FakeModelCfg: def __init__(self): - self.cross_entropy_fusion_impl = "te" + self.cross_entropy_fusion_impl = "native" self.context_parallel_size = 1 def finalize(self): diff --git a/tests/unit_tests/recipes/test_cross_entropy_defaults.py b/tests/unit_tests/recipes/test_cross_entropy_defaults.py new file mode 100644 index 0000000000..de64cd351e --- /dev/null +++ b/tests/unit_tests/recipes/test_cross_entropy_defaults.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ast +from pathlib import Path + +import pytest + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_SHIPPED_DEFAULT_ROOTS = ( + Path("src/megatron/bridge/recipes"), + Path("src/megatron/bridge/diffusion/recipes"), + Path("src/megatron/bridge/models"), + Path("scripts/performance"), +) + + +def _iter_shipped_python_files() -> list[Path]: + files: list[Path] = [] + for root in _SHIPPED_DEFAULT_ROOTS: + absolute_root = _REPO_ROOT / root + if absolute_root.exists(): + files.extend(sorted(absolute_root.rglob("*.py"))) + return files + + +def _is_cross_entropy_impl_target(target: ast.expr) -> bool: + if isinstance(target, ast.Attribute): + return target.attr == "cross_entropy_fusion_impl" + return isinstance(target, ast.Name) and target.id == "cross_entropy_fusion_impl" + + +def _is_te_literal(value: ast.expr | None) -> bool: + return isinstance(value, ast.Constant) and value.value == "te" + + +def _find_te_cross_entropy_defaults(path: Path) -> list[str]: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + failures: list[str] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Assign) and _is_te_literal(node.value): + for target in node.targets: + if _is_cross_entropy_impl_target(target): + failures.append(f"{path.relative_to(_REPO_ROOT)}:{node.lineno}") + elif isinstance(node, ast.AnnAssign) and _is_te_literal(node.value): + if _is_cross_entropy_impl_target(node.target): + failures.append(f"{path.relative_to(_REPO_ROOT)}:{node.lineno}") + elif isinstance(node, ast.Call): + for keyword in node.keywords: + if keyword.arg == "cross_entropy_fusion_impl" and _is_te_literal(keyword.value): + failures.append(f"{path.relative_to(_REPO_ROOT)}:{node.lineno}") + return failures + + +@pytest.mark.unit +def test_shipped_defaults_do_not_select_te_cross_entropy_fusion() -> None: + """Shipped defaults should prefer native CE fusion because MCore rejects the TE path.""" + failures: list[str] = [] + for path in _iter_shipped_python_files(): + failures.extend(_find_te_cross_entropy_defaults(path)) + + assert not failures, "Shipped defaults must use cross_entropy_fusion_impl='native' instead of 'te':\n" + "\n".join( + failures + ) diff --git a/tests/unit_tests/recipes/test_llama_recipes.py b/tests/unit_tests/recipes/test_llama_recipes.py index 698272e0cf..02f37b1da2 100644 --- a/tests/unit_tests/recipes/test_llama_recipes.py +++ b/tests/unit_tests/recipes/test_llama_recipes.py @@ -82,7 +82,7 @@ def _safe_overrides_for(name: str) -> dict: class _FakeModelCfg: def __init__(self): - self.cross_entropy_fusion_impl = "te" + self.cross_entropy_fusion_impl = "native" self.context_parallel_size = 1 def finalize(self): @@ -186,11 +186,7 @@ def test_each_llama_recipe_builds_config(recipe_func: Callable, monkeypatch: pyt assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 if "llama3" in recipe_func.__name__.lower(): - # Pretrain configs use "te", SFT/PEFT configs use "native" - expected_impl = ( - "native" if ("sft" in recipe_func.__name__.lower() or "peft" in recipe_func.__name__.lower()) else "te" - ) - assert cfg.model.cross_entropy_fusion_impl == expected_impl + assert cfg.model.cross_entropy_fusion_impl == "native" @pytest.mark.parametrize("recipe_func", _LLAMA3_SFT_FUNCS) diff --git a/tests/unit_tests/recipes/test_qwen_recipes.py b/tests/unit_tests/recipes/test_qwen_recipes.py index 540091f23c..ada66149ff 100644 --- a/tests/unit_tests/recipes/test_qwen_recipes.py +++ b/tests/unit_tests/recipes/test_qwen_recipes.py @@ -136,7 +136,7 @@ def test_each_qwen_recipe_builds_config(recipe_func: Callable, monkeypatch: pyte assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 if "qwen3" in recipe_name and "pretrain" in recipe_name and "next" not in recipe_name: - assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cross_entropy_fusion_impl == "native" # SFT and PEFT-specific assertions if is_sft_or_peft: diff --git a/tests/unit_tests/recipes/test_run_plugins.py b/tests/unit_tests/recipes/test_run_plugins.py index eaf62269e4..b471614132 100644 --- a/tests/unit_tests/recipes/test_run_plugins.py +++ b/tests/unit_tests/recipes/test_run_plugins.py @@ -105,7 +105,7 @@ def create_test_config(**kwargs): hidden_size=4096, ffn_hidden_size=14336, num_attention_heads=32, - cross_entropy_fusion_impl="te", + cross_entropy_fusion_impl="native", tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=pipeline_model_parallel_size, pipeline_dtype=pipeline_dtype,