From d09d79a7de9502b49f3e98d72558b41973f73e47 Mon Sep 17 00:00:00 2001 From: Taksh Date: Mon, 20 Apr 2026 15:20:19 +0530 Subject: [PATCH] Fix bias='lora_only' state dict collection in get_peft_state_maybe_zero_3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lora_only branch of get_peft_state_maybe_zero_3 has three related bugs in the same loop: 1. `for k, t in maybe_lora_bias:` iterates dictionary keys (strings), not (key, value) pairs. Unpacking a string into `k, t` raises ValueError as soon as the dict is non-empty. 2. `if bias_name in lora_bias_names` reuses `bias_name` leaked from the previous loop — always the last lora-derived bias name — so the condition has nothing to do with the current key `k`. 3. `to_return[bias_name] = t` writes to that same leaked name instead of `k`, which means the bias param keyed by `k` is stored under the wrong name (or overwritten). Mirror the peft reference (peft/utils/save_and_load.py get_peft_model_state_dict): iterate `maybe_lora_bias.items()` and match against `k` directly. --- fastchat/train/train_lora.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/fastchat/train/train_lora.py b/fastchat/train/train_lora.py index 9ecb47c29..61b0437e6 100644 --- a/fastchat/train/train_lora.py +++ b/fastchat/train/train_lora.py @@ -92,9 +92,9 @@ def get_peft_state_maybe_zero_3(named_params, bias): lora_bias_names.add(bias_name) elif "bias" in k: maybe_lora_bias[k] = t - for k, t in maybe_lora_bias: - if bias_name in lora_bias_names: - to_return[bias_name] = t + for k, t in maybe_lora_bias.items(): + if k in lora_bias_names: + to_return[k] = t else: raise NotImplementedError to_return = {k: maybe_zero_3(v) for k, v in to_return.items()}