From 85669451d0e4cfb2370994bdcad38b190cfbb417 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 20 Feb 2025 13:32:01 -0800 Subject: [PATCH 1/2] 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)) From 3d793ecf6887512fd81f0f0c6bd156c06a6eaf88 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 20 Feb 2025 15:55:55 -0800 Subject: [PATCH 2/2] Fix logits processor bugs with spec dec (#1291) * Fix logits processor bugs with spec dec * bump patch --- llms/mlx_lm/_version.py | 2 +- llms/mlx_lm/utils.py | 19 ++++++------------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/llms/mlx_lm/_version.py b/llms/mlx_lm/_version.py index b2f98e6f..89e6cd00 100644 --- a/llms/mlx_lm/_version.py +++ b/llms/mlx_lm/_version.py @@ -1,3 +1,3 @@ # Copyright © 2023-2024 Apple Inc. -__version__ = "0.21.0" +__version__ = "0.21.5" diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py index 78a2e802..1fae76fa 100644 --- a/llms/mlx_lm/utils.py +++ b/llms/mlx_lm/utils.py @@ -409,8 +409,7 @@ def speculative_generate_step( for processor in logits_processors: logits = processor(tokens, logits) - logprobs = logits - mx.logsumexp(logits, keepdims=True) - logprobs = logprobs.squeeze(0) + logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) y = sampler(logprobs) return y, logprobs @@ -429,16 +428,14 @@ def speculative_generate_step( prev_tokens = ( mx.concat([prev_tokens, y]) if prev_tokens is not None else y ) - y, logprobs = _process_and_sample( - prev_tokens, logits[:, i : i + 1, :] - ) + y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :]) out_y.append(y) out_logprobs.append(logprobs) return mx.concatenate(out_y, axis=0), mx.concatenate( out_logprobs, axis=0 ) else: - return _process_and_sample(None, logits) + return _process_and_sample(None, logits.squeeze(0)) def _prefill(model, cache, y): while y.size > prefill_step_size: @@ -476,13 +473,9 @@ def speculative_generate_step( num_draft = min(max_tokens - ntoks, num_draft_tokens) draft_tokens = _draft_generate(draft_y, num_draft) if prev_tokens is not None: - prev_tokens = prev_tokens[ - : prev_tokens.size - draft_y.size - num_draft + 1 - ] + prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1] y = mx.concatenate([y, draft_tokens]) - tokens, logprobs = _step(model, model_cache, y, num_draft + 1) - mx.eval(tokens, draft_tokens) draft_tokens = draft_tokens.tolist() tokens = tokens.tolist() @@ -514,8 +507,8 @@ def speculative_generate_step( [mx.array(draft_tokens[-1:], mx.uint32), draft_y] ) - if prev_tokens is not None and n < num_draft: - prev_tokens = prev_tokens[: -(num_draft - n)] + if prev_tokens is not None: + prev_tokens = prev_tokens[: -max(num_draft - n, 1)] _rewind_cache(num_draft, n) finally: _rewind_cache(num_draft, n)