bridge.peft.module_matcher#

Module Contents#

Classes#

ModuleMatcher

Implements the LoRA (Low-Rank Adaptation) module for parameter-efficient fine-tuning.

Data#

HAVE_TE

API#

bridge.peft.module_matcher.HAVE_TE#

‘all(…)’

class bridge.peft.module_matcher.ModuleMatcher#

Implements the LoRA (Low-Rank Adaptation) module for parameter-efficient fine-tuning.

LoRA uses a low-rank projection to adapt the weights of a pre-trained model to a new downstream task. This class facilitates the application of LoRA to specific modules within the model architecture.

Parameters:

target_modules (List[str], optional) – A list of module names to apply LoRA to. Defaults to all linear layers [‘linear_qkv’, ‘linear_proj’, ‘linear_fc1’, ‘linear_fc2’]. - ‘linear_qkv’: Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention. - ‘linear_proj’: Apply LoRA to the linear layer used for projecting the output of self-attention. - ‘linear_fc1’: Apply LoRA to the first fully-connected layer in MLP. - ‘linear_fc2’: Apply LoRA to the second fully-connected layer in MLP. Target modules can also contain wildcards. For example, you can specify target_modules=[’.layers.0..linear_qkv’, ‘.layers.1..linear_qkv’] to add LoRA to only linear_qkv on the first two layers.

target_modules: List[str]#

‘field(…)’

exclude_modules: List[str]#

‘field(…)’

canonical_mapping: Dict[str, Set]#

‘field(…)’

match(
m: torch.nn.Module,
name: Optional[str] = None,
prefix: Optional[str] = None,
) → Optional[tuple[str, str]]#

Determines whether a given module matches specified target patterns.

This function checks if the provided module m should be included based on predefined mapping rules (canonical_mapping, target_modules, and exclude_modules). It returns the matching pattern if a match is found; otherwise, it returns None.

Parameters:

  • m (nn.Module) – The module being checked.

  • name (str, optional) – The module’s name.

  • prefix (str, optional) – A prefix to be used in constructing full_name.

Returns:

A tuple containing (matching_pattern, full_name) if a match is found; otherwise, None.

Return type:

Optional[Tuple[str, str]]

Matching Logic:

  1. If canonical_mapping is defined, it checks:

  • Whether name exactly matches a pattern.

  • Whether full_name matches any regex pattern in canonical_mapping.

  1. If target_modules is defined, it follows the same logic as canonical_mapping.

  2. If neither canonical_mapping nor target_modules are defined, it ensures:

  • name is not in exclude_modules.

  • full_name does not match any target_modules patterns.

  • m is an instance of nn.Linear.

Notes:

  • exclude_modules should only be non-empty if neither canonical_mapping nor target_modules are set.

  • The function asserts that exclude_modules is empty when using canonical_mapping or target_modules.