From e7decacdaf0d9c988e8fabd0f8e21dd9f48d4192 Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 6 Jun 2025 04:59:29 +0000 Subject: [PATCH] fix compiled checkpoint loading --- src/ctx_to_lora/modeling/hypernet.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 7adf258..5187694 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -664,8 +664,16 @@ class ModulatedPretrainedModel(nn.Module): f"but the loaded name is: {self.base_model_name_or_path}" ) self._init_model() + + def remove_compile_prefix(sd: dict[str, Tensor]) -> dict[str, Tensor]: + COMPILED_PREFIX = "_orig_mod." + for k in list(sd.keys()): + if k.startswith(COMPILED_PREFIX): + sd[k[len(COMPILED_PREFIX) :]] = sd.pop(k) + return sd + load_result = self.hypernet.load_state_dict( - state_dict, strict=False, *args, **kwargs + remove_compile_prefix(state_dict), strict=True, *args, **kwargs ) logger.info(f"load result: {load_result}") return load_result