`bridge.peft.canonical_lora`#

Module Contents#

Classes#

`ModuleDict`	nn.ModuleDict with a sharded_state_dict implementation for checkpointing
`LoRALinearSplitQKV`	An adapter wrapper for `linear_qkv` where q, k, v are three separate adapters. This module that adds the output of the adapters to the output of the wrapped module while taking care of shape.
`LoRALinearSplitFC1UpGate`	An adapter wrapper for `linear_fc1` where up_proj and gate_proj are two separate adapters. This module that adds the output of the adapters to the output of the wrapped module while taking care of shape.
`CanonicalLoRA`	Implements the LoRA (Low-Rank Adaptation) module for parameter-efficient fine-tuning. Canonical LoRA applies LoRA on Q, K, V projection matrices separately, as well as Up and Gate projection matrices separately. This follows more closely with Huggingface’s implementation of LoRA.

Data#

logger

API#

bridge.peft.canonical_lora.logger#: ‘getLogger(…)’

class bridge.peft.canonical_lora.ModuleDict#

Bases: torch.nn.ModuleDict

nn.ModuleDict with a sharded_state_dict implementation for checkpointing

sharded_state_dict( prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None, ) → megatron.core.dist_checkpointing.mapping.ShardedStateDict#

Retrieve the sharded state dictionary of the wrapped module and adapter.

This method is used for distributed checkpointing, combining the sharded states of both the main module and the adapter.

Parameters:

prefix (str) – A prefix added to parameter and buffer names. Defaults to ‘’.
sharded_offsets (Tuple[Tuple[int, int, int]]) – Offsets for sharded parameters. Defaults to an empty tuple.
metadata (Optional[dict]) – Additional metadata for the sharded state. Defaults to None.

Returns:

The combined sharded state dictionary.

Return type:

ShardedStateDict

class bridge.peft.canonical_lora.LoRALinearSplitQKV#

Bases: megatron.bridge.peft.adapter_wrapper.AdapterWrapper

An adapter wrapper for linear_qkv where q, k, v are three separate adapters. This module that adds the output of the adapters to the output of the wrapped module while taking care of shape.

This class is designed to be used with LoRA (Low-Rank Adaptation) and similar techniques where the adapter’s output is added to the main module’s output. It extends the AdapterWrapper class to provide a specific implementation of the forward method.

forward( x: torch.Tensor, ) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

class bridge.peft.canonical_lora.LoRALinearSplitFC1UpGate#

Bases: megatron.bridge.peft.adapter_wrapper.AdapterWrapper

An adapter wrapper for linear_fc1 where up_proj and gate_proj are two separate adapters. This module that adds the output of the adapters to the output of the wrapped module while taking care of shape.

forward( x: torch.Tensor, ) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

class bridge.peft.canonical_lora.CanonicalLoRA#

Bases: megatron.bridge.peft.base.PEFT, megatron.bridge.peft.module_matcher.ModuleMatcher

Implements the LoRA (Low-Rank Adaptation) module for parameter-efficient fine-tuning. Canonical LoRA applies LoRA on Q, K, V projection matrices separately, as well as Up and Gate projection matrices separately. This follows more closely with Huggingface’s implementation of LoRA.

Parameters:

target_modules (List[str], optional) – A list of module names to apply LoRA to. Defaults to all linear layers [‘linear_q’, ‘linear_k’, ‘linear_v’, ‘linear_proj’, ‘linear_fc1_up’, ‘linear_fc1_gate’, ‘linear_fc2’]. - ‘linear_q’, ‘linear_k’, ‘linear_v’: Apply LoRA to the linear layer used for query, key, and value projections in self-attention. This is fused into one matrix in LoRA, but left as three separate matrices in Canonical LoRA. - ‘linear_proj’: Apply LoRA to the linear layer used for projecting the output of self-attention. - ‘linear_fc1_up’, ‘linear_fc1_gate’: Apply LoRA to the Up proj and Gate proj layers. These two together constitute the first fully-connected layer in MLP in LoRA. - ‘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_q’, ‘.layers.1..linear_q’] to add LoRA to only linear_q 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’.
lora_A_init_method (str) – Initialization method for LoRA A matrix. Defaults to “xavier”.
lora_B_init_method (str) – Initialization method for LoRA B matrix. Defaults to “zero”.

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’

__post_init__() → None#

Initialize the canonical mapping and call the parent post_init.

Construct a mapping from the target module as supported in LoRA() to the specific parts of the layer for which adapter is applied.

For example, if user specifies target_module = [‘linear_q’, ‘linear_k’, ‘linear_proj’, ‘linear_fc1_up’], then canonical_lora_mapping = { “linear_qkv”: {‘linear_q’, ‘linear_k’}, “linear_proj”: {‘linear_proj’}, # the value of this key does not matter “linear_fc1”: {‘linear_fc1_up’}, }

If user specifies target_module = [’.layers.0..linear_q’, ‘.layers.1..linear_q’], then canonical_lora_mapping = { “’.layers.0..linear_qkv’”: {‘linear_q’}, “’.layers.1..linear_qkv’”: {‘linear_q’}, }

transform( m: 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 (Optional[str]) – Name of the module (if applicable). Defaults to None.
prefix (Optional[str]) – 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

bridge.peft.canonical_lora#

Module Contents#

Classes#

Data#

API#

`bridge.peft.canonical_lora`#