allow registration of custom checkpoint conversion mappings (#42634)

* allow registration of custom checkpoint conversion mappings * add tests * chore: lint * move tests to test_core_model_loading.py * fixup --------- Co-authored-by: Arthur <arthur.zucker@gmail.com>

Author: winglian

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/utils/test_core_model_loading.py

@@ -18,6 +18,7 @@
 import torch.nn as nn
 
 from transformers import PretrainedConfig
+from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping
 from transformers.core_model_loading import (
  Chunk,
  Concatenate,
@@ -505,5 +506,43 @@ def __init__(self):
  torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2)
 
 
+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)
+
+
 if __name__ == "__main__":
  unittest.main()