fix(muse-glimmer): decouple TE norms from attention backend

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

skills/adding-model-support/SKILL.md

@@ -185,6 +185,9 @@ Transformer Engine version pinned by Megatron-Core. Use the backend-provided `TE
 corresponding TE module) for affine LayerNorm and RMSNorm whenever TE implements the required
 semantics. Preserve `eps`, parameter dtype/device initialization, sequence-parallel metadata, and
 `zero_centered_gamma`; standard-gamma and zero-centered-gamma norms are not interchangeable.
+The attention/linear backend and normalization choice need not be coupled: a local attention spec
+may still use `TENorm`. Retain a semantics-equivalent local fallback when model construction must
+work in an environment where Transformer Engine itself is unavailable.
 
 Do not introduce a dummy affine parameter merely to reach a fused kernel. For example, current TE
 RMSNorm always owns a learnable `weight`, so a genuinely weightless/scaleless RMSNorm must retain a

src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py

@@ -24,6 +24,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from megatron.core import parallel_state
+from megatron.core.extensions.transformer_engine import HAVE_TE, TENorm
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.backends import get_backend
 from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules
@@ -132,20 +133,18 @@ def forward(self, hidden_states: Tensor) -> Tensor:
 
 
 def _centered_rms_norm_cls(config: TransformerConfig) -> type[nn.Module]:
-    """Select the backend-native centered RMSNorm when it supports Muse semantics."""
-    if config.transformer_impl == "local":
-        return MuseGlimmerCenteredRMSNorm
-    return get_backend(config.transformer_impl).layer_norm(rms_norm=True)
+    """Use TE for affine normalization whenever it is available."""
+    return TENorm if HAVE_TE else MuseGlimmerCenteredRMSNorm
 
 
 def _standard_rms_norm(config: TransformerConfig, hidden_size: int, eps: float) -> nn.Module:
-    """Build a standard-gamma RMSNorm with the configured model backend."""
-    if config.transformer_impl == "local":
+    """Build a standard-gamma RMSNorm with TE when it is available."""
+    if not HAVE_TE:
         return MuseGlimmerRMSNorm(config, hidden_size, eps=eps)
 
     norm_config = copy(config)
     norm_config.layernorm_zero_centered_gamma = False
-    return get_backend(config.transformer_impl).layer_norm(rms_norm=True)(
+    return TENorm(
         norm_config,
         hidden_size,
         eps=eps,
@@ -156,15 +155,15 @@ def _vision_layer_norm(
     config: MuseGlimmerVisionModelConfig,
     transformer_config: TransformerConfig,
 ) -> nn.Module:
-    """Build the vision LayerNorm with TE unless the local backend is explicit."""
-    if transformer_config.transformer_impl == "local":
+    """Build the vision LayerNorm with TE when it is available."""
+    if not HAVE_TE:
         return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
 
     norm_config = copy(transformer_config)
     norm_config.normalization = "LayerNorm"
     norm_config.layernorm_zero_centered_gamma = False
     norm_config.sequence_parallel = False
-    return get_backend(transformer_config.transformer_impl).layer_norm()(
+    return TENorm(
         norm_config,
         config.hidden_size,
         eps=config.layer_norm_epsilon,

src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py

@@ -216,7 +216,7 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM
             params_dtype=params_dtype,
             fp16=params_dtype == torch.float16,
             bf16=params_dtype == torch.bfloat16,
-            transformer_impl="transformer_engine",
+            transformer_impl="local",
             output_multiplier=float(text_config.output_multiplier),
             final_logit_softcapping=float(text_config.final_logit_softcapping),
         )

tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py

@@ -113,7 +113,7 @@ def test_config_conversion_uses_model_config_path_only() -> None:
     assert isinstance(model_config, MuseGlimmerModelConfig)
     assert isinstance(model_config, HybridModelConfig)
     assert model_config.get_builder_cls() is MuseGlimmerModelBuilder
-    assert model_config.transformer_impl == "transformer_engine"
+    assert model_config.transformer_impl == "local"
     assert model_config.num_layers == 4
     assert model_config.hidden_size == 64
     assert model_config.ffn_hidden_size == 128
@@ -209,7 +209,7 @@ def test_hybrid_attention_layers_use_full_transformer_flops() -> None:
 @pytest.mark.parametrize(
     ("transformer_impl", "expected_qkv", "expected_attention", "expected_norm"),
     [
-        ("local", ColumnParallelLinear, DotProductAttention, MuseGlimmerCenteredRMSNorm),
+        ("local", ColumnParallelLinear, DotProductAttention, TENorm),
         ("transformer_engine", TEColumnParallelLinear, TEDotProductAttention, TENorm),
     ],
 )
@@ -358,6 +358,20 @@ def test_weightless_rms_norm_does_not_add_a_transformer_engine_parameter() -> No
     assert norm.state_dict() == {}
 
 
+def test_affine_norms_retain_a_local_fallback_without_transformer_engine() -> None:
+    model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config())
+    config = model_config.transformer
+
+    with patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False):
+        stack_spec = get_muse_glimmer_hybrid_stack_spec(config)
+        final_norm = _standard_rms_norm(config, config.hidden_size, config.layernorm_epsilon)
+        vision_model = MuseGlimmerVisionModel(model_config.vision, config)
+
+    assert stack_spec.submodules.attention_layer.submodules.input_layernorm is MuseGlimmerCenteredRMSNorm
+    assert isinstance(final_norm, MuseGlimmerRMSNorm)
+    assert isinstance(vision_model.ln_pre, torch.nn.LayerNorm)
+
+
 def test_centered_norm_sharded_state_dict_identifies_tensor_parallel_replicas() -> None:
     config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer
     norm = MuseGlimmerCenteredRMSNorm(config, hidden_size=4)
@@ -403,11 +417,12 @@ def test_tiny_vision_model_preserves_expected_token_count() -> None:
     model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config())
     model_config.transformer_impl = "local"
     vision_config = model_config.vision
-    model = MuseGlimmerVisionModel(vision_config, model_config.transformer)
     pixel_values = torch.randn(16, 24)
     grid_thw = torch.tensor([[1, 4, 4]])
 
-    output = model(pixel_values, grid_thw)
+    with patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False):
+        model = MuseGlimmerVisionModel(vision_config, model_config.transformer)
+        output = model(pixel_values, grid_thw)
 
     assert output.shape == (4, 64)
 
@@ -460,14 +475,17 @@ def test_vision_layer_recomputation_is_training_only() -> None:
     model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config())
     model_config.transformer_impl = "local"
     vision_config = model_config.vision
-    model = MuseGlimmerVisionModel(vision_config, model_config.transformer, recompute_layers=True)
     pixel_values = torch.randn(16, 24)
     grid_thw = torch.tensor([[1, 4, 4]])
 
-    with patch(
-        "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.torch_checkpoint",
-        side_effect=lambda function, *args, **kwargs: function(*args),
-    ) as checkpoint:
+    with (
+        patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False),
+        patch(
+            "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.torch_checkpoint",
+            side_effect=lambda function, *args, **kwargs: function(*args),
+        ) as checkpoint,
+    ):
+        model = MuseGlimmerVisionModel(vision_config, model_config.transformer, recompute_layers=True)
         training_output = model(pixel_values, grid_thw)
         model.eval()
         inference_output = model(pixel_values, grid_thw)
@@ -524,7 +542,9 @@ def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGl
     assert "decoder.layers.0.self_attention.post_layernorm.weight" in names
     assert "decoder.layers.0.mlp.post_layernorm.weight" in names
     assert not any(name.startswith("language_model.") for name in names)
-    assert isinstance(model.decoder.final_norm, MuseGlimmerRMSNorm)
+    assert isinstance(model.decoder.final_norm, TERMSNorm)
+    assert isinstance(model.vision_tower.ln_pre, TELayerNorm)
+    assert isinstance(model.vision_tower.layers[0].norm1, TELayerNorm)
     assert names["vision_tower.patch_embedder.patch_embedding.weight"].dtype == torch.float32
     assert all(
         getattr(parameter, "average_gradients_across_tp_domain", False)