Feature: QDoRA (#891)

* feat: QDoRA with tests and a small bug fix for recalculation of self.m

* some simplifications and fixes

---------

Co-authored-by: Awni Hannun <awni@apple.com>

llms/tests/test_finetune.py

@@ -11,7 +11,7 @@ import mlx.nn as nn
 import mlx.optimizers as opt
 from mlx.utils import tree_flatten
 from mlx_lm import lora, tuner
-from mlx_lm.tuner.dora import DoRAEmbedding
+from mlx_lm.tuner.dora import DoRAEmbedding, DoRALinear
 from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear
 from mlx_lm.tuner.trainer import evaluate
 from mlx_lm.tuner.utils import build_schedule
@@ -164,6 +164,147 @@ class TestDora(unittest.TestCase):
         self.assertFalse(mx.array_equal(embedding.weight, new_embedding.weight))
         self.assertFalse(mx.array_equal(embedding(tokens), dora_emb(tokens)))
 
+    def test_llama(self):
+        from mlx_lm.models import llama
+
+        hidden_size = 1024
+        intermediate_size = 2048
+        args = llama.ModelArgs(
+            model_type="llama",
+            hidden_size=hidden_size,
+            num_hidden_layers=4,
+            intermediate_size=intermediate_size,
+            num_attention_heads=4,
+            rms_norm_eps=1e-5,
+            vocab_size=10_000,
+        )
+
+        dora_layers = 4
+
+        def check_config(params):
+            n_keys = 2
+            if "keys" in params:
+                n_keys = len(params["keys"])
+            model = llama.Model(args)
+            model.freeze()
+            tuner.utils.linear_to_lora_layers(model, dora_layers, params, use_dora=True)
+            trainable_params = sum(
+                v.size for _, v in tree_flatten(model.trainable_parameters())
+            )
+            self.assertEqual(
+                trainable_params,
+                dora_layers
+                * (params["rank"] * hidden_size * 2 * n_keys + n_keys * hidden_size),
+            )
+
+        params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
+        check_config(params)
+
+        params["rank"] = 1
+        check_config(params)
+
+        params["keys"] = ["self_attn.k_proj"]
+        check_config(params)
+
+    def test_dora_m_parameter(self):
+        dora_lin = DoRALinear(input_dims=100, output_dims=100)
+        self.assertTrue(
+            mx.allclose(dora_lin.m, mx.linalg.norm(dora_lin.linear.weight, axis=1))
+        )
+
+        # Recomputes m when changing Linear
+        inital_m = dora_lin.m
+        lin = nn.Linear(10, 10)
+        dora_lin.set_linear(lin)
+        self.assertTrue(mx.allclose(dora_lin.m, mx.linalg.norm(lin.weight, axis=1)))
+
+        # Works with quantized weights
+        quantized_linear = nn.QuantizedLinear(512, 512)
+        dora_lin.set_linear(quantized_linear)
+        dequantized_weight = mx.dequantize(
+            quantized_linear.weight,
+            quantized_linear.scales,
+            quantized_linear.biases,
+            quantized_linear.group_size,
+            quantized_linear.bits,
+        )
+        self.assertTrue(
+            mx.allclose(dora_lin.m, mx.linalg.norm(dequantized_weight, axis=1))
+        )
+
+    def test_dora_from_linear(self):
+        in_dims = 256
+        out_dims = 256
+        r = 4
+
+        linear = nn.Linear(in_dims, out_dims)
+        dora_lin = DoRALinear.from_base(linear, r)
+        self.assertTrue(mx.allclose(dora_lin.m, mx.linalg.norm(linear.weight, axis=1)))
+        self.assertEqual(dora_lin.lora_a.shape, (in_dims, r))
+        self.assertEqual(dora_lin.lora_b.shape, (r, out_dims))
+        self.assertEqual(dora_lin.m.shape, (out_dims,))
+
+        quantized_linear = nn.QuantizedLinear(in_dims, out_dims)
+        dequantized_weight = mx.dequantize(
+            quantized_linear.weight,
+            quantized_linear.scales,
+            quantized_linear.biases,
+            quantized_linear.group_size,
+            quantized_linear.bits,
+        )
+        dora_quant_lin = DoRALinear.from_base(quantized_linear, r)
+        self.assertTrue(
+            mx.allclose(dora_quant_lin.m, mx.linalg.norm(dequantized_weight, axis=1))
+        )
+        self.assertEqual(dora_quant_lin.lora_a.shape, (in_dims, r))
+        self.assertEqual(dora_quant_lin.lora_b.shape, (r, out_dims))
+        self.assertEqual(dora_quant_lin.m.shape, (out_dims,))
+
+    def test_dora_to_linear(self):
+        in_dims = 256
+        out_dims = 256
+        r = 4
+
+        linear = nn.Linear(in_dims, out_dims, bias=True)
+        dora_lin = DoRALinear.from_base(linear, r)
+        to_linear = dora_lin.fuse()
+        self.assertTrue(mx.allclose(linear.weight, to_linear.weight))
+        self.assertTrue(mx.allclose(linear.bias, to_linear.bias))
+
+        def dequantize_weight(quantized_linear):
+            return mx.dequantize(
+                quantized_linear.weight,
+                quantized_linear.scales,
+                quantized_linear.biases,
+                quantized_linear.group_size,
+                quantized_linear.bits,
+            )
+
+        quantized_linear = nn.QuantizedLinear(in_dims, out_dims, bias=True)
+        dora_quantized_linear = DoRALinear.from_base(quantized_linear, r)
+        # Dequantize
+        to_linear_from_quantized = dora_quantized_linear.fuse(de_quantize=True)
+        self.assertTrue(
+            mx.allclose(quantized_linear.bias, to_linear_from_quantized.bias)
+        )
+        self.assertTrue(
+            mx.allclose(
+                dequantize_weight(quantized_linear), to_linear_from_quantized.weight
+            )
+        )
+
+    def test_dora_dtype(self):
+        in_dims = 256
+        out_dims = 256
+        r = 4
+
+        linear = nn.Linear(in_dims, out_dims, bias=True)
+        linear.set_dtype(mx.float16)
+        dora_lin = DoRALinear.from_base(linear, r)
+
+        x = mx.random.uniform(shape=(2, 256)).astype(mx.float16)
+        self.assertEqual(dora_lin(x).dtype, mx.float16)
+
 
 class TestScheduleConfig(unittest.TestCase):
     def test_join(self):