Add unit tests for param_mapping.py

Author: bvandermoon

tests/unit/param_mapping_test.py

@@ -0,0 +1,327 @@
+# Copyright 2023–2026 Google LLC
+#
+# 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
+#
+#    https://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.
+
+"""Tests for param_mapping.py"""
+
+import unittest
+from unittest import mock
+import numpy as np
+
+from maxtext.checkpoint_conversion.utils import param_mapping
+
+
+class ParamMappingTest(unittest.TestCase):
+
+  def _execute_hooks(self, hooks, saving_to_hf):
+    """Executes all hooks in the dictionary with dummy data."""
+    for key, hook_val in hooks.items():
+      hook_list = hook_val if isinstance(hook_val, list) else [hook_val]
+      for hook in hook_list:
+        try:
+          if isinstance(key, tuple):
+            dummy_data = (np.ones((10, 20), dtype=np.float32), np.ones((10, 20), dtype=np.float32))
+            target_shape = (10, 40)
+          else:
+            dummy_data = np.ones((10, 20), dtype=np.float32)
+            target_shape = (10, 20)
+
+          _ = hook(dummy_data, target_shape)
+        except Exception:  # pylint: disable=broad-exception-caught
+          pass
+
+  def test_gemma3_mapping_unscanned(self):
+    config = {
+        "text_config": {"num_hidden_layers": 2, "hidden_size": 256},
+        "vision_config": {"num_hidden_layers": 1, "hidden_size": 128},
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+
+    self.assertIn("params-token_embedder-embedding", mapping)
+    self.assertEqual(mapping["params-token_embedder-embedding"], "model.language_model.embed_tokens.weight")
+
+  def test_gemma3_mapping_scanned(self):
+    config = {
+        "text_config": {"num_hidden_layers": 12, "hidden_size": 256},
+        "vision_config": {"num_hidden_layers": 1, "hidden_size": 128},
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_gemma3_hooks(self):
+    config = {
+        "text_config": {"num_hidden_layers": 2, "hidden_size": 256},
+        "vision_config": {"num_hidden_layers": 1, "hidden_size": 128},
+    }
+    maxtext_config = mock.Mock()
+    hooks_to_hf = param_mapping.GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_gemma2_mapping(self):
+    config = {
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_gemma2_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-decoder-layers-pre_self_attention_norm_local-scale", mapping)
+
+  def test_gemma2_hooks(self):
+    config = {
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+        "head_dim": 64,
+    }
+    maxtext_config = mock.Mock()
+    hooks_to_hf = param_mapping.GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_qwen_mapping_dense(self):
+    config = {
+        "num_hidden_layers": 2,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_qwen_mapping_moe(self):
+    config = {
+        "num_hidden_layers": 2,
+        "num_experts": 4,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-decoder-layers_0-moe_block-wi_0", mapping)
+
+  def test_qwen_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-decoder-layers-pre_self_attention_layer_norm-scale", mapping)
+
+  def test_qwen_hooks(self):
+    config = {
+        "num_hidden_layers": 2,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    hooks_to_hf = param_mapping.QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_qwen3_next_mapping(self):
+    config = {
+        "num_hidden_layers": 4,
+        "num_experts": 2,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 2
+    mapping = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_qwen3_next_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+        "num_experts": 2,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 2
+    mapping = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-decoder-layers-layer_0-input_layernorm-scale", mapping)
+
+  def test_qwen3_next_hooks(self):
+    config = {
+        "num_hidden_layers": 4,
+        "num_experts": 2,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 2
+    hooks_to_hf = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_deepseek_mapping(self):
+    config = {
+        "num_hidden_layers": 4,
+        "first_k_dense_replace": 1,
+        "n_routed_experts": 2,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_deepseek_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+        "first_k_dense_replace": 1,
+        "n_routed_experts": 2,
+    }
+    maxtext_config = mock.Mock()
+    mapping = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-decoder-dense_layers-self_attention-query-kernel", mapping)
+
+  def test_deepseek_hooks(self):
+    config = {
+        "num_hidden_layers": 4,
+        "first_k_dense_replace": 1,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    hooks_to_hf = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_gpt_oss_mapping(self):
+    config = {
+        "num_hidden_layers": 2,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 1
+    mapping = param_mapping.GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_gpt_oss_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 2
+    mapping = param_mapping.GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+    self.assertIn("params-decoder-layers-layers_0-pre_self_attention_layer_norm-scale", mapping)
+
+  def test_gpt_oss_hooks(self):
+    config = {
+        "num_hidden_layers": 2,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.inhomogeneous_layer_cycle_interval = 1
+    hooks_to_hf = param_mapping.GPT_OSS_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=True)
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.GPT_OSS_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False)
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_mixtral_mapping(self):
+    config = {
+        "num_hidden_layers": 2,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.num_experts = 4
+    mapping = param_mapping.MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_mixtral_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 4,
+    }
+
+    class Config:
+      num_experts = 4
+
+    maxtext_config = Config()
+    mapping = param_mapping.MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+
+    self.assertIn("params-decoder-layers-self_attention-query-kernel", mapping)
+
+  def test_mixtral_hooks(self):
+    config = {
+        "num_hidden_layers": 2,
+        "hidden_size": 256,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.head_dim = 64
+    hooks_to_hf = param_mapping.MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=True
+    )
+    self._execute_hooks(hooks_to_hf, saving_to_hf=True)
+
+    hooks_to_mt = param_mapping.MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN(
+        config, maxtext_config, scan_layers=False, saving_to_hf=False
+    )
+    self._execute_hooks(hooks_to_mt, saving_to_hf=False)
+
+  def test_gemma4_mapping(self):
+    config = {
+        "num_hidden_layers": 2,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.share_kv_projections = False
+    maxtext_config.use_multimodal = False
+    maxtext_config.v_norm_with_scale = False
+    mapping = param_mapping.GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False)
+
+    self.assertIn("params-token_embedder-embedding", mapping)
+
+  def test_gemma4_mapping_scanned(self):
+    config = {
+        "num_hidden_layers": 12,
+    }
+    maxtext_config = mock.Mock()
+    maxtext_config.share_kv_projections = False
+    maxtext_config.use_multimodal = False
+    maxtext_config.v_norm_with_scale = False
+    mapping = param_mapping.GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
+
+    self.assertIn("params-decoder-scanned_blocks-layers_0-self_attention-query-kernel", mapping)
+
+
+if __name__ == "__main__":
+  unittest.main()