src/megatron/bridge/peft/base.py (1)

99-120: ⚠️ Potential issue | 🟠 Major

Preflight target validation before mutating the model.

Running _validate_target_matches() only after self._walk_model(model, self.transform) makes the error path unsafe: a config like ["linear_qkv", "typo"] will already have frozen/wrapped the valid targets before the ValueError is raised. It also breaks reapplying matcher-based PEFT configs, because CanonicalLoRA.transform() and DoRA.transform() both return early for already wrapped modules before calling match(), so the second pass records no hits and fails validation. A read-only match walk before freeze_model() / self.transform avoids both behaviors.

💡 Suggested approach

     def __call__(self, model: ModelType, training: bool = True) -> ModelType:
         if isinstance(self, ModuleMatcher):
             self._reset_target_match_state()
+
+            def _validate_only(
+                module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None
+            ) -> nn.Module:
+                self.match(module, name, prefix)
+                return module
+
+            self._walk_model(model, _validate_only)
+            self._validate_target_matches()
+            self._reset_target_match_state()
 
         self.freeze_model(model, training=training)
         self._walk_model(model, self.transform)
@@
-        if isinstance(self, ModuleMatcher):
-            self._validate_target_matches()
-
         return model

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/peft/base.py` around lines 99 - 120, Run a read-only
preflight match pass for ModuleMatcher before mutating the model: if
isinstance(self, ModuleMatcher) call a non-mutating walk that invokes the
matcher (i.e., calls match() via the same walk logic) and then call
self._validate_target_matches() immediately after that preflight; only after
validation proceed to call self.freeze_model(model, training=training),
self._walk_model(model, self.transform) and the rest (including
maybe_enable_recompute_inputs_grad and setting training modes). This ensures
CanonicalLoRA.transform and DoRA.transform won’t short-circuit earlier matches
and prevents partial mutation before validation.