feat(mlx): Make adamw_8bit actually use 8-bit optimizer state - #962
feat(mlx): Make adamw_8bit actually use 8-bit optimizer state#962BardiaKoopah wants to merge 2 commits into
Conversation
`trainer.py:745-758` accepts `adamw_8bit` and rewrites it to plain `adamw` with
no message, so a user asking for an 8-bit optimizer on MLX silently gets full
fp32 optimizer state. `adam_8bit` was not accepted at all.
Delivered. Both names now route to `QuantizedMomentAdamW` / `QuantizedMomentAdam`
in a new `unsloth_zoo/mlx/optimizers_quantized.py`, holding Adam's FIRST moment
as 8-bit packed arrays. Optimizer state 1,052,684 -> 675,852 B: 35.8% measured on
the test model, 35.9% on a wider one. Final loss 0.08318 against an fp32 baseline
of 0.08317 on the same seed. Both classes wrap the stock MLX update, so bias
correction and AdamW's decoupled decay are unchanged. Routing prints once, naming
what you get: first moment 8-bit (group_size=64), second moment float32, and only
2-D parameters whose last dim is a multiple of the group size.
Documented limitation: the second moment is NOT quantized. `v` is a running mean
of squared gradients, so it spans orders of magnitude and sits near zero early.
MLX's affine `mx.quantize` reconstructs poorly there, `sqrt(v)` lands in the
update denominator, and the step explodes -- measured, both for int8 m+v and
fp32 m + int8 v:
1.0590 -> 0.8627 -> 19.3755 -> 10.2466 -> nan
That needs a dynamic/exponential mapping with stochastic rounding, which
bitsandbytes implements and MLX does not expose. A guard test asserts `v` stays a
single float32 array. `bits` is constrained to 8 with a ValueError rather than
shipped untested; `group_size` accepts 32/64/128, all three exercised.
`paged_adamw_8bit` and `adamw_bnb_8bit` keep collapsing to `adamw` deliberately:
"paged" promises CPU offload and "bnb" names a library MLX does not use, so
routing them here would swap one wrong answer for another.
Two test files on purpose: the numeric tests need real `mx.quantize`, so one file
would skip entirely on Linux CI and be a gate no-op (unslothai#669/unslothai#739). `..._routing.py`
runs under the existing torch shim and executes in the gate; `..._state.py` holds
the behaviour tests and runs on Apple Silicon.
Tested, 32 tests: loss decreases over 8 steps; within 1e-4 of fp32; asserted byte
reduction; save/reload/resume reproduces the next step exactly; compiled and
uncompiled bit-identical under `mx.compile(inputs=state, outputs=state)`; a model
with no eligible parameter still trains. Full suite twice per side, totals
2979 -> 3011.
Textual overlap with unslothai#819 on `trainer.py:56` and the `_build_optimizer` chain --
same author, whichever merges second rebases.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dc36474205
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # regress silently). test_mlx_quantized_optimizer_routing pins which optim | ||
| # names reach the 8-bit path; it runs under the torch shim so it EXECUTES | ||
| # here rather than skipping (the numeric half needs real MLX and runs on | ||
| # the Apple Silicon job). Executing them here keeps fixture/source drift visible. |
There was a problem hiding this comment.
For PRs/scheduled MLX CI, this new comment says the numeric half runs on Apple Silicon, but .github/workflows/mlx-ci.yml still invokes only the existing fixed MLX test files and never runs tests/test_mlx_quantized_optimizer_state.py; on Linux that file is only collected/skipped because of pytest.importorskip("mlx.core"). As a result, regressions in the new QuantizedMomentAdam(W) update, checkpoint, and compile behavior won't fail CI unless the Apple Silicon job is extended to run that file.
Useful? React with 👍 / 👎.
| def apply_single(self, gradient, parameter, state): | ||
| packed = state["m"] | ||
| # tuple OR list: tree_map/tree_unflatten rebuild the triple as a list. | ||
| quantized = isinstance(packed, (tuple, list)) |
There was a problem hiding this comment.
When adamw_8bit resumes from a checkpoint written by the old implementation (which collapsed the same setting to plain AdamW) or from a user switching an AdamW checkpoint to the 8-bit optimizer, state['m'] loads as a plain mx.array. Because this flag is derived only from the saved container type, it stays false and line 136 never quantizes the eligible moment, so the run continues with full-precision first-moment state despite the 8-bit optimizer selection; use is_quantizable(parameter) to quantize migrated array moments after the parent update.
Useful? React with 👍 / 👎.
Same-backend paired differencesDifferences are
Memory and wall time
Training-call time excludes full-dataset formatting/tokenization. These are single paired measurements, so small speed differences are not throughput claims. State bytes are exact recursive array sums; peak allocation is allocator-reported and, on MLX unified memory, is not resident system memory. Dataset, order, and protocol
|
No functional change: comment and docstring text only, verified by ast.parse and the test suite.
2 participants
The silent ignore
trainer.py:745-758acceptsadamw_8bitand rewrites it to plainadamwwith no message, so a user asking for an 8-bit optimizer on MLX silently gets full fp32 optimizer state.adam_8bitwas not accepted at all.Delivered
Both names now route to
QuantizedMomentAdamW/QuantizedMomentAdamin a newunsloth_zoo/mlx/optimizers_quantized.py, holding Adam's first moment as 8-bit packed arrays.Optimizer state 1,052,684 → 675,852 B — 35.8% measured on the test model, 35.9% on a wider one. Final loss 0.08318 against an fp32 baseline of 0.08317 on the same seed.
Both classes wrap the stock MLX update, so bias correction and AdamW's decoupled decay are unchanged. Routing prints once, naming what you get: first moment 8-bit (
group_size=64), second moment float32, and only 2-D parameters whose last dim is a multiple of the group size.Documented limitation: the second moment is not quantized
vis a running mean of squared gradients, so it spans orders of magnitude and sits near zero early. MLX's affinemx.quantizereconstructs poorly there,sqrt(v)lands in the update denominator, and the step explodes. Measured, both for int8 m+v and for fp32 m + int8 v:Fixing this needs a dynamic/exponential mapping with stochastic rounding — what bitsandbytes implements and MLX does not expose. A guard test asserts
vstays a single float32 array so this cannot be silently switched on.bitsis constrained to 8 with aValueErrorrather than shipped untested.group_sizeaccepts 32 / 64 / 128, all three exercised (finals 0.75205 / 0.75203 / 0.75202, all monotonic).Deliberate exclusions
paged_adamw_8bitandadamw_bnb_8bitkeep collapsing toadamw. "paged" promises CPU offload and "bnb" names a library MLX does not use, so routing them here would swap one wrong answer for another.Two test files, on purpose
The numeric tests need real
mx.quantize, so a single file would skip entirely on Linux CI and be a gate no-op — the #669/#739 pattern the behavioral-gates step exists to prevent.tests/test_mlx_quantized_optimizer_routing.pyruns under the existing torch shim and executes in the gate.tests/test_mlx_quantized_optimizer_state.pyholds the behaviour tests and runs on Apple Silicon.Tested
32 tests. Loss decreases over 8 steps; final loss within 1e-4 of fp32 on the same seed; asserted byte reduction; save → reload → resume reproduces the next step exactly; compiled and uncompiled bit-identical under
mx.compile(inputs=state, outputs=state); a model with no quantization-eligible parameter still trains;bits != 8and invalidgroup_sizerejected.Three of the routing tests fail against unpatched
trainer.py—adamw_8bitresolving toadamw,adam_8bitraisingUnsupported MLX optimizer, andadamw_8bitabsent from the advertised list.Full zoo suite run twice per side: totals 2979 → 3011, delta 32 = the tests added. No environment writes. No
Dataset.mappath is involved, so the forced-forkmatrix does not apply here.One caveat on the suite comparison:
tests/test_unsloth_zoo_lora_merge.pyandtests/test_moe_fused_narrow_expert_merge.pyare nondeterministic independent of this change — the same file run twice on unmodifiedmaingave 14 then 12 failures. Four of their tests appeared to regress on a two-run-per-side comparison and cleared on two additional baseline passes; neither file references MLX or the new module. Two runs per side is not always enough to separate signal from that file's noise.Overlap
Textual overlap with #819 on
trainer.py:56and the_build_optimizerchain — same author, semantically disjoint (the optimizers #819 adds hold no first moment). Whichever merges second rebases.