Fix final-layer LoRA target matching

train_omni_turn_lora.py

@@ -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