chore(lora): support mixtral in lora example (#343)

2025-10-23 22:18:06 +08:00 · 2024-01-20 06:07:45 -08:00
parent 527cea4027
commit 1415595409
6 changed files with 279 additions and 4 deletions
--- a/lora/models.py
+++ b/lora/models.py
@@ -328,7 +328,12 @@ def load(path_or_hf_repo: str):

    model = Model(model_args)
    if quantization is not None:
-        nn.QuantizedLinear.quantize_module(model, **quantization)
+        nn.QuantizedLinear.quantize_module(
+            model,
+            **quantization,
+            linear_class_predicate=lambda m: isinstance(m, nn.Linear)
+            and m.weight.shape[0] != 8,
+        )

    model.load_weights(list(weights.items()))