From 85669451d0e4cfb2370994bdcad38b190cfbb417 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 20 Feb 2025 13:32:01 -0800 Subject: [PATCH] Fix num layers in fine tune (#1294) --- llms/mlx_lm/lora.py | 8 +++++++- llms/mlx_lm/tuner/utils.py | 7 +------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/llms/mlx_lm/lora.py b/llms/mlx_lm/lora.py index abc5dfa9..def3b6dd 100644 --- a/llms/mlx_lm/lora.py +++ b/llms/mlx_lm/lora.py @@ -181,8 +181,14 @@ def train_model( training_callback: TrainingCallback = None, ): model.freeze() + if args.num_layers > len(model.layers): + raise ValueError( + f"Requested to train {args.num_layers} layers " + f"but the model only has {len(model.layers)} layers." + ) + if args.fine_tune_type == "full": - for l in model.layers[-min(args.num_layers, 0) :]: + for l in model.layers[-max(args.num_layers, 0) :]: l.unfreeze() elif args.fine_tune_type in ["lora", "dora"]: # Convert linear layers to lora/dora layers and unfreeze in the process diff --git a/llms/mlx_lm/tuner/utils.py b/llms/mlx_lm/tuner/utils.py index 9f8d2925..f5df11e3 100644 --- a/llms/mlx_lm/tuner/utils.py +++ b/llms/mlx_lm/tuner/utils.py @@ -52,11 +52,6 @@ def linear_to_lora_layers( use_dora (bool): If True, uses DoRA instead of LoRA. Default: ``False`` """ - if num_layers > len(model.layers): - raise ValueError( - f"Requested {num_layers} LoRA layers " - f"but the model only has {len(model.layers)} layers." - ) def to_lora(layer): if isinstance(layer, (nn.Linear, nn.QuantizedLinear)): @@ -154,7 +149,7 @@ def linear_to_lora_layers( else: raise ValueError(f"Lora does not support {model.model_type}") - for l in model.layers[-min(num_layers, 0) :]: + for l in model.layers[-max(num_layers, 0) :]: lora_layers = [(k, to_lora(m)) for k, m in l.named_modules() if k in keys] if lora_layers: l.update_modules(tree_unflatten(lora_layers))