Add LoRA fine-tuning to the MLX backend - #51
Closed
ovuruska wants to merge 5 commits into
Closed
Conversation
Documents the repo layout, Google/pyink coding style, the per-backend integration pattern, numerical-fidelity rules for ports, and test conventions so coding agents follow existing conventions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add tabfm/src/mlx/ following the existing backend pattern (model.py + tabfm_v1_0_0.py). The MLX port mirrors the PyTorch module/parameter naming and Linear [out, in] layout, so it loads the PyTorch v1.0.0 safetensors release directly with no weight conversion step. - Faithful forward-path port with the same fp32 upcasts as JAX/PyTorch (RMSNorm, Fourier expansion, PerDimScale softplus, RoPE phases); SDPA at scale=1.0 with checkpoint-loaded RoPE frequencies. - sklearn wrappers dispatch MLX models through the shared eager execution path (_predict_step_mlx); JAX path untouched. - torch<->mlx parity test asserts < 1e-4 max abs diff in float32 for classification and regression; sklearn fit/predict integration test. - Released-checkpoint compatibility verified at the safetensors-header level: all tensor names/shapes of the classification (913) and regression (918) checkpoints map one-to-one onto the MLX module tree. - CI installs the mlx extra (manylinux wheels exist for cp311) so the new tests also run on ubuntu-latest. Constraint: MLX has no one_hot op and its gelu variants use erf/sigmoid approximations -- one-hot is built by comparing against the class range and jax.nn.gelu's tanh approximation is implemented manually. Rejected: separate MLX weight release on Hugging Face -- names and layouts match the PyTorch release exactly, so re-hosting ~13GB of identical tensors adds no value. Confidence: high Scope-risk: low -- additive; the only shared-code change is renaming the PyTorch branch of _batch_forward to a generic eager branch. Not-tested: end-to-end load() of the full 6.5GB pre-trained checkpoint (verified via safetensors header instead); MLX CPU wheel behavior on ubuntu-latest CI. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With a bare `pip install -e .[pytorch]`, tabfm_v1_0_0_pytorch.load() fails with `NameError: name 'safetensors' is not defined` inside PyTorchModelHubMixin._load_as_safetensor: huggingface-hub only imports safetensors conditionally and does not depend on it, and neither torch nor any core dep pulls it in. Found while end-to-end testing the v1.0.0 release checkpoint locally. Confidence: high Scope-risk: low Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Pin the mlx extra to >=0.31 (the port needs boolean SDPA masks, nan_to_num, Module.set_dtype; an older preinstalled mlx would otherwise satisfy the extra and crash at runtime). - Materialize weights eagerly in load() via mx.eval so checkpoint errors surface at load time and the process-wide cached model is safe to share across threads. - Simplify OneHotAndLinear: the sentinel-class remap and the sliced-off extra comparison column were torch-F.one_hot artifacts; direct comparison against the class range is behaviorally identical (parity test still <1e-4). - Add the missing safetensors pin to the requirements.txt lock. - Sync CLAUDE.md with this PR (mlx backend in the layout, CI install line) and add the license header to tabfm/src/mlx/__init__.py. - Wrap an over-length line in Encoder. Full suite: 76 passed with jax+torch+mlx installed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Parameter-efficient fine-tuning on top of the frozen pre-trained model: apply_lora freezes the base and wraps target attention projections (default: q_proj/v_proj in the ICL predictor) with trainable low-rank adapters; train_lora optimizes only the adapters with in-context episodes (context/target row splits, loss on targets only, -100 sentinel masking of target labels); fit_lora reuses the sklearn wrapper's fit() preprocessing so training matches the predict-time input distribution; merge_lora folds adapters back into plain Linears for zero-overhead inference; save_adapters/load_adapters round-trip only the adapter tensors. Adapters are float32 regardless of the bf16 base compute (adapter gradients are smaller than bf16 accumulation noise). LoRA-B starts at zero, so an adapted model is exactly the base model at init. Tests: base-frozen invariant (bit-identical after training), loss descent, merge parity, adapter save/load round-trip, wrapper integration, regression path smoke. Full suite: 82 passed. Constraint: mlx core ships no LoRA layer (mlx-lm keeps its own) -- implemented LoRALinear locally with the standard init. Rejected: adapting all towers by default -- the ICL predictor holds most parameters (24 blocks at 8x base width in v1.0.0) and is where label-conditioned adaptation acts; scope="all" remains available. Confidence: high Scope-risk: low -- new module; no existing code path changes except conftest test gating. Not-tested: accuracy gain over zero-shot on the real checkpoint (validation experiment queued; results to be added to the PR). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds parameter-efficient LoRA fine-tuning to the MLX backend. The pre-trained base weights are frozen and never updated — training optimizes only low-rank adapter matrices.
Design
lora.apply_lora(model, rank, alpha, target_keys, scope)— freezes the whole model first, then wraps target attention projections (defaultq_proj/v_projin the ICL predictor, where most parameters and the label-conditioned computation live) withLoRALinear. Because adapters are added after the freeze,trainable_parameters()is exactly the adapter set.lora.train_lora(model, X, y, ...)— trains with in-context episodes mirroring how TabFM is used: each step samples rows, splits into context (labels visible) and targets, takes cross-entropy/MSE on target positions only. Target labels are masked with the-100sentinel, so label leakage is structurally impossible.lora.fit_lora(estimator, X, y, ...)— the sklearn-wrapper bridge: runs the estimator's ordinaryfit()(label/feature encoders + ensemble generator + preprocessing — which never touches model weights), then trains the adapters on the same preprocessed matrix the wrapper feeds the model at predict time.lora.merge_lora(model)— folds adapters into the baseLinears (W += (α/r)·(A·B)ᵀ) for zero-overhead inference;save_adapters/load_adaptersround-trip only the adapter tensors (KBs, not GBs).Binitializes to zero → an adapted model is bit-equivalent to the base until trained.Tests (all passing; full suite 82 passed)
fit_lora→predictthroughTabFMClassifier.Validation on real weights
Accuracy-vs-zero-shot validation on the released 6.5GB checkpoint is running; results will be posted here.
🤖 Generated with Claude Code