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.

Functions

_should_treat_linear_fc1_as_unfused

Return True when CanonicalLoRA should keep linear_fc1 as a single adapter.

Data

logger

API

bridge.peft.canonical_lora.logger

‘getLogger(…)’

bridge.peft.canonical_lora._should_treat_linear_fc1_as_unfused(full_name: str) → bool

Return True when CanonicalLoRA should keep linear_fc1 as a single adapter.

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.

_interleave_qkv(

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

) → torch.Tensor

Interleave QKV outputs to match Megatron’s packed ordering.

forward(

x: torch.Tensor,

*args: Any,

**kwargs: Any,

) → 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.

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,

*args: Any,

**kwargs: Any,

) → 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”.

• normalize_moe_lora (bool) – When True, expert linear layers use dim // moe_router_topk as the LoRA rank while non-expert layers keep the full dim. This normalizes the total adapter capacity for MoE models so it is comparable to a dense model. Defaults to False.

• share_expert_adapters (bool) – When True, grouped MoE expert linears share one adapter across all local experts on the EP rank. Set to False to create one adapter per local expert instead. Defaults to True.

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’

normalize_moe_lora: bool

False

share_expert_adapters: bool

True

__post_init__() → None

Eagerly build canonical_mapping from the initial target_modules.

canonical_mapping is also re-derived in _init_target_match_state at PEFT.__call__ time so post-construction mutation of target_modules is respected. Building eagerly here additionally surfaces the linear_qkv / linear_fc1 asserts at construction time and supports callers that inspect canonical_mapping before applying the PEFT transform.

_init_target_match_state() → None

Build the canonical mapping and validation aliases from the current target_modules.

This rebuilds both canonical_mapping (used by match to drive transforms) and the alias bookkeeping inherited from ModuleMatcher (used to validate that every requested target matched at least one module). Running here rather than in __post_init__ ensures any post-construction mutation of self.target_modules is reflected.

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