cp: Clean up fusion checks (938) into r0.2.0

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

@@ -109,6 +109,7 @@ class DeepSeekModelProvider(MLATransformerConfig, GPTModelProvider):
     bias_activation_fusion: bool = True
     bias_dropout_fusion: bool = True
     masked_softmax_fusion: bool = True
+    gradient_accumulation_fusion: bool = True
     cross_entropy_loss_fusion: bool = True
     cross_entropy_fusion_impl: str = "te"
     moe_permute_fusion: bool = is_te_min_version("2.1.0") if HAVE_TE else False
@@ -180,6 +181,7 @@ class DeepSeekV3ModelProvider(DeepSeekModelProvider):
     moe_router_num_groups: int = 8
     moe_router_group_topk: int = 4
     moe_router_topk_scaling_factor: float = 2.5
+    moe_aux_loss_coeff: float = 1e-4
     make_vocab_size_divisible_by: int = 1280
     moe_router_score_function: str = "sigmoid"
     moe_router_enable_expert_bias: bool = True

src/megatron/bridge/models/gemma/gemma3_provider.py

@@ -40,7 +40,6 @@
 from torch import Tensor
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.utils import fusions
 from megatron.bridge.utils.import_utils import safe_import_from
 
 
@@ -81,7 +80,7 @@ class Gemma3ModelProvider(GPTModelProvider):
     # mlp
     gated_linear_unit: bool = True
     add_bias_linear: bool = False
-    activation_func: Callable = field(default_factory=lambda: fast_gelu)  # identical to openai_gelu
+    activation_func: Callable = fast_gelu  # identical to openai_gelu
 
     # Do not change
     is_vision_language: bool = False
@@ -91,8 +90,6 @@ class Gemma3ModelProvider(GPTModelProvider):
         default_factory=lambda: gemma3_layer_spec
     )
     scatter_embedding_sequence_parallel: bool = True
-    apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
-    masked_softmax_fusion: bool = field(default_factory=fusions.can_enable_masked_softmax_fusion)
 
     # Data type settings to match HF models
     bf16: bool = True

src/megatron/bridge/models/gpt_provider.py

@@ -165,17 +165,12 @@ class GPTModelProvider(TransformerConfig, ModelProviderMixin[MCoreGPTModel]):
     account_for_loss_in_pipeline_split: bool = False
 
     # Fusions
-    masked_softmax_fusion: bool = field(default_factory=fusions.can_enable_masked_softmax_fusion)
+    masked_softmax_fusion: bool = True
     cross_entropy_loss_fusion: bool = True  # Generally beneficial, no specific dependencies
     gradient_accumulation_fusion: bool = field(default_factory=fusions.can_enable_gradient_accumulation_fusion)
-    bias_activation_fusion: bool = False  # Disabled by default as it can interfere with certain architectures
-    persist_layer_norm: bool = False
-    bias_dropout_fusion: bool = field(default_factory=fusions.can_enable_bias_dropout_fusion)
-    apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
 
     # If True, restore the modelopt_state that contains quantization, sparsity, speculative decoding transformation state.
     # When resuming modelopt_state, we also change the transformer_layer_spec to `megatron.core.post_training.modelopt.gpt.model_specs` which is a combination of local spec + TEDotProductAttention.
-
     restore_modelopt_state: bool = False
 
     def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGPTModel:
@@ -339,6 +334,7 @@ class GPTProvider126M(GPTModelProvider):
     num_attention_heads: int = 12
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
 
 
 @dataclass
@@ -356,6 +352,7 @@ class GPTProvider5B(GPTModelProvider):
     num_attention_heads: int = 32
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
 
 
 @dataclass
@@ -373,6 +370,7 @@ class GPTProvider7B(GPTModelProvider):
     num_attention_heads: int = 32
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
 
 
 @dataclass
@@ -390,6 +388,7 @@ class GPTProvider20B(GPTModelProvider):
     num_attention_heads: int = 48
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
 
 
 @dataclass
@@ -407,6 +406,7 @@ class GPTProvider40B(GPTModelProvider):
     num_attention_heads: int = 64
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
 
 
 @dataclass
@@ -426,4 +426,5 @@ class GPTProvider175B(GPTModelProvider):
     attention_dropout: float = 0.0
     bias_activation_fusion: bool = True
     bias_dropout_add_fusion: bool = True
+    use_transformer_engine_full_layer_spec: bool = True
     layernorm_zero_centered_gamma: bool = True

src/megatron/bridge/models/llama/llama_provider.py

@@ -24,7 +24,6 @@
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.llama.llama4_utils import get_llama4_layer_spec
-from megatron.bridge.utils import fusions
 
 
 logger = logging.getLogger(__name__)
@@ -51,10 +50,10 @@ class LlamaModelProvider(GPTModelProvider):
     share_embeddings_and_output_weights: bool = False
     # Fusions
     bias_activation_fusion: bool = True
-    masked_softmax_fusion: bool = field(default_factory=fusions.can_enable_masked_softmax_fusion)
-    bias_dropout_fusion: bool = field(default_factory=fusions.can_enable_bias_dropout_fusion)
-    apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
-    gradient_accumulation_fusion: bool = field(default_factory=fusions.can_enable_gradient_accumulation_fusion)
+    masked_softmax_fusion: bool = True
+    persist_layer_norm: bool = True
+    bias_dropout_fusion: bool = True
+    apply_rope_fusion: bool = True
     use_transformer_engine_op_fuser: Optional[bool] = None
 
 
@@ -123,9 +122,10 @@ class Llama3ModelProvider(LlamaModelProvider):
     gated_linear_unit: bool = True
     # Fusions
     bias_activation_fusion: bool = True
-    masked_softmax_fusion: bool = field(default_factory=fusions.can_enable_masked_softmax_fusion)
-    bias_dropout_fusion: bool = field(default_factory=fusions.can_enable_bias_dropout_fusion)
-    apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
+    masked_softmax_fusion: bool = True
+    persist_layer_norm: bool = True
+    bias_dropout_fusion: bool = True
+    apply_rope_fusion: bool = True
     share_embeddings_and_output_weights: bool = False
     position_embedding_type: str = "rope"
     rotary_percent: float = 1.0
@@ -271,6 +271,7 @@ class Llama32ModelProvider1B(Llama31ModelProvider):
     scale_factor: float = 32.0
     share_embeddings_and_output_weights: bool = True
     rotary_base: int = 500_000
+    seq_length: int = 131072
     num_layers: int = 16
     hidden_size: int = 2048
     ffn_hidden_size: int = 8192
@@ -290,6 +291,7 @@ class Llama32ModelProvider3B(Llama31ModelProvider):
     scale_factor: int = 32
     share_embeddings_and_output_weights: bool = True
     rotary_base: int = 500_000
+    seq_length: int = 131072
     num_layers: int = 28
     hidden_size: int = 3072
     ffn_hidden_size: int = 8192

src/megatron/bridge/models/nemotron/nemotron_provider.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 
 import logging
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Callable, Optional
 
 import torch
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.utils import fusions
 
 
 logger = logging.getLogger(__name__)
@@ -44,12 +43,12 @@ class NemotronModelProvider(GPTModelProvider):
     hidden_dropout: float = 0.0
     attention_dropout: float = 0.0
     rotary_percent: float = 0.5
-    masked_softmax_fusion: bool = field(default_factory=fusions.can_enable_masked_softmax_fusion)
+    masked_softmax_fusion: bool = True
     persist_layer_norm: bool = True
     bias_dropout_add_fusion: bool = False
     layernorm_zero_centered_gamma: bool = True
     cross_entropy_loss_fusion: bool = True
-    apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
+    apply_rope_fusion: bool = True
 
     # Nemotron3Config4B as default configs
     num_layers: int = 32

src/megatron/bridge/models/qwen/qwen_provider.py

@@ -246,6 +246,7 @@ class Qwen3ModelProvider(GPTModelProvider):
     kv_channels: Optional[int] = 128
     num_query_groups: int = 8
     seq_length: int = 40960
+    max_position_embeddings: int = 40960
     init_method_std: int = 0.02
     hidden_dropout: float = 0.0
     attention_dropout: float = 0.0
@@ -352,6 +353,7 @@ class Qwen3MoEModelProvider(GPTModelProvider):
     kv_channels: Optional[int] = 128
     num_query_groups: int = 8
     seq_length: int = 40960
+    max_position_embeddings: int = 40960
     init_method_std: int = 0.02
     hidden_dropout: float = 0.0
     attention_dropout: float = 0.0

src/megatron/bridge/utils/fusions.py

@@ -31,50 +31,6 @@
 LOG_FUSION_DISABLE = os.environ.get("MEGATRON_SUPPRESS_FUSION_WARNINGS", "0") != "1"
 
 
-def can_enable_apply_rope_fusion() -> bool:
-    """Check if RoPE (Rotary Position Embedding) fusion can be enabled.
-
-    Returns:
-        bool: True if RoPE fusion is available and compatible.
-    """
-    # Check for Transformer Engine availability
-    try:
-        import transformer_engine  # noqa: F401
-        from megatron.core.utils import get_te_version, is_te_min_version
-
-        if not is_te_min_version("2.2.0.dev0"):
-            if LOG_FUSION_DISABLE:
-                logger.warning(
-                    "apply_rope_fusion requires Transformer Engine >= 2.2.0.dev0. "
-                    f"Current version: {get_te_version()}. Fusion disabled."
-                )
-            return False
-    except ImportError:
-        if LOG_FUSION_DISABLE:
-            logger.warning("apply_rope_fusion requires Transformer Engine but it is not installed. Fusion disabled.")
-        return False
-
-    # Check for RoPE fusion kernel availability
-    try:
-        from megatron.core.models.common.embeddings.rope_utils import (
-            fused_apply_rotary_pos_emb,
-            fused_apply_rotary_pos_emb_thd,
-        )
-
-        if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None:
-            if LOG_FUSION_DISABLE:
-                logger.warning("apply_rope_fusion kernels are not available in megatron.core. Fusion disabled.")
-            return False
-        return True
-    except ImportError:
-        if LOG_FUSION_DISABLE:
-            logger.warning(
-                "apply_rope_fusion requires RoPE fusion kernels from megatron.core but they are not available. "
-                "Fusion disabled."
-            )
-        return False
-
-
 def can_enable_gradient_accumulation_fusion() -> bool:
     """Check if gradient accumulation fusion can be enabled.
 
@@ -94,59 +50,29 @@ def can_enable_gradient_accumulation_fusion() -> bool:
         return False
 
 
-def can_enable_bias_dropout_fusion() -> bool:
-    """Check if bias dropout fusion can be enabled.
-
-    Returns:
-        bool: True if bias dropout fusion is available.
-    """
-    try:
-        from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train  # noqa: F401
-
-        return True
-    except ImportError:
-        if LOG_FUSION_DISABLE:
-            logger.warning(
-                "bias_dropout_fusion requires fused_bias_dropout from megatron.core.fusions "
-                "but it is not available. Fusion disabled."
-            )
-        return False
-
-
-def can_enable_masked_softmax_fusion() -> bool:
-    """Check if masked softmax fusion can be enabled.
-
-    Returns:
-        bool: True if masked softmax fusion kernels are available.
-    """
-    try:
-        # Try to import the CUDA kernels that are required for masked softmax fusion
-        import scaled_masked_softmax_cuda  # noqa: F401
-        import scaled_upper_triang_masked_softmax_cuda  # noqa: F401
-
-        return True
-    except ImportError:
-        if LOG_FUSION_DISABLE:
-            logger.warning(
-                "masked_softmax_fusion requires CUDA kernels (scaled_masked_softmax_cuda) "
-                "but they are not available. This typically happens when Megatron-Core is not "
-                "built with CUDA extensions. Fusion disabled."
-            )
-        return False
-
-
 def validate_rope_fusion_compatibility(config: TransformerConfig) -> bool:
     """Validate if RoPE fusion is compatible with the current model configuration.
 
     Args:
-        model_provider: The GPTModelProvider instance to validate.
+        config: The TransformerConfig instance to validate.
 
     Returns:
         bool: True if RoPE fusion is compatible, False otherwise.
     """
     if not config.apply_rope_fusion:
         return True
 
+    # Check if position embedding type is RoPE (similar to arguments.py logic)
+    position_embedding_type = getattr(config, "position_embedding_type", "learned_absolute")
+    if position_embedding_type != "rope":
+        if LOG_FUSION_DISABLE:
+            logger.warning(
+                f"apply_rope_fusion is only compatible with RoPE position embeddings. "
+                f"Current position_embedding_type: {position_embedding_type}. "
+                f"Consider disabling apply_rope_fusion."
+            )
+        return False
+
     # Check for multi_latent_attention incompatibility
     if getattr(config, "multi_latent_attention", False):
         if LOG_FUSION_DISABLE:
@@ -156,24 +82,4 @@ def validate_rope_fusion_compatibility(config: TransformerConfig) -> bool:
             )
         return True
 
-    # Check TE version for rotary_interleaved
-    if getattr(config, "rotary_interleaved", False):
-        try:
-            from megatron.core.utils import get_te_version, is_te_min_version
-
-            if not is_te_min_version("2.2.0.dev0"):
-                if LOG_FUSION_DISABLE:
-                    logger.warning(
-                        "apply_rope_fusion with rotary_interleaved requires TE >= 2.2.0.dev0. "
-                        f"Current TE version: {get_te_version()}. Consider disabling apply_rope_fusion."
-                    )
-                return False
-        except ImportError:
-            if LOG_FUSION_DISABLE:
-                logger.warning(
-                    "apply_rope_fusion with rotary_interleaved requires Transformer Engine. "
-                    "Consider disabling apply_rope_fusion."
-                )
-            return False
-
     return True

tests/unit_tests/models/llama/test_llama_provider.py

@@ -56,7 +56,6 @@ def test_llama_model_provider_initialization(self):
         assert provider.normalization == "RMSNorm"
         assert provider.add_bias_linear is False
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.persist_layer_norm is False
 
     def test_llama_model_provider_with_custom_rope(self):
         """Test LlamaModelProvider with custom RoPE configuration."""