`bridge.peft.lora_layers`#

Module Contents#

Classes#

`LoRALinear`	An adapter wrapper that adds the output of the adapter to the output of the wrapped module.
`LoRATopKRouter`	Adapter wrapper that applies LoRA to router gating logits.
`TEFusedLoRALinear`	LoRA adapter wrapper using Transformer Engine operation fuser
`LinearAdapter`	Linear with LoRA that preserves the base weight and bias checkpoint keys.

API#

class bridge.peft.lora_layers.LoRALinear#

Bases: megatron.bridge.peft.adapter_wrapper.AdapterWrapper

An adapter wrapper that adds the output of the adapter to the output of the wrapped module.

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.

property weight: torch.Tensor#: Return the effective base weight, including the LoRA delta when enabled.

property bias: torch.Tensor | None#: Return the wrapped linear bias.

forward(

x: torch.Tensor,

*args: Any,

**kwargs: Any,

) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

Forward pass that combines the wrapped module output with the adapter output.

Parameters:

x – Input tensor.
*args – Additional positional arguments for the wrapped module.
**kwargs – Additional keyword arguments for the wrapped module.

Returns:

 - Combined output (linear_output + adapter_output) if adapter is enabled,
   otherwise just the linear_output
 - Bias term (if present, otherwise None)

Return type:

A tuple containing

class bridge.peft.lora_layers.LoRATopKRouter#

Bases: megatron.bridge.peft.adapter_wrapper.AdapterWrapper

Adapter wrapper that applies LoRA to router gating logits.

forward(x: torch.Tensor, *args: Any, **kwargs: Any)#: Forward pass that adds LoRA delta to router logits before routing.

class bridge.peft.lora_layers.TEFusedLoRALinear(to_wrap: torch.nn.Module, adapter: torch.nn.Module)#

Bases: bridge.peft.lora_layers.LoRALinear

LoRA adapter wrapper using Transformer Engine operation fuser

Initialization

_make_fused_branches() → tuple[transformer_engine.pytorch.ops.Sequential, transformer_engine.pytorch.ops.Sequential]#: Construct fused modules for main and LoRA branches

_make_main_branch( *, in_features: int, out_features: int, tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, accumulate_into_main_grad: bool, ) → transformer_engine.pytorch.ops.Sequential#: Construct fused module for main branch (norm + fork + linear)

_make_lora_branch( *, in_features: int, out_features: int, tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, accumulate_into_main_grad: bool, ) → transformer_engine.pytorch.ops.Sequential#: Construct fused module for LoRA branch (linear_in + linear_out + add)

forward(x: torch.Tensor) → tuple[torch.Tensor, None]#

class bridge.peft.lora_layers.LinearAdapter( orig_linear: torch.nn.Linear, dim: int = 8, alpha: int = 32, dropout: float = 0.0, dropout_position: Literal[pre, post] = 'pre', lora_A_init_method: Literal[xavier, uniform] = 'xavier', lora_dtype: Optional[torch.dtype] = None, )#

Bases: torch.nn.Linear

Linear with LoRA that preserves the base weight and bias checkpoint keys.

Parameters:

orig_linear – The linear module to augment.
dim – LoRA’s dimension (in_features -> dim -> out_features).
alpha – LoRA’s scaling alpha.
dropout – Dropout probability (default: 0.0).
dropout_position – Where to apply dropout relative to LoRA (choices: [‘pre’, ‘post’], default=’pre’).
lora_A_init_method – Initialization method for lora_A (choices: [‘xavier’, ‘uniform’]).
lora_dtype – Adapter weight dtype. Defaults to the original linear’s weight dtype.

Initialization

Initialize LinearAdapter by copying from original Linear and adding LoRA components.

Parameters:

orig_linear – The original Linear module to adapt.
dim – LoRA rank dimension.
alpha – LoRA scaling factor.
dropout – Dropout probability.
dropout_position – When to apply dropout (‘pre’ or ‘post’ LoRA computation).
lora_A_init_method – Initialization method for LoRA matrix A.
lora_dtype – Data type for LoRA weights.

enable_adapter_layers() → None#: Enable the adapter layers, allowing them to contribute to the forward pass output.

disable_adapter_layers() → None#: Disable the adapter layers, making the forward pass return only the base module output.

_init_adapter( dim: int = 8, alpha: int = 32, dropout: float = 0.0, dropout_position: Literal[pre, post] = 'pre', lora_A_init_method: Literal[xavier, uniform] = 'xavier', lora_dtype: Optional[torch.dtype] = None, ) → None#

Initialize the LoRA weights and freeze the base parameters.

Parameters:

dim – LoRA’s dimension (in_features -> dim -> out_features).
alpha – LoRA’s scaling alpha.
dropout – Dropout probability (default: 0.0).
dropout_position – Where to apply dropout relative to LoRA (choices: [‘pre’, ‘post’], default=’pre’).
lora_A_init_method – Initialization method for lora_A (choices: [‘xavier’, ‘uniform’]).
lora_dtype – Adapter weight dtype. Defaults to the base weight dtype.

forward(x: torch.Tensor) → torch.Tensor#

Forward pass combining Linear output with LoRA adaptation.

Parameters:: x – Input tensor.
Returns:: Combined output from original linear layer and LoRA adaptation.

bridge.peft.lora_layers#

Module Contents#

Classes#

API#

`bridge.peft.lora_layers`#