From 8bcc310e95576f82037a366cc4b85bcbad50137a Mon Sep 17 00:00:00 2001 From: MarkovChain-why Date: Wed, 11 Mar 2026 02:11:16 +0800 Subject: [PATCH] Fix: inject LoRA adapters before loading LoRA weights in from_pretrained MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When `defer_lora_injection=True`, LoRA adapters are not created during model init — they are deferred until after base weights are loaded so that the pretrained weight key paths match the model's state dict. However, `VLA.from_pretrained` (the inference/deployment path) skips adapter injection when `lora_weights_path` is provided: it jumps straight to `load_lora_weight()` without first calling `inject_lora_after_loading()`. This means the LoRA weight keys (lora_A, lora_B) have no matching parameters in the model, so the loaded weights are silently dropped. The training path in `base.py` does not have this bug — it correctly calls `inject_lora_after_loading()` after loading pretrained weights (line 723-726). Fix: call `inject_lora_after_loading()` before `load_lora_weight()` when a LoRA weights path is provided. Co-Authored-By: Claude Opus 4.6 --- groot/vla/model/dreamzero/base_vla.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/groot/vla/model/dreamzero/base_vla.py b/groot/vla/model/dreamzero/base_vla.py index 274b8640..7784f9b8 100644 --- a/groot/vla/model/dreamzero/base_vla.py +++ b/groot/vla/model/dreamzero/base_vla.py @@ -324,6 +324,10 @@ def from_pretrained_for_tuning( ) if lora_weights_path is not None: + # Inject LoRA adapters first (base weights are already loaded with correct key paths) + if hasattr(model, 'action_head') and hasattr(model.action_head, 'inject_lora_after_loading'): + print("Injecting LoRA adapters before loading LoRA weights") + model.action_head.inject_lora_after_loading() print(f"Loading LoRA weights from: {lora_weights_path}") model.load_lora_weight(lora_weights_path) else: