|
43 | 43 | is_expert_linear, |
44 | 44 | is_grouped_expert_linear, |
45 | 45 | is_modelopt_linear, |
| 46 | + wildcard_match, |
46 | 47 | ) |
47 | 48 |
|
48 | 49 |
|
@@ -159,6 +160,63 @@ class LoRA(PEFT, ModuleMatcher): |
159 | 160 | experts_shared_outer_loras: bool = False |
160 | 161 | use_transformer_engine_op_fuser: bool = False |
161 | 162 |
|
| 163 | + def __call__(self, model, training: bool = True): |
| 164 | + """Apply LoRA, promoting targeted parameter-backed projections first. |
| 165 | +
|
| 166 | + Parameter-backed projections use the same explicit ``target_modules`` |
| 167 | + patterns as ordinary modules. Turning such a parameter into the regular |
| 168 | + LoRA module only when it is targeted lets the existing LoRA wrapper and |
| 169 | + adapter export path handle it without changing the base model's |
| 170 | + checkpoint layout for non-LoRA runs. |
| 171 | + """ |
| 172 | + |
| 173 | + self._promote_parameter_targets(model) |
| 174 | + return super().__call__(model, training=training) |
| 175 | + |
| 176 | + def _promote_parameter_targets(self, model) -> None: |
| 177 | + """Expose supported bare parameters as linear modules before walking. |
| 178 | +
|
| 179 | + The promoted module keeps the original ``nn.Parameter`` object as its |
| 180 | + weight, so initialization, freezing, and optimizer ownership remain |
| 181 | + unchanged. Parameters that do not match ``target_modules`` retain the |
| 182 | + existing PEFT behavior. |
| 183 | + """ |
| 184 | + |
| 185 | + target_patterns = self.target_modules |
| 186 | + if isinstance(target_patterns, str): |
| 187 | + target_patterns = (target_patterns,) |
| 188 | + else: |
| 189 | + target_patterns = tuple(target_patterns or ()) |
| 190 | + if not target_patterns: |
| 191 | + return |
| 192 | + |
| 193 | + models = model if isinstance(model, list) else [model] |
| 194 | + for root in models: |
| 195 | + if isinstance(root, torch.nn.parallel.DistributedDataParallel): |
| 196 | + root = root.module |
| 197 | + for prefix, parent in list(root.named_modules()): |
| 198 | + for parameter_name, parameter in list(parent.named_parameters(recurse=False)): |
| 199 | + full_name = f"{prefix}.{parameter_name}" if prefix else parameter_name |
| 200 | + if not any( |
| 201 | + parameter_name == pattern or wildcard_match(pattern, full_name) for pattern in target_patterns |
| 202 | + ): |
| 203 | + continue |
| 204 | + if parameter.ndim != 2: |
| 205 | + raise ValueError( |
| 206 | + f"{full_name} must be a 2-D matrix for LoRA, got shape {tuple(parameter.shape)}" |
| 207 | + ) |
| 208 | + promoted = nn.Linear( |
| 209 | + parameter.shape[1], |
| 210 | + parameter.shape[0], |
| 211 | + bias=False, |
| 212 | + device=parameter.device, |
| 213 | + dtype=parameter.dtype, |
| 214 | + ) |
| 215 | + promoted.weight = parameter |
| 216 | + delattr(parent, parameter_name) |
| 217 | + setattr(parent, parameter_name, promoted) |
| 218 | + logger.info("Promoted parameter-backed LoRA target: %s", full_name) |
| 219 | + |
162 | 220 | def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None) -> nn.Module: |
163 | 221 | """ |
164 | 222 | Applies LoRA to a specific module within the model architecture. |
|
0 commit comments