From ffcf2d6b6151d3d0820ead11e73df84552f08742 Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Tue, 29 Sep 2026 17:33:25 -0700 Subject: [PATCH 1/2] fix(model): store Muse Glimmer final norm uncentered for exact bf16 round trip HF Muse applies its final RMSNorm gain directly (w), unlike the decoder-layer norms (1 + w). The bridge mapped decoder.final_norm.weight through RMSNorm2ZeroCenteredRMSNormMapping, which stores w - 1 in bf16. For Muse's final-norm gains (about -5 to -3) that subtraction rounds, so 1,120 of 6,656 elements of model.language_model.norm.weight changed by one bf16 step after an HF -> Megatron -> HF round trip. Build the Megatron final norm with zero_centered_gamma=False and map it with ReplicatedMapping so the released weights are stored as-is. Add a regression test that round-trips bf16 gains in [-5, -3) through the mapping bit-exactly. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Kamran Jafari --- .../muse_glimmer/modeling_muse_glimmer.py | 11 +++++++++ .../muse_glimmer/muse_glimmer_bridge.py | 3 +-- .../models/muse_glimmer/test_muse_glimmer.py | 23 ++++++++++++++----- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index bdac50530b..f9d69dfb6d 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/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: diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index 61097b3595..7ac1e5d3e1 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/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", ), diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index a9e79505dd..e8f30ebe5a 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/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() + # Muse final-norm gains sit in [-5, -3); a bf16 ``w - 1`` round trip loses a bit in this range. + 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) From c9b763d5b6325b0e9b77102bb9e16067cf1b33c9 Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Tue, 29 Sep 2026 20:11:05 -0700 Subject: [PATCH 2/2] test(model): correct Muse final-norm gain range in comment The released final-norm gains span about -4.9 to 4.4, not [-5, -3). Most elements that drifted under the old zero-centered mapping (1,061 of 1,120) lie in [-4, -3), where bf16 w - 1 rounds to a coarser step. The test input range is unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Kamran Jafari --- tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index e8f30ebe5a..461c4a9b33 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -624,7 +624,7 @@ def test_final_norm_mapping_round_trips_bf16_exactly(tiny_hybrid_model: MuseGlim 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() - # Muse final-norm gains sit in [-5, -3); a bf16 ``w - 1`` round trip loses a bit in this range. + # 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)