burn_core/module/param/
lora.rs

1use super::{Param, Reparameterization};
2use crate as burn;
3use crate::module::Module;
4use burn_tensor::Tensor;
5
6/// A LoRA (Low-Rank Adaptation) adapter attached to a frozen weight [parameter](Param).
7///
8/// When present on a `Param<Tensor<2>>`, the parameter materializes its effective value as
9/// `base + scale * (a @ b)`, where `base` is the frozen (and optionally quantized) weight and
10/// `a`/`b` are the trainable low-rank factors. The frozen base is the stored value of the
11/// parameter; the adapter factors are surfaced to the optimizer, autodiff and record systems as
12/// regular parameters with their own [`ParamId`](super::ParamId)s through the module
13/// visitor/mapper traversal.
14#[derive(Debug, Module)]
15pub struct LoraAdapter {
16    /// Down-projection factor with shape `[d_in, rank]` (trainable).
17    pub a: Param<Tensor<2>>,
18    /// Up-projection factor with shape `[rank, d_out]` (trainable).
19    pub b: Param<Tensor<2>>,
20    /// Scaling factor applied to the low-rank product, typically `alpha / rank`.
21    pub scale: f64,
22}
23
24impl Reparameterization for LoraAdapter {
25    const NAME: &'static str = "lora";
26
27    fn materialize<const D: usize>(&self, base: Tensor<D>) -> Tensor<D> {
28        let delta = self.delta().reshape(base.shape());
29        // Compose at the factors' dtype: cast a mismatched dense base, but leave
30        // a packed base untouched so mixed addition dequantizes it directly to
31        // the delta's dtype rather than through the device default.
32        let base = if base.dtype().is_float() && base.dtype() != delta.dtype() {
33            base.cast(delta.dtype())
34        } else {
35            base
36        };
37        base + delta
38    }
39}
40
41impl LoraAdapter {
42    /// Compute the low-rank delta `scale * (a @ b)` with shape `[d_in, d_out]`.
43    ///
44    /// `a` and `b` are read through [`Param::val`], so the delta always reflects the current
45    /// (optimizer-updated) factors and keeps them as autodiff leaves for backpropagation.
46    pub fn delta(&self) -> Tensor<2> {
47        self.a.val().matmul(self.b.val()).mul_scalar(self.scale)
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    use crate::test_device;
55    use burn_tensor::DType;
56
57    #[test]
58    fn materialize_casts_a_dense_base_to_the_factor_dtype() {
59        let device = test_device();
60        let base = Tensor::<2>::ones([4, 4], (&device, DType::F32));
61        let adapter = LoraAdapter {
62            a: Param::from_tensor(Tensor::<2>::ones([4, 2], (&device, DType::F16))),
63            b: Param::from_tensor(Tensor::<2>::zeros([2, 4], (&device, DType::F16))),
64            scale: 1.0,
65        };
66
67        let materialized = adapter.materialize(base);
68
69        assert_eq!(materialized.dtype(), DType::F16);
70    }
71}
burn_core/module/param/lora.rs

burn_core/module/param/
lora.rs
