mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-25 01:41:19 +08:00
103 lines
2.8 KiB
Python
103 lines
2.8 KiB
Python
# Copyright © 2024 Apple Inc.
|
|
|
|
import math
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
|
|
class LoRALinear(nn.Module):
|
|
@staticmethod
|
|
def from_linear(
|
|
linear: nn.Linear,
|
|
r: int = 8,
|
|
alpha: float = 16,
|
|
dropout: float = 0.0,
|
|
scale: float = 10.0,
|
|
):
|
|
# TODO remove when input_dims and output_dims are attributes
|
|
# on linear and quantized linear
|
|
output_dims, input_dims = linear.weight.shape
|
|
if isinstance(linear, nn.QuantizedLinear):
|
|
input_dims *= 32 // linear.bits
|
|
lora_lin = LoRALinear(
|
|
input_dims=input_dims,
|
|
output_dims=output_dims,
|
|
r=r,
|
|
alpha=alpha,
|
|
dropout=dropout,
|
|
scale=scale,
|
|
)
|
|
lora_lin.linear = linear
|
|
return lora_lin
|
|
|
|
def to_linear(self, de_quantize: bool = False):
|
|
linear = self.linear
|
|
bias = "bias" in linear
|
|
weight = linear.weight
|
|
is_quantized = isinstance(linear, nn.QuantizedLinear)
|
|
|
|
# Use the same type as the linear weight if not quantized
|
|
dtype = weight.dtype
|
|
|
|
if is_quantized:
|
|
dtype = mx.float16
|
|
weight = mx.dequantize(
|
|
weight,
|
|
linear.scales,
|
|
linear.biases,
|
|
linear.group_size,
|
|
linear.bits,
|
|
)
|
|
output_dims, input_dims = weight.shape
|
|
fused_linear = nn.Linear(input_dims, output_dims, bias=bias)
|
|
|
|
lora_b = (self.scale * self.lora_b.T).astype(dtype)
|
|
lora_a = self.lora_a.T.astype(dtype)
|
|
fused_linear.weight = weight + lora_b @ lora_a
|
|
if bias:
|
|
fused_linear.bias = linear.bias
|
|
|
|
if is_quantized and not de_quantize:
|
|
fused_linear = nn.QuantizedLinear.from_linear(
|
|
fused_linear,
|
|
linear.group_size,
|
|
linear.bits,
|
|
)
|
|
|
|
return fused_linear
|
|
|
|
def __init__(
|
|
self,
|
|
input_dims: int,
|
|
output_dims: int,
|
|
r: int = 8,
|
|
alpha: float = 16,
|
|
dropout: float = 0.0,
|
|
scale: float = 10.0,
|
|
bias: bool = False,
|
|
):
|
|
super().__init__()
|
|
|
|
# Regular linear layer weights
|
|
self.linear = nn.Linear(input_dims, output_dims, bias=bias)
|
|
|
|
self.dropout = nn.Dropout(p=dropout)
|
|
|
|
# Scale for low-rank update
|
|
self.scale = scale * (alpha / r)
|
|
|
|
# Low rank lora weights
|
|
scale = 1 / math.sqrt(input_dims)
|
|
self.lora_a = mx.random.uniform(
|
|
low=-scale,
|
|
high=scale,
|
|
shape=(input_dims, r),
|
|
)
|
|
self.lora_b = mx.zeros(shape=(r, output_dims))
|
|
|
|
def __call__(self, x):
|
|
y = self.linear(x)
|
|
z = (self.dropout(x) @ self.lora_a) @ self.lora_b
|
|
return y + (self.scale * z).astype(x.dtype)
|