bridge.peft.lora#

Module Contents#

Classes#

LoRA

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

LoRAMerge

Implements the LoRA weight merge for parameter-efficient fine-tuning.

Data#

logger

API#

bridge.peft.lora.logger#

‘getLogger(…)’

class bridge.peft.lora.LoRA#

Bases: megatron.bridge.peft.base.PEFT, megatron.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.

  • exclude_modules (List[str], optional) – A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name does not match any string in exclude_modules. If used, will require target_modules to be empty list or None.

  • dim (int) – Dimension of the low-rank projection space. Defaults to 32.

  • alpha (int) – Weighting factor for the low-rank projection. Defaults to 32.

  • dropout (float) – Dropout rate for the low-rank projection. Defaults to 0.0.

  • dropout_position (Literal['pre', 'post'], optional) – Position for applying dropout. Can be ‘pre’ (before the low-rank projection) or ‘post’ (after). Defaults to ‘pre’.

  • a2a_experimental (bool) – Enables the experimental All-to-All (A2A) communication strategy. Defaults to False.

  • lora_A_init_method (str) – Initialization method for the low-rank matrix A. Defaults to “xavier”.

  • lora_B_init_method (str) – Initialization method for the low-rank matrix B. Defaults to “zero”.

  • lora_dtype (torch.dtype) – Parameter data type for LoRA weights. Default None (will use model’s dtype).

target_modules: List[str]#

‘field(…)’

dim: int#

32

alpha: int#

32

dropout: float#

0.0

dropout_position: Literal[pre, post]#

‘pre’

lora_A_init_method: str#

‘xavier’

lora_B_init_method: str#

‘zero’

a2a_experimental: bool#

False

lora_dtype: torch.dtype#

None

transform(
module: torch.nn.Module,
name: Optional[str] = None,
prefix: Optional[str] = None,
) → torch.nn.Module#

Applies LoRA to a specific module within the model architecture.

Parameters:

  • m (nn.Module) – The module to apply LoRA to.

  • name (str, optional) – Name of the module (if applicable). Defaults to None.

  • prefix (str, optional) – Prefix for the module name (if applicable). Defaults to None.

Returns:

The modified module with LoRA applied, or the original module if not a target.

Return type:

nn.Module

class bridge.peft.lora.LoRAMerge#

Bases: megatron.bridge.peft.base.PEFT

Implements the LoRA weight merge for parameter-efficient fine-tuning.

transform(
module: torch.nn.Module,
name: Optional[str] = None,
prefix: Optional[str] = None,
) → torch.nn.Module#

Merges the LoRA adapter with the base model weights.

Parameters:

  • m (nn.Module) – The module to apply LoRA merge to.

  • name (str, optional) – Name of the module to merge. Defaults to None.

  • prefix (str, optional) – Prefix for the module name. Defaults to None.

Returns:

The modified module with the LoRA adapter merged into the base model weights.

Return type:

nn.Module