allow registration of custom checkpoint conversion mappings

src/transformers/conversion_mapping.py

@@ -186,10 +186,22 @@ def _build_checkpoint_conversion_mapping():
 
 def get_checkpoint_conversion_mapping(model_type):
     global _checkpoint_conversion_mapping_cache
-    _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping()
+    if _checkpoint_conversion_mapping_cache is None:
+        _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping()
     return deepcopy(_checkpoint_conversion_mapping_cache.get(model_type))
 
 
+def register_checkpoint_conversion_mapping(
+    model_type: str, mapping: list[WeightConverter | WeightRenaming], overwrite: bool = False
+) -> None:
+    global _checkpoint_conversion_mapping_cache
+    if _checkpoint_conversion_mapping_cache is None:
+        _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping()
+    if model_type in _checkpoint_conversion_mapping_cache and not overwrite:
+        raise ValueError(f"Model type {model_type} already exists in the checkpoint conversion mapping.")
+    _checkpoint_conversion_mapping_cache[model_type] = mapping
+
+
 # DO NOT MODIFY, KEPT FOR BC ONLY
 VLMS = [
     "aria",

tests/test_conversion_mapping.py

@@ -0,0 +1,43 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# 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.
+import unittest
+
+from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping
+from transformers.core_model_loading import WeightRenaming
+
+
+class TestConversionMapping(unittest.TestCase):
+    def test_register_checkpoint_conversion_mapping(self):
+        register_checkpoint_conversion_mapping("foobar", [
+            WeightRenaming(".block_sparse_moe.gate", ".mlp.gate"),
+        ])
+        self.assertEqual(len(get_checkpoint_conversion_mapping("foobar")), 1)
+
+    def test_register_checkpoint_conversion_mapping_overwrites(self):
+        register_checkpoint_conversion_mapping("foobarbaz", [
+            WeightRenaming(".block_sparse_moe.gate", ".mlp.gate"),
+        ])
+        with self.assertRaises(ValueError):
+            register_checkpoint_conversion_mapping("foobarbaz", [
+                WeightRenaming(".block_sparse_moe.foo", ".mlp.foo"),
+                WeightRenaming(".block_sparse_moe.bar", ".mlp.bar"),
+            ])
+
+        register_checkpoint_conversion_mapping("foobarbaz", [
+            WeightRenaming(".block_sparse_moe.foo", ".mlp.foo"),
+            WeightRenaming(".block_sparse_moe.bar", ".mlp.bar"),
+        ], overwrite=True)
+
+        self.assertEqual(len(get_checkpoint_conversion_mapping("foobarbaz")), 2)
+