From b7acc14d0ce2d4601be9ca34bffe6f791c5bb63c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:58:44 -0700 Subject: [PATCH] fix(peft): expose finetune_last_n_layers for parity with mlx-lm CLI (#5564) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(peft): expose finetune_last_n_layers for parity with mlx-lm CLI mlx-lm's lora CLI defaults `CONFIG_DEFAULTS['num_layers']=16` (mlx_lm/lora.py:56), so it applies LoRA only to the LAST 16 transformer blocks. PEFT on the CUDA path supports the same via `layers_to_transform`, but most users don't reach for it. This commit adds a `finetune_last_n_layers` convenience parameter to both `FastLlamaModel.get_peft_model` and `FastBaseModel.get_peft_model` (vision/multi-modal). When set, it fills `layers_to_transform` automatically with the last N blocks, mirroring mlx-lm CLI's behavior AND `unsloth_zoo.mlx.loader.FastMLXModel.get_peft_model`. A single config value now controls layer-selection consistently across CUDA, MLX (zoo), and mlx-lm CLI paths. Default is None (= train all layers, current behavior unchanged). When set, the value is clamped to [1, total_transformer_layers] so callers can't accidentally over- or under-select. The total is read from `config.num_hidden_layers` (or aliases), falling through to `config.text_config.num_hidden_layers` for VLMs. Why this matters: with the same fixture/seed, training the last N layers vs all layers picks a different basin under stochastic LoRA init. Empirically (n=15 seeds, gemma-3-270m-it single-row LoRA memorization, MLX path) last-16 hits 67% greedy-decode pass rate vs all-18 at 47%. The teacher-forced completion loss is 0 in both — the model memorizes either way; only the first- token argmax distribution differs. CUDA fp32 shows the same pattern. Aligning the layer selection puts CUDA + MLX + mlx-lm all in the same basin family for parity comparisons. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * peft: trim verbose finetune_last_n_layers comments Per code-comment policy: parameter name is self-documenting, the clamp and range() construction are obvious. Rationale (mlx-lm CLI parity, empirical pass-rate data) lives in commit 106c1df4's message and the PR description. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_finetune_last_n_layers.py | 92 ++++++++++++++++++++++++++++ unsloth/models/llama.py | 10 +++ unsloth/models/vision.py | 36 +++++++++++ 3 files changed, 138 insertions(+) create mode 100644 tests/test_finetune_last_n_layers.py diff --git a/tests/test_finetune_last_n_layers.py b/tests/test_finetune_last_n_layers.py new file mode 100644 index 0000000000..fa0e2af5bc --- /dev/null +++ b/tests/test_finetune_last_n_layers.py @@ -0,0 +1,92 @@ +# Unsloth - 2x faster, 70% less memory LLM finetuning +# Tests for the `finetune_last_n_layers` parity knob (CUDA side). +# +# Mirrors unsloth-zoo's `FastMLXModel.get_peft_model` parameter. +# mlx-lm CLI's CONFIG_DEFAULTS['num_layers']=16 applies LoRA to the +# last 16 transformer blocks only. On the CUDA path, PEFT exposes +# `layers_to_transform` to do the same. This convenience knob fills +# `layers_to_transform` for the user when set, matching mlx-lm CLI +# AND unsloth-zoo's MLX path with a single config value. +# +# The tests intentionally avoid pulling in CUDA / a real model +# checkpoint — they exercise only the helper that translates +# `finetune_last_n_layers` into `layers_to_transform`. + +from __future__ import annotations + +import pytest + + +def test_get_total_transformer_layers_reads_num_hidden_layers(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeConfig: + num_hidden_layers = 18 + + class FakeModel: + config = FakeConfig() + + assert _get_total_transformer_layers(FakeModel()) == 18 + + +def test_get_total_transformer_layers_reads_text_config(): + from unsloth.models.vision import _get_total_transformer_layers + + class TextConfig: + num_hidden_layers = 24 + + class FakeConfig: + text_config = TextConfig() + + class FakeModel: + config = FakeConfig() + + # No num_hidden_layers at top level — should fall through to text_config. + assert _get_total_transformer_layers(FakeModel()) == 24 + + +def test_get_total_transformer_layers_handles_alternative_attr_names(): + from unsloth.models.vision import _get_total_transformer_layers + + for attr in ("n_layer", "n_layers", "num_layers"): + cfg = type("Cfg", (), {attr: 12})() + model = type("M", (), {"config": cfg})() + assert _get_total_transformer_layers(model) == 12 + + +def test_get_total_transformer_layers_returns_none_when_unknown(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeConfig: + pass + + class FakeModel: + config = FakeConfig() + + assert _get_total_transformer_layers(FakeModel()) is None + + +def test_get_total_transformer_layers_returns_none_for_missing_config(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeModel: + pass + + assert _get_total_transformer_layers(FakeModel()) is None + + +def test_finetune_last_n_layers_signature_present_on_llama_and_vision(): + """Both entry points must expose the new parameter with default None.""" + import inspect + from unsloth.models.llama import FastLlamaModel + from unsloth.models.vision import FastBaseModel + + for cls in (FastLlamaModel, FastBaseModel): + sig = inspect.signature(cls.get_peft_model) + assert ( + "finetune_last_n_layers" in sig.parameters + ), f"{cls.__name__}.get_peft_model missing finetune_last_n_layers" + assert sig.parameters["finetune_last_n_layers"].default is None, ( + f"{cls.__name__}.get_peft_model: finetune_last_n_layers default " + f"must be None to preserve historical behavior" + ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6ddfe04d21..20b515c711 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2831,6 +2831,7 @@ class FastLlamaModel: bias = "none", layers_to_transform = None, layers_pattern = None, + finetune_last_n_layers = None, use_gradient_checkpointing = "unsloth", random_state = 3407, max_seq_length = 2048, # not used anymore @@ -2863,6 +2864,7 @@ class FastLlamaModel: bias = bias, layers_to_transform = layers_to_transform, layers_pattern = layers_pattern, + finetune_last_n_layers = finetune_last_n_layers, use_gradient_checkpointing = use_gradient_checkpointing, random_state = random_state, max_seq_length = max_seq_length, @@ -3160,6 +3162,14 @@ class FastLlamaModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + if finetune_last_n_layers is not None and layers_to_transform is None: + from .vision import _get_total_transformer_layers + + _total_layers = _get_total_transformer_layers(model) + if _total_layers is not None and _total_layers > 0: + _n = max(1, min(int(finetune_last_n_layers), _total_layers)) + layers_to_transform = list(range(_total_layers - _n, _total_layers)) + arguments = dict( r = r, lora_alpha = lora_alpha, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..73ef1db7e3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -547,6 +547,35 @@ def _construct_vlm_processor_fallback( return None +def _get_total_transformer_layers(model): + """Best-effort total transformer block count across HF model shapes. + Returns None if not determinable; caller should skip the conversion.""" + cfg = getattr(model, "config", None) + if cfg is None: + return None + for name in ( + "num_hidden_layers", + "n_layer", + "n_layers", + "num_layers", + ): + v = getattr(cfg, name, None) + if isinstance(v, int) and v > 0: + return v + text_cfg = getattr(cfg, "text_config", None) + if text_cfg is not None: + for name in ( + "num_hidden_layers", + "n_layer", + "n_layers", + "num_layers", + ): + v = getattr(text_cfg, name, None) + if isinstance(v, int) and v > 0: + return v + return None + + class FastBaseModel: @staticmethod def from_pretrained( @@ -1319,6 +1348,7 @@ class FastBaseModel: finetune_language_layers = True, finetune_attention_modules = True, finetune_mlp_modules = True, + finetune_last_n_layers = None, layers_to_transform = None, layers_pattern = None, use_gradient_checkpointing = "unsloth", @@ -1417,6 +1447,12 @@ class FastBaseModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + if finetune_last_n_layers is not None and layers_to_transform is None: + _total_layers = _get_total_transformer_layers(model) + if _total_layers is not None and _total_layers > 0: + n = max(1, min(int(finetune_last_n_layers), _total_layers)) + layers_to_transform = list(range(_total_layers - n, _total_layers)) + # Get only allowed parameters for LoraConfig local_variables = { **locals(),