LoRALinear

class torchtune.modules.peft.LoRALinear(in_dim: int, out_dim: int, rank: int, alpha: float, dropout: float = 0.0, use_bias: bool = False, quantize_base: bool = False, **quantization_kwargs)

LoRA linear layer as introduced in LoRA: Low-Rank Adaptation of Large Language Models.

LoRA perturbs a given layer via a low-rank approximation where only the rank decomposition matrices are trainable. In a linear layer instead of \(x \mapsto W_0x\) a LoRALinear layer is defined as \(x \mapsto W_0x + (\alpha / r)BAx\), where \(r\) is the rank of the matrices \(A\) and \(B\) and \(\alpha\) is a scaling factor. As in the original implementation, we support dropout before multiplication by the low-rank matrices.

Parameters:

• in_dim (int) – input dimension

• out_dim (int) – output dimension

• rank (int) – rank of the low-rank approximation

• alpha (float) – scaling factor for the low-rank approximation

• dropout (float) – dropout probability. Default: 0.0

• use_bias (bool) – whether to include bias in the original linear layer. Default: False

• quantize_base (bool) – Whether to quantize base linear weight or not. Default: False

• **quantization_kwargs – Keyword arguments to pass to to_nf4 when quantizing the base linear weight. Examples of valid arguments are block_size and scaler_block_size, which control the granularity of weight quantization and scaler quantization respectively. This is only used if quantize_base is True. Default None

Raises:

ValueError – If quantize_base is False, but quantization kwargs are provided.

adapter_params() → List[str]

Return a list of strings corresponding to the names of the nn.Parameter s in the model coming from the adapter.

For LoRA this means lora_a.weight and lora_b.weight.

forward(x: Tensor) → Tensor

Parameters:

x (torch.Tensor) – input tensor with shape (..., in_dim)

Returns:

output tensor with shape (..., out_dim)

Return type:

torch.Tensor

to_empty(*, device: Optional[Union[str, device, int]], recurse: bool = True)

Move the parameters and buffers to the specified device without copying storage.

Parameters:

• device (torch.device) – The desired device of the parameters and buffers in this module.

• recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module