chore(beep boop 🤖): Bump uv.lock (main, mcore-main) (2026-04-18)

.main.commit

@@ -1 +1 @@
-23265d2db37efa883b3983bce321c631898822f3
+997896883b04dee1259d92889e7344be528eefa8

src/megatron/bridge/diffusion/models/flux/flux_provider.py

@@ -69,6 +69,7 @@ class FluxProvider(TransformerConfig, ModelProviderMixin[VisionModule]):
     attention_dropout: float = 0
 
     # FLUX-specific layer configuration
+    qk_layernorm: bool = True
     num_joint_layers: int = 19
     num_single_layers: int = 38
 

src/megatron/bridge/diffusion/models/wan/wan_provider.py

@@ -33,6 +33,7 @@
 @dataclass
 class WanModelProvider(TransformerConfig, ModelProviderMixin[VisionModule]):  # noqa: D101
     crossattn_emb_size: int = 1536  # cross attention emebedding size after linear projection
+    qk_layernorm: bool = True
     add_bias_linear: bool = True
     gated_linear_unit: bool = False
 

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

@@ -68,6 +68,7 @@ class Gemma3ModelProvider(GPTModelProvider):
     layernorm_epsilon: float = 1e-6
 
     # attention
+    qk_layernorm: bool = True
     window_size: tuple = 512  # local
     interleaved_attn_pattern: tuple = (5, 1)  # (local, global)
     attention_dropout: float = 0.0
@@ -217,8 +218,8 @@ def gemma3_layer_spec(config) -> ModuleSpec:
                 submodules=SelfAttentionSubmodules(
                     linear_qkv=TELayerNormColumnParallelLinear,
                     core_attention=Gemma3TEDotProductAttention,  # mixed gloabl/local attn
-                    q_layernorm=TENorm,
-                    k_layernorm=TENorm,
+                    q_layernorm=TENorm if config.qk_layernorm else None,
+                    k_layernorm=TENorm if config.qk_layernorm else None,
                     linear_proj=TERowParallelLinearLayerNorm,  # post attn RMSNorm
                 ),
             ),

src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py

@@ -150,9 +150,9 @@ def provider_bridge(self, hf_pretrained):
             provider.rotary_percent = rotary_dim / head_dim
 
         # Full-dimension QK norm via custom layer spec (see minimax_m2_provider.py).
-        # qk_layernorm stays False to avoid the default per-head TENorm; our custom
-        # spec injects FullDimQNorm/FullDimKNorm directly into SelfAttention.
-        provider.qk_layernorm = False
+        # qk_layernorm=True so mcore creates QK norms; the spec overrides the default
+        # TENorm with FullDimQNorm/FullDimKNorm for full-dimension normalization.
+        provider.qk_layernorm = True
         provider.transformer_layer_spec = minimax_m2_layer_spec
 
         # MoE settings — sigmoid routing with expert bias (same pattern as DeepSeek V3)

tests/unit_tests/models/minimax_m2/test_minimax_m2_bridges.py

@@ -132,8 +132,8 @@ def test_provider_bridge_sets_custom_layer_spec(self, mock_pretrained):
         from megatron.bridge.models.minimax_m2.minimax_m2_provider import minimax_m2_layer_spec
 
         assert provider.transformer_layer_spec is minimax_m2_layer_spec
-        # qk_layernorm is disabled at the provider level; the custom spec injects norms instead.
-        assert provider.qk_layernorm is False
+        # qk_layernorm must be True so mcore accepts the spec's FullDimQNorm/FullDimKNorm overrides.
+        assert provider.qk_layernorm is True
 
     def test_provider_bridge_dtype_bfloat16(self, mock_pretrained):
         bridge = MiniMaxM2Bridge()