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

Delta-only LoRA adapter for a plain nn.Linear, mirroring :class:ParallelLinearAdapter’s role.

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)#

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:

When the wrapped module returns a Megatron-style tuple, a (combined_output, bias) tuple; when it returns a bare tensor (e.g. a plain nn.Linear), a bare tensor so the wrapper stays a drop-in replacement in simple (non-parallel) models.

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.Module

Delta-only LoRA adapter for a plain nn.Linear, mirroring :class:ParallelLinearAdapter’s role.

This adapter holds only the low-rank LoRA delta (linear_in -> linear_out) and produces just the scaled adaptation term. It is intended to be wrapped together with the original linear by :class:LoRALinear (the to_wrap / adapter pattern), so that base weights and adapter weights live in distinct submodules and adapter state is checkpointed under the adapter. prefix.

Parameters:

  • orig_linear – The linear module to augment (only its shape/dtype/device are used; its weights are not copied).

  • 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 the LoRA delta weights from the original Linear’s shape and dtype.

Parameters:

  • orig_linear – The original Linear module to adapt (weights are not copied).

  • 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.

_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,
device: Optional[torch.device] = None,
base_dtype: Optional[torch.dtype] = None,
) → None#

Initialize the LoRA delta weights.

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.

  • device – Device for the LoRA weights.

  • base_dtype – Base weight dtype, used when lora_dtype is not provided.

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

Compute the scaled LoRA delta only (no base-weight term).

Parameters:

x – Input tensor.

Returns:

The scaled low-rank adaptation scale * linear_out(linear_in(x)) with dropout applied per dropout_position.