model: support Pi0.5 (#30633)

2026-07-11 07:50:58 +08:00
parent e8646701c1
commit 649ce5dd3d
61 changed files with 8593 additions and 108 deletions
@@ -179,14 +179,20 @@ class skip_init_modules:
    def __enter__(self):
        # Save originals
        self._orig_reset = {}
-        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d):
+        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Embedding):
            self._orig_reset[cls] = cls.reset_parameters
            cls.reset_parameters = lambda self: None  # skip init
+        from transformers.modeling_utils import PreTrainedModel
+
+        self._pretrained_model_cls = PreTrainedModel
+        self._orig_post_init = PreTrainedModel.post_init
+        PreTrainedModel.post_init = lambda self: None

    def __exit__(self, exc_type, exc_value, traceback):
        # restore originals
        for cls, orig in self._orig_reset.items():
            cls.reset_parameters = orig
+        self._pretrained_model_cls.post_init = self._orig_post_init


 def _normalize_component_type(module_type: str) -> str: