fix(gemma4_moe): re-tie lm_head to active embed_tokens on MoE path

The MoE path replaces language_model after HF __init__, orphaning the lm_head<->embed_tokens tie that HF set up. Re-tie lm_head to the active embed_tokens when tie_word_embeddings is set (Gemma defaults True). Add CPU tied/untied tests.

Refs #2512

Author: Achyuthan-S

nemo_automodel/components/models/gemma4_moe/model.py

@@ -809,6 +809,15 @@ def __init__(
         # Expose moe_config for the MoE parallelizer assertion
         self.model.moe_config = self.model.language_model.moe_config
 
+        # HF's super().__init__() tied lm_head.weight to the *original* text
+        # embed_tokens, but the language_model replacement above swapped in a
+        # fresh embed_tokens and orphaned that alias. Re-tie lm_head to the
+        # now-active embedding when the config requests tied embeddings (Gemma
+        # defaults to tie_word_embeddings=True). The shared Parameter survives
+        # the in-place cast in initialize_weights().
+        if getattr(text_config, "tie_word_embeddings", False):
+            self.lm_head.weight = self.model.language_model.embed_tokens.weight
+
         self.vocab_size = text_config.vocab_size
         # State dict adapter for HF ↔ NeMo weight conversion
         if self.backend.enable_hf_state_dict_adapter:

tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Weight-tying tests for the Gemma4 MoE conditional-generation model.
+
+HF's ``super().__init__()`` ties ``lm_head.weight`` to the original text
+``embed_tokens``. The MoE path then replaces ``language_model`` with
+``Gemma4MoETextModelBackend`` (a fresh ``embed_tokens``), which orphans that
+alias. The model re-ties ``lm_head`` to the now-active embedding when
+``tie_word_embeddings`` is set (Gemma defaults to ``True``); these tests pin
+that behavior for both the tied and untied configs.
+
+Runs on CPU (no CUDA / TE / DeepEP required).
+"""
+
+import torch
+from transformers.models.gemma4.configuration_gemma4 import Gemma4Config, Gemma4TextConfig
+
+from nemo_automodel.components.models.common import BackendConfig
+from nemo_automodel.components.models.gemma4_moe.model import (
+    Gemma4ForConditionalGeneration,
+    Gemma4MoETextModelBackend,
+)
+
+
+def _make_text_config(**overrides):
+    """Tiny Gemma4TextConfig (2 layers, small hidden, tiny vocab, few experts)."""
+    defaults = dict(
+        vocab_size=256,
+        hidden_size=64,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_hidden_layers=2,
+        intermediate_size=128,
+        rms_norm_eps=1e-6,
+        max_position_embeddings=256,
+        enable_moe_block=True,  # routes construction through the NeMo MoE backend
+        num_experts=4,
+        top_k_experts=2,
+        moe_intermediate_size=64,
+        layer_types=["full_attention", "sliding_attention"],
+        sliding_window=128,
+        hidden_activation="gelu_pytorch_tanh",
+        torch_dtype="bfloat16",
+    )
+    defaults.update(overrides)
+    return Gemma4TextConfig(**defaults)
+
+
+def _make_cpu_backend():
+    """CPU-friendly backend: no TE, no DeepEP, plain torch kernels."""
+    return BackendConfig(
+        linear="torch",
+        attn="sdpa",
+        rms_norm="torch",
+        experts="torch",
+        dispatcher="torch",
+        fake_balanced_gate=False,
+        enable_hf_state_dict_adapter=False,
+    )
+
+
+def _build(tie_word_embeddings: bool) -> Gemma4ForConditionalGeneration:
+    config = Gemma4Config(text_config=_make_text_config(tie_word_embeddings=tie_word_embeddings))
+    model = Gemma4ForConditionalGeneration(config, backend=_make_cpu_backend())
+    # Sanity: construction routed through the real NeMo MoE backend (the path
+    # that replaces language_model and breaks HF's tie).
+    assert isinstance(model.model.language_model, Gemma4MoETextModelBackend)
+    return model
+
+
+def test_tied_lm_head_shares_active_embedding_after_construction():
+    """tie_word_embeddings=True: lm_head must alias the *active* MoE embed_tokens."""
+    model = _build(tie_word_embeddings=True)
+    assert model.lm_head.weight is model.model.language_model.embed_tokens.weight
+
+
+def test_tied_lm_head_survives_initialize_weights():
+    """The tie set in __init__ must survive the bf16 cast in initialize_weights()."""
+    model = _build(tie_word_embeddings=True)
+    model.initialize_weights(dtype=torch.bfloat16, buffer_device=torch.device("cpu"))
+
+    embed = model.model.language_model.embed_tokens.weight
+    lm_head = model.lm_head.weight
+    assert lm_head is embed
+    assert lm_head.dtype == torch.bfloat16
+
+
+def test_untied_lm_head_is_separate():
+    """tie_word_embeddings=False: lm_head must keep its own storage."""
+    model = _build(tie_word_embeddings=False)
+    assert model.lm_head.weight is not model.model.language_model.embed_tokens.weight
+    assert model.lm_head.weight.data_ptr() != model.model.language_model.embed_tokens.weight.data_ptr()