b7acc14d0c* 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>
92 lines
3 KiB
Python
92 lines
3 KiB
Python
# 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"
|
|
)
|