feat(mlx): PEFT adapter interop for the MLX backend - #957
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0a74fe355f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds PEFT ↔ MLX LoRA adapter interoperability so adapters can be trained/served across Apple silicon (MLX) and CUDA (PEFT/Transformers), while also making saving_utils.py importable on hosts without heavy CUDA-side deps.
Changes:
- Introduces PEFT↔MLX adapter detection/validation and directory conversion utilities, plus PEFT export support.
- Refactors
saving_utils.pyto use lazy imports for torch/bitsandbytes/transformers/peft to keep MLX-only installs importable. - Extends MLX save/load paths to preserve per-module rank/scale metadata and full-state module provenance across reload/export.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| unsloth_zoo/saving_utils.py | Adds adapter-format detection + PEFT↔MLX conversion/export helpers; refactors imports to be lazy. |
| unsloth_zoo/mlx/utils.py | Adds PEFT-format export path to save_lora_adapters, full-state persistence, and PEFT attach/bind logic. |
| unsloth_zoo/mlx/loader.py | Routes PEFT adapter dirs through new import path; supports per-module rank/scale on reload; improves quantization flag behavior. |
| tests/test_mlx_adapter_interop_metal.py | Adds end-to-end tests validating conversion, import/export refusal rules, and round-trips. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…T and MLX formats FastMLXModel.from_pretrained now loads PEFT-format adapter directories (adapter_model.safetensors with lora_A/lora_B keys) directly: the weight file drives attachment — module set and per-module ranks come from tensor shapes, per-module scales from lora_alpha/alpha_pattern with peft's ordered first-match resolution — the base is frozen like a peft is_trainable load, caller quantization controls (flags, q_* knobs, or config dicts, in any spelling) are honored for the base load including bnb-remapped bases, and any unbindable tensor aborts the load instead of returning a partial adapter. Configs are whitelist-validated with named rejections for features whose semantics cannot be preserved (DoRA, modules_to_save, embedding LoRA, expert target_parameters, aLoRA, QALoRA, lora_bias, layer_replication, and base-mutating initializers such as PiSSA/OLoRA/CorDA/LoftQ). unsloth_zoo.saving_utils gains directory-to-directory converters in both directions with bitwise dtype-preserving round trips for plain linear LoRA, a root-anchored model.layers gate (plus a module_types oracle for out-of-stack paths) so exported paths always bind in peft, exact-match full-path target_modules, lora_parameters keys pinning so stock mlx-lm does not grow adapter topology, per-module rank/scale metadata that the loader honors on reload (rebuilding wrappers whose global-parameter reconstruction disagrees), and create-only destinations so two adapter formats can never mix in one directory. saving_utils' heavy imports (torch, bitsandbytes, peft, transformers, safetensors.torch, and torch-chained internal modules) are now resolved lazily so the module imports on MLX-only Apple installs while the torch fallback serves hosts without mlx; behavior on CUDA hosts is unchanged and the full test suite matches the pre-change baseline exactly.
save_lora_adapters (and the bound model method) gains adapter_format="mlx"|"peft" with the default unchanged. PEFT export saves the canonical mlx artifact to scratch and converts through the shared directory converter using the live module tree as the linear-vs-embedding oracle — out-of-stack linear adapters such as an untied lm_head convert when every wrapped path demonstrates the Hugging-Face-mirroring model.layers layout, and anything the format cannot represent (custom or fused wrappers named with the offending class, stacked switch factors, mixed per-module lora_dropout) refuses by name before tensor collection. A public export_peft_adapter wrapper (keyword-only after src/dst) joins saving_utils' __all__ with base_weights_source reserved for full-weight emission. Conversion destinations are hardened: complete artifacts stage in uuid-unique owned temp directories and publish with one atomic rename, so readers can never observe a partial adapter, a mid-write failure leaves nothing behind, and a retry at the same path works; raw destination spellings — trailing separators and dangling symlinks under both spellings included — are refused before path resolution; os.PathLike destinations are supported; and reconstructed rank/alpha patterns use caret-anchored escaped keys with modal global references so peft's matcher resolves every module exactly.
Linear DoRA adapters now import, attach, train, and export across the PEFT/MLX boundary: peft's lora_magnitude_vector maps to mlx-lm's m without transpose (forward equivalence verified numerically to float32 epsilon before enabling), use_dora and magnitude presence must agree everywhere (all-or-nothing per artifact in both directions), magnitude shapes are validated symmetrically against the module out-dim in both converters and at attach, and exports write peft's canonical suffix-free magnitude key. Refusals guard every semantics-dropping corner: DoRA with nonzero lora_dropout (peft applies dropout inside the DoRA correction while mlx-lm scales the undropped base plus the dropped update, so continued training would diverge), embedding DoRA in both directions (the two decompositions are incompatible), and mixed plain-LoRA/DoRA models at both adapter-save entry points including trainer checkpoints (one fine_tune_type per artifact; a mixed save would reload every module as DoRA).
… interop Imports consume PEFT full-module state — modules_to_save snapshots (bias and composite submodules included) and auto-saved embedding/output-head weights in every genuine on-disk form (plain keys, .base_layer keys on LoRA-wrapped modules, and the wrapper-internal duplicate/pristine copies, which fold or drop with consistency checks) — using dtype-tolerant verify-or-apply semantics and per-origin trainability: modules_to_save stays trainable, auto-saved embedding replacements stay frozen. An origin-tagged full_state_modules map persists this state through MLX save/reload (wrapped modules keep their tree-native inner paths; reload re-verifies shapes and dtypes against a pre-bind snapshot because load_weights(strict=False) validates neither) and reconstructs the PEFT modules_to_save list from origin tags on re-export. Embedding LoRA converts in both directions (lora_embedding_A/B <-> LoRAEmbedding by the shared transpose math), with PEFT-matching dropout semantics (PEFT applies none on embeddings, so nonzero lora_dropout refuses rather than silently diverging) and auto-saved embedding emission sourced from the artifact itself, from base_weights_source (validated against factor-implied geometry), or omitted with a log. Models that route output logits through a tied embedding refuse embedding adapters and tied-module replacements on every surface (attach, both converters, live export, native reload). The routing rule is conservative: a declared tie or a tree without any known output head (lm_head/output/embed_out spellings, root or nested) counts as tied, since module presence cannot prove routing. Tied output-head duplicates fold after verifying they exactly match the embedding snapshot. Converted artifacts carry an unsloth_peft_converted marker that survives load/save cycles, so derivative artifacts keep these checks; natively trained tied embedding adapters remain loadable under MLX semantics. Non-floating factor or state dtypes, truncated or internally inconsistent snapshots, declared-but-tensorless modules_to_save entries, resized vocabularies, quantized full-state targets, embedding DoRA, and dual linear+embedding factor sets are refused with named reasons instead of converting lossily. Embedding and head module spellings are enumerated from the mlx-lm source inventory.
Checkpoints whose config declares a vision, audio or dflash sub-config load through mlx-vlm, which nests the text tower under language_model. A PEFT adapter trained against the Hugging Face text model names its modules one level shallower, so none of them matched and the import refused every entry — for ordinary text-only LoRAs as much as multimodal ones. This affected much of the current model landscape, including Qwen3.5, Gemma 3/4, Qwen2.5-VL and Llama-3.2-Vision. The import now resolves that nesting once, from the LoRA pairs alone, and binds every path through it. A prefix is only considered when it is one of the attribute names mlx-vlm uses for a text tower, so a vision or audio tower carrying identically shaped projections can never be selected; an adapter that fits both the root and a nested tower is refused rather than guessed at, and anything else falls through to the existing per-path report. Full-state entries are resolved separately so a legitimately absent counterpart — a tied model has no output head — explains itself instead of suppressing the nesting every adapter tensor needs. Export is symmetric but strictly provenance-driven: the resolved nesting is remembered on the model, persisted in the MLX artifact, and restored on reload, so a save/reload/re-save cycle keeps it. Only that provenance, or an explicit caller assertion, unnests an adapter on the way out; the Hugging Face counterpart of a nested tower cannot be recovered by inspection, since mlx-vlm's language_model.model.layers path corresponds to model.language_model.layers for a multimodal class but to plain model.layers for a text class. A natively trained nested adapter therefore keeps the existing layout refusal rather than emitting paths that bind the wrong modules. Unnesting is checked to be injective, so an artifact carrying entries for more than one tower is refused instead of silently discarding one, and the mixed-dropout guard now compares live module names in the same space it scans. Validated across machines on Qwen3.5-0.8B: a CUDA-trained adapter imports and reproduces its source logits, continues training on Apple silicon, exports back, and reloads into PEFT on the CUDA host, with per-module rank and alpha patterns, DoRA and quantized bases exercised on the same nested checkpoint.
Exporting an embedding LoRA without base_weights_source logged through the shared logger, which resolves on first use by importing a module that imports torch. saving_utils is otherwise importable and usable without torch, so that export raised ModuleNotFoundError on a torch-free Apple install. Drop the informational log; the behaviour it described is unchanged.
3406d9f to
98b1405
Compare
|
Rebased onto main to pick up #953 ( This branch already makes if isinstance(module, _bnb_linear4bit()):The helper returns It also filters to real classes, which fixes something worth flagging separately: on a host with no CUDA build, Reproduced on main as of 2430906 on an Apple-silicon host. That is a pre-existing issue in the stubbed path rather than anything this branch introduces, and it is incidental to this PR — happy to split it into its own fix if you would rather the change did not ride along here. |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 98b1405b50
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
The mixed-dropout export guard reimplemented the wrapper dropout lookup inline and diverged from _get_mlx_dropout_probability, which already prefers _p_1, falls back to .p, tolerates non-numeric values and treats a missing dropout as zero. A wrapper without a dropout object therefore contributed nothing to the comparison, so a model mixing it with a nonzero dropout looked uniform. Call the shared helper instead.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d801009e53
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…path Older mlx-lm wheels reject scale/dropout keywords on LoRA and DoRA from_base. The native adapter paths already install the compatibility shim for that gap, but the PEFT import built its wrappers directly, so every import raised TypeError on those wheels.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7bfe5ba3d8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7bfe5ba3d8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
The MLX backend builds a causal language model, but task_type sat among the config keys treated as safe to ignore, so a classification or feature-extraction adapter could attach its backbone factors and then answer with vocabulary logits. Refuse a declared non-causal task by name. peft_type and task_type also arrive as peft enums when the caller passes a live LoraConfig.to_dict() rather than adapter JSON, and str() on those reads 'PeftType.LORA', which refused every such config. Compare their values.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 056a6bb4a6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
What this does
LoRA adapters could not cross between Apple silicon and CUDA machines. An adapter trained with PEFT on a GPU box could not be loaded by the MLX backend, and an adapter trained on MLX could not be loaded by
transformers/PEFT anywhere else. This adds conversion in both directions, so the same LoRA can be trained on one platform and continued or served on the other.Two entry points change behaviour:
FastMLXModel.from_pretrained(<peft adapter dir>)now imports PEFT adapters directly, attaching and binding them onto the live MLX tree.save_lora_adapters(model, path, adapter_format="peft")writes a standard Hugging Face adapter. The default stays"mlx", so training checkpoints and every existing caller are untouched.Supported and refused
Converted: plain linear LoRA, linear DoRA (the magnitude vector maps to mlx-lm's
m), embedding LoRA vialora_embedding_A/B, per-modulerank_pattern/alpha_pattern, rsLoRA scaling, and full-module state — bothmodules_to_savesnapshots and PEFT's auto-saved embeddings, tracked through an origin-tagged map so a save round-trips them with the right trainability.Refused by name, rather than partially or lossily converted: expert
target_parametersfactors, embedding DoRA, embedding adapters on tied-output models, and any config field carrying semantics this converter does not implement. That last one is deliberate — an unknown non-empty field is refused rather than ignored, so for example an aLoRA adapter cannot silently import as an always-on LoRA.How the mapping works
The two formats encode the same math in different layouts:
Conversion is rename plus transpose only, so dtypes are preserved exactly and a round trip is bitwise identical for the supported subset. This is measured, not assumed — see the validation section.
Two design points are worth a reviewer's attention.
Nested text towers. A checkpoint carrying a vision, audio, or dflash sub-config is routed through mlx-vlm, which nests the text tower under
language_model.. Ordinary text-only Hugging Face LoRAs therefore matched nothing on such models. Import now resolves that nesting once from the LoRA pairs alone and binds every path through it, all or nothing. Candidate prefixes are restricted to known text-tower attribute names, so a vision tower can never be selected even when its projection shapes would fit.Export never guesses the nesting. The Hugging Face counterpart of a nested tower is not recoverable by inspection: mlx-vlm's
language_model.model.layers.*corresponds tomodel.language_model.layers.*under a multimodal class but to plainmodel.layers.*under a text class. Export therefore unnests only when the adapter carries provenance recorded at import, and refuses otherwise, pointing at the MLX-format save instead. Guessing would emit an adapter whose weights bind to nothing and silently train from zero.Directory-only conversion is correspondingly conservative: without a caller-supplied
module_typesmap it converts only root-anchoredmodel.layers.N.paths, the one layout where mlx-lm mirrors the Hugging Face tree. GPT-2-style stacks rename roots (model.h.*versustransformer.h.*), so an unanchored heuristic would emit adapters PEFT cannot bind.Adapter directories are published by staging into a unique sibling directory and renaming once, so a reader never observes a partial adapter and a crash leaves either nothing or the complete artifact.
Validation
Beyond the 46 tests added here, the loop was driven end to end on real hardware between an Apple silicon host and a CUDA host, in five orderings, on a nested model (
Qwen/Qwen3.5-0.8B, which routes through mlx-vlm) withQwen/Qwen3-0.6Bas a non-nested control:Agreement at every hop was top-5 5/5 and correlation ≥ 0.9996 on held-out probe logits (both sides bf16, so exact equality is not the right bar; one hop showed 4/5 from a near-tie reordering).
Two properties were measured rather than assumed:
peft.save_pretraineddrops the provenance keys, because PEFT rebuildsadapter_config.jsonfrom its ownLoraConfig. This is harmless: on a second crossing the tower prefix is re-derived from the adapter's own target paths. Provenance accelerates the export direction; it is never a correctness dependency of import.0.000e+00across a full crossing.Local suites:
python -m pytest tests/test_mlx_adapter_interop_metal.py -q # 46 passedThe full suite matches the pre-change baseline (the same two known-flaky MoE tests, unrelated to this change).
Notes for reviewers
adapter_formatdefaults to"mlx".module_typesparameter onconvert_mlx_dir_to_peftis an assertion interface for callers that can establish what a weights file cannot: the live module type, that the base does not route tied output logits through it, and that the same dotted path exists in the Hugging Face tree. Its contract is documented on the function.declared-True OR no-head-module, checkinglm_head,output, andembed_outspellings. A declaredTruestays conservative even where a forward would ignore it, because a wrongly-admitted embedding adapter silently changes output logits.