Fix final-layer LoRA target matching

This commit is contained in:
Xin Wang
2026-08-24 15:38:45 +08:00
parent 7fdda297d3
commit aa4278e13e
2 changed files with 37 additions and 17 deletions
+4 -2
View File
@@ -1083,8 +1083,10 @@ The first follow-up to the sparse-LoRA regression should remove sparse
supervision as a variable and constrain adaptation near the Thinker output.
Use dense 5-Hz caches, reserve `Session_12` as development data, and train on
`Session_4` through `Session_11`. `--lora-last-n-layers 8` discovers the
language-layer indices and stores PEFT `layers_to_transform=[28,...,35]` for
the 36-layer Qwen2.5-Omni-3B Thinker. Combining it with
language-layer indices and passes the 24 selected full module paths directly
to PEFT for the 36-layer Qwen2.5-Omni-3B Thinker. Explicit paths are required
because the extracted Thinker exposes names beginning with `layers.0...`,
which some PEFT `layers_pattern` versions fail to recognize. Combining it with
`--lora-target-modules q_proj,v_proj,o_proj` reduces the adapter from all 252
attention/MLP projections to 24 output-side attention projections.
+33 -15
View File
@@ -809,7 +809,7 @@ def _prepare_language_lora(
if (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
}
)
layers_to_transform: list[int] | None = None
selected_layer_indices: list[int] | None = None
matched = matched_all
if last_n_layers:
if not layer_indices:
@@ -821,8 +821,8 @@ def _prepare_language_lora(
f"--lora-last-n-layers={last_n_layers} exceeds the "
f"{len(layer_indices)} matched Thinker layers"
)
layers_to_transform = layer_indices[-last_n_layers:]
selected = set(layers_to_transform)
selected_layer_indices = layer_indices[-last_n_layers:]
selected = set(selected_layer_indices)
matched = [
name
for name in matched_all
@@ -835,21 +835,39 @@ def _prepare_language_lora(
raise RuntimeError(
"Final-layer selection removed every matched LoRA module"
)
config_kwargs: dict[str, Any] = {}
if layers_to_transform is not None:
config_kwargs.update(
layers_to_transform=layers_to_transform,
layers_pattern="layers",
)
# PEFT's layers_pattern matcher requires additional nesting before the
# ModuleList name in some versions. The extracted Omni Thinker exposes
# modules directly as ``layers.28...``, so layers_to_transform can reject
# every otherwise valid target. Explicit full module paths use PEFT's
# exact/suffix-list matching and are stable for this submodel hierarchy.
peft_target_modules = matched if selected_layer_indices is not None else list(
target_modules
)
config = LoraConfig(
r=rank,
lora_alpha=alpha,
lora_dropout=dropout,
target_modules=list(target_modules),
target_modules=peft_target_modules,
bias="none",
**config_kwargs,
)
language_model = get_peft_model(language_model, config)
targeted_module_names = getattr(
getattr(language_model, "base_model", None), "targeted_module_names", None
)
if targeted_module_names is not None:
actual = list(targeted_module_names)
missing = [
expected
for expected in matched
if not any(
name == expected or name.endswith(f".{expected}") for name in actual
)
]
if missing:
raise RuntimeError(
"PEFT did not inject all selected LoRA modules; first missing="
f"{missing[0]}, missing={len(missing)}, targeted={len(actual)}"
)
language_model.config.use_cache = False
if gradient_checkpointing:
try:
@@ -861,14 +879,14 @@ def _prepare_language_lora(
if verbose:
language_model.print_trainable_parameters()
print(
f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
f"LoRA matched {len(matched)} language linear modules; first={matched[0]}, "
f"last={matched[-1]}"
)
if layers_to_transform is not None:
if selected_layer_indices is not None:
print(
"LoRA layer restriction: "
f"{layers_to_transform[0]}--{layers_to_transform[-1]} "
f"({len(layers_to_transform)} of {len(layer_indices)} layers)"
f"{selected_layer_indices[0]}--{selected_layer_indices[-1]} "
f"({len(selected_layer_indices)} of {len(layer_indices)} layers)"
)
return language_model