fix(model): store Muse Glimmer final norm uncentered for exact bf16 round trip

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

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import copy
 import functools
 from typing import TYPE_CHECKING, Any
 
@@ -565,6 +566,16 @@ def __init__(
             vp_stage=vp_stage,
         )
         self.model_config = config
+        if hasattr(self.decoder, "final_norm"):
+            # HF stores the final norm gain directly, unlike the zero-centered decoder-layer
+            # norms. Keep it uncentered so bf16 weights load without a lossy ``w - 1``.
+            final_norm_config = copy.copy(config.transformer)
+            final_norm_config.layernorm_zero_centered_gamma = False
+            self.decoder.final_norm = TENorm(
+                config=final_norm_config,
+                hidden_size=final_norm_config.hidden_size,
+                eps=final_norm_config.layernorm_epsilon,
+            )
         if hasattr(self, "output_layer"):
             extend_instance(self.output_layer, MuseGlimmerOutputLayerMixin)
         if pre_process:

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

@@ -31,7 +31,6 @@
     GatedMLPMapping,
     MegatronParamMapping,
     ReplicatedMapping,
-    RMSNorm2ZeroCenteredRMSNormMapping,
     merge_qkv_weights,
     split_qkv_weights,
 )
@@ -354,7 +353,7 @@ def mapping_registry(self) -> MegatronMappingRegistry:
                 megatron_param="output_layer.weight",
                 hf_param="lm_head.weight",
             ),
-            RMSNorm2ZeroCenteredRMSNormMapping(
+            ReplicatedMapping(
                 megatron_param="decoder.final_norm.weight",
                 hf_param=f"{text_prefix}.norm.weight",
             ),

tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py

@@ -23,7 +23,7 @@
 
 from megatron.bridge import AutoBridge
 from megatron.bridge.models.conversion.param_mapping import (
-    RMSNorm2ZeroCenteredRMSNormMapping,
+    ReplicatedMapping,
     merge_qkv_weights,
     split_qkv_weights,
 )
@@ -338,10 +338,7 @@ def test_affine_norms_use_transformer_engine_directly() -> None:
     assert set(vision_model.ln_pre.state_dict()) == {"weight", "bias", "_extra_state"}
 
     registry = MuseGlimmerBridge().mapping_registry()
-    assert isinstance(
-        registry.megatron_to_hf_lookup("decoder.final_norm.weight"),
-        RMSNorm2ZeroCenteredRMSNormMapping,
-    )
+    assert isinstance(registry.megatron_to_hf_lookup("decoder.final_norm.weight"), ReplicatedMapping)
     for name, _ in vision_model.named_parameters():
         assert registry.megatron_to_hf_lookup(f"vision_tower.{name}") is not None
 
@@ -493,7 +490,7 @@ def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGl
     assert dict(model.decoder.layers[0].self_attention.q_layernorm.named_parameters()) == {}
     assert dict(model.decoder.layers[0].self_attention.k_layernorm.named_parameters()) == {}
     assert isinstance(model.decoder.final_norm, TERMSNorm)
-    assert model.decoder.final_norm.zero_centered_gamma is True
+    assert model.decoder.final_norm.zero_centered_gamma is False
     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
@@ -621,3 +618,17 @@ def test_qkvg_mapping_executes_against_hybrid_qkv_module(tiny_hybrid_model: Muse
     torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.k_proj.weight"], key)
     torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.v_proj.weight"], value)
     torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.gate_proj.weight"], gate)
+
+
+def test_final_norm_mapping_round_trips_bf16_exactly(tiny_hybrid_model: MuseGlimmerModel) -> None:
+    final_norm = tiny_hybrid_model.decoder.final_norm
+    mapping = MuseGlimmerBridge().mapping_registry().megatron_to_hf_lookup("decoder.final_norm.weight")
+    hidden_size = final_norm.weight.numel()
+    # Most released gains that drifted lie in [-4, -3), where bf16 ``w - 1`` rounds to a coarser step.
+    hf_weight = torch.linspace(-5.0, -3.0, hidden_size).to(torch.bfloat16)
+
+    megatron_weight = mapping.hf_to_megatron(hf_weight, final_norm)
+    exported = mapping.megatron_to_hf(megatron_weight, final_norm)
+
+    assert torch.equal(megatron_weight, hf_weight)
+    assert torch.equal(exported["model.language_model.norm.weight"], hf_weight)