Feat/mps gemm 4bit bf16 - #2070
Closed
eaglstun wants to merge 23 commits into
Closed
Conversation
Adds the CPU-as-oracle parity net that gates the native Metal port
(docs/apple_silicon/PORT_PLAN.md §3):
- tests/test_mps_parity.py: 184 tests mirroring the test_ops.py style.
Bit-exactness asserts for quantize_blockwise/quantize_4bit codes,
documented per-dtype tolerances (fp32 1e-6/1e-5, fp16 1e-3/1e-2,
bf16 1e-2/4e-2) for dequant/matmul/optimizer parity, NF4/FP4 and
blockwise round-trip reconstruction across blocksize {64,128,256,512}
incl. partial-block tails, and loud-failure tests pinning the ops that
raise NotImplementedError on mps. Skips cleanly when MPS is absent.
- docs/apple_silicon/MPS_STATUS.md: op-by-op resolution matrix (mps reg
vs default reg vs unregistered), what actually runs on this machine
(Hub kernels inert: kernels pkg not installed despite macOS 26;
native csrc path confirmed dead; no silent CPU fallback possible),
measured baseline, and chosen tolerances.
Baseline on macOS 26.4.1 / torch 2.12.1: 183 passed, 1 xfailed (strict).
The xfail documents a real cross-backend divergence found by the sweep:
the default kernel (used on mps) applies coupled weight decay for lion,
while the cpu and CUDA kernels apply decoupled decay per the Lion paper
— the default backend is the outlier.
No Metal/C++/CMake changes; Phase 1 is audit-only by design.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The shared 32-bit optimizer path in the `default` backend folded weight decay into the gradient (coupled L2) for Lion, which corrupts the sign update. Lion uses *decoupled* (AdamW-style) weight decay per Chen et al. 2023: shrink the param directly (p *= 1 - lr*wd), outside the sign. This is what the cpu backend and the CUDA 8-bit blockwise kernel already do; the default backend (used by MPS and any device without a dedicated kernel) was the outlier. - Drop LION (id 4) from the coupled-decay group; keep MOMENTUM/RMSPROP/ADAGRAD. - Apply decoupled decay in the LION branch, matching cpu/ops.py. - Add test_lion32bit_weight_decay: the existing test_optimizer32bit never caught this because it constructs optimizers with the default weight_decay=0, where coupled and decoupled coincide. New test fails without the fix, passes with it. NOTE: the CUDA 32-bit kernel (csrc/kernels.cu::kOptimizer32bit1State) and the Triton 1-state kernel (backends/triton/kernels_optim.py) carry the same coupled- decay bug for Lion and need a separate upstream fix (they can't be built/tested on Apple Silicon). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…one op end to end)
Proves the entire native MPS pipe with a single op, per PORT_PLAN.md §3 Phase 2.
On a source build (cmake -DCOMPUTE_BACKEND=mps -S . -B . && cmake --build .),
bitsandbytes::quantize_blockwise on mps runs through a hand-written Metal kernel
and is BIT-EXACT vs the CPU oracle (codes AND absmax) across fp32/fp16/bf16 x
blocksize {64,128,256,512} incl. partial-block tails. Full parity suite on the
native build: 199 passed, 1 xfailed.
Validated the pre-existing kernel first (plan step 1): its scalar binary-search
math is correct (0/200k vs torch.bucketize), but the kernel as a whole was the
wrong shape for the op (no per-block absmax, no scaling, never writes absmax,
unrelated NUM_BLOCK grid-stride). Replaced, not reused.
- csrc/mps_kernels.metal: new quantize_blockwise kernel, one thread per block,
per-block absmax + reciprocal-scale + searchsorted-left over the code table's
255 midpoint bounds (reproduces torch.bucketize right=False). Correctness-first
one-thread-per-block; SIMD-group reduction deferred to a perf phase.
- CMakeLists.txt: compile the metallib with -fno-fast-math so division is
correctly rounded and no FMA contraction -> bucket selection matches CPU.
- csrc/mps_ops.mm: replaced the "Not implemented" stub with a real encode path
(cached device/queue/library/pipeline, commandBuffer -> encoder -> bind ->
dispatchThreads -> commit -> waitUntilCompleted). Stable extern "C" entry.
Install-safe metallib load via dladdr (next to the loaded dylib) + BNB_MPS_METALLIB
override, replacing the CWD-relative NSURL path (plan §4).
- cextension.py: MpsBNBNativeLibrary + get_mps_library() -- loads
libbitsandbytes_mps.dylib and confirms the metallib is present; returns None
(never raises) when absent.
- backends/mps/ops.py: route quantize_blockwise to native when loaded, else the
pure-torch fallback. Inputs forced to fresh offset-0 fp32 buffers, with
torch.mps.synchronize() before the (separate-queue) dispatch.
- tests: TestNativeMetalPath asserts the native path is bit-exact and, under
BNB_MPS_REQUIRE_NATIVE=1, fails hard if native didn't load; plus a
graceful-degradation test with the native handle forced off.
Buffer bridging needs no libtorch linkage: a torch MPS tensor's data_ptr() IS
its id<MTLBuffer> (verified: [length] == byte size, class AGXG16XFamilyBuffer),
so the .mm casts the ctypes void* directly. Ruled out bytesNoCopy (data_ptr not
page-aligned) and memcpy (data_ptr not CPU-readable). Details + open packaging
risk (metallib/dylib not yet force-included in a wheel) in MPS_STATUS.md §7.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tive Metal, and a load-time buffer-contract guard Graduates the remaining quant/dequant ops onto hand-written Metal, reusing the Phase-2 pipe, per PORT_PLAN.md §3 Phase 3. On a source build all four ops (quantize_blockwise from P2, plus these three) run through Metal and are BIT-EXACT vs the CPU oracle. Full parity suite with BNB_MPS_REQUIRE_NATIVE=1: 293 passed, 1 xfailed. Sub-task 0 (harden the foundation): the (__bridge id<MTLBuffer>) cast rides on an undocumented torch contract (MPS tensor data_ptr() == its id<MTLBuffer>). Added a cheap one-time load-time guard -- extern "C" bnb_mps_check_buffer_contract + MpsBNBNativeLibrary.verify_buffer_contract() -- that confirms a real MPS tensor's data_ptr() resolves to a genuine id<MTLBuffer> (protocol conformance + [length] check, @try/@catch). If a future torch breaks it, get_mps_library() disables native and logs a clear error (fall back, no crash, no silent corruption); BNB_MPS_REQUIRE_NATIVE=1 turns that into a hard test failure. Verified: real->1, null->0, oversize->0. Ops graduated (each bit-exact, each with native + graceful-fallback tests): - dequantize_blockwise: NEW mps registration (was missing -> fell to default). Kernel out[i]=code[A[i]]*absmax[i/bs]; Python casts fp32 out to dtype. - dequantize_4bit: native swap. High nibble->even index, low->odd; out[j]=code4[nib]*absmax[j/bs]. Also feeds gemv/gemm dequant (matmul still pure-torch F.linear -- fused matmul NOT started). - quantize_4bit: native swap. Per-block absmax + searchsorted over the 15 sorted-code midpoint bounds + order remap (FP4) + nibble pack. Storage-dtype reinterpret mirrors the reference (incl. quant_storage=bf16). Two reference subtleties reproduced for true bit-exactness (both latent risks): - Tail-block asymmetry: the reference stores the partial block's absmax CLAMPED and scales by DIRECT DIVIDE, while full blocks store raw and reciprocal-multiply (differ by up to 1 ulp under -fno-fast-math -> can flip a bucket). Hardened quantize_blockwise (P2 used reciprocal-multiply everywhere, passed only by luck) and quantize_4bit to branch on is_tail. - Odd-numel padding nibble: the reference pads scaled with 0.0 and QUANTIZES it (nonzero NF4/FP4 index) as the final low nibble -- kernel now quantizes 0.0 for that slot instead of writing a literal 0 (was a real 1-code mismatch until fixed). Tolerances: integer/packed outputs (codes, nibbles, absmax) asserted bit-exact (torch.equal / view-as-uint8 to dodge NaN!=NaN on bf16 reinterpret); float dequant outputs use the Phase-1 per-dtype tolerances but measure bit-exact. Unchanged debt (not regressed): per-call offset-0 input copy; wheel-packaging gap (metallib/dylib not force-included in a wheel). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…kwise + 4-bit Brings quantize_blockwise, dequantize_blockwise, quantize_4bit, dequantize_4bit onto hand-written Metal kernels on the mps backend, all bit-exact vs the CPU oracle, with a load-time buffer-contract guard and CPU-parity test harness. gemv_4bit/gemm_4bit fused matmul and wheel packaging remain follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…s from pip install
The package-data glob "libbitsandbytes*.*" catches every shared library (all
prefixed "lib...", including libbitsandbytes_mps.dylib) but MISSES
bitsandbytes.metallib, which has no "lib" prefix. A wheel built before this fix
carried the dylib but not the shader archive, so get_mps_library() found the
dylib, failed the metallib.exists() gate, and silently fell back to pure-torch.
Fix (packaging only, one line): add a "*.metallib" entry to package-data.
package-data = { "*" = ["libbitsandbytes*.*", "*.metallib", "py.typed"] }
Verified the .dylib is genuinely covered by the existing glob (checked the built
wheel, not assumed). Build flow (matches how bnb ships prebuilt CUDA .so's):
cmake-build the MPS artifacts into bitsandbytes/ first, then
`BNB_SKIP_CMAKE=1 python -m build --wheel` to package the pre-built artifacts
without re-running cmake (BNB_SKIP_CMAKE is the real switch for the
scikit_build_core.setuptools shim; wheel.cmake=false only applies to the native
backend). Clean build/ first or a stale cpu dylib gets swept in.
Evidence:
- unzip -l dist/*.whl shows BOTH bitsandbytes/bitsandbytes.metallib and
bitsandbytes/libbitsandbytes_mps.dylib at the correct in-package path.
- Isolated throwaway venv (pip install --no-deps the wheel, run from outside the
worktree so import resolves to the INSTALLED package): both artifacts present,
get_mps_library() loads native with metallib_path resolved via dladdr inside
the venv, verify_buffer_contract() passes, native quantize_blockwise bit-exact
vs CPU, and the parity subset runs 236 passed with BNB_MPS_REQUIRE_NATIVE=1.
No code/kernel changes. MPS_STATUS.md §9 documents the change + verification.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Docs only; no code/kernel/build changes. Every supported/unsupported claim is derived from docs/apple_silicon/MPS_STATUS.md ground truth. - README.md accelerator table (Metal / mps row): corrected to reality using the table's existing legend. LLM.int8() ❌ (int8_double_quant raises NotImplementedError on mps; no native int8 path), QLoRA 4-bit 🐢 (4-bit quant/dequant native + bit-exact, but the matmul is unfused: dequant-through- Metal + F.linear -- "Slow Implementation Supported"), 8-bit Optimizers ❌ (optimizer_update_8bit_blockwise raises NotImplementedError, not in development). Also closed the pre-existing missing </tr> on that row. - docs/source/installation.mdx: new "Apple Silicon / Metal (Preview)" section modeled on "AMD ROCm (Preview)" -- requirements (macOS 14+, M1+, PyTorch>=2.4 with MPS), compile-from-source recipe (cmake -DCOMPUTE_BACKEND=mps -> build -> pip install), the in-source-build metallib note, and that native MPS now ships in the built wheel. Added to the table of contents. Honest about what's not supported (int8, 8-bit optimizers) and that 4-bit matmul is unfused. - docs/apple_silicon/README.md: concise user/dev guide -- status at a glance, requirements, install/build, supported-op matrix (native vs fallback vs unsupported), numerics/parity, known limitations, pointers to PORT_PLAN.md / MPS_STATUS.md / NEXT_MATMUL_PLAN.md. - docs/apple_silicon/NEXT_MATMUL_PLAN.md: executable spec (PORT_PLAN.md-style) teeing up native 4-bit matmul fusion cold -- current state, the design fork (MPSMatrixMultiplication on dequantized B vs a hand-fused dequant+matmul kernel) with trade-offs and a recommended per-op split, phased plan, CPU-as-oracle parity approach reusing the existing harness, what to reuse from the built pipe, and the per-call offset-0 input-copy debt to also close. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Benchmarked today's unfused gemv_4bit/gemm_4bit (native dequant of B -> F.linear) on this machine to ground the NEXT_MATMUL_PLAN.md design fork with real numbers instead of intuition. Findings (MPS_STATUS.md §10): - gemv (M=1) is 80-90% dequant-bound; the GEMM is ~10%. Option A (MPSMatMul on materialized B_dq) can only touch that 10% -> fusion (Option B) is the win. - gemm has a fixed ~0.75ms dequant floor; the GEMM overtakes it near M~512. Large-M gemm -> Option A (don't hand-beat Apple's tuned GEMM). - The existing dequant kernel runs at ~54 GB/s vs ~400 GB/s achievable, so the M2 target is memory bandwidth, not "fusion" per se: a fused gemv that reads packed B no faster than today wins nothing. Decision: Option B (fused) for gemv_4bit = Phase M2; Option A (MPSMatrixMultiplication) for large-M gemm_4bit = Phase M3. Matches the plan's recommended split, now evidence-backed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… no B_dq) One SIMD-group (32 threads) per output element n: threads stride the packed row in uint4 units (16 B = 32 nibbles), dequantize nf4/fp4 in registers (code[nib] * absmax, rounded to the activation dtype to match the oracle's B_dq.to(dtype)), and accumulate the dot product in fp32 with split accumulators + explicit fma, reduced via simd_sum. Per-dtype kernel variants (fp32/fp16/bf16) bind A and out in torch's own dtype, so the steady-state call launches zero torch cast kernels. The dequantized B is never written to device memory. - csrc/mps_kernels.metal: gemv_4bit_body<T> + fp32/fp16/bf16 kernels - csrc/mps_ops.mm: bnb_mps_gemv_4bit dispatch (own queue, blocking wait, per-dtype pipeline; BNB_MPS_PROFILE=1 logs kernel-only GPU time) - cextension.py: argtypes (guarded by hasattr so stale dylibs keep working) - backends/mps/ops.py: route _gemv_4bit_impl to native when available; guards (M==1, K % 32 == 0, pow2 blocksize, 16-entry code, packed-size check) fall back to the unchanged dequant + F.linear path - tests: fused-native spy test (asserts native is actually hit), unaligned-K fallback test, gemv added to the graceful-fallback matrix Parity: tests/test_mps_parity.py -k gemv all green under BNB_MPS_REQUIRE_NATIVE=1 (nf4+fp4, bs 64/256, fp32/fp16/bf16). Speedup vs dequant+F.linear baseline (nf4/bs64, M1 shapes): 3.4-6.2x wall-clock (fp16 4096x4096: 0.31ms vs 1.64ms). Kernel-only ~0.11ms on the ~25 MB shapes when clocked up (~230 GB/s read vs the standalone dequant kernel's ~54 GB/s); per-call cross-queue sync now dominates wall-clock. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…on, one command buffer Option A from the M1 decision: bnb_mps_gemm_4bit (csrc/mps_ops.mm) encodes, on ONE command buffer with ONE commit + ONE blocking wait, (1) a chunked 4-bit dequant kernel (dequantize_4bit_chunked_fp32/fp16: one thread per 32-element uint4 chunk, writing (T)(code[nib]*absmax) into a growable private scratch MTLBuffer -- the same rounding as the oracle's B_dq.to(dtype)), (2) a shape-cached MPSMatrixMultiplication computing A[M,K] . B_dq[N,K]^T (row-major, transposeRight), and (3) an optional bias epilogue kernel out[m,n] += bias[n]. The single sync is the structural win: the previous dequant + F.linear tail paid the ~0.15-0.25ms cross-queue sync twice per call. bf16 is excluded by the Python router: MPSMatrixMultiplication hard-asserts on anything but fp32/fp16/int8/int16 (probed on macOS 26.4.1), so bf16 keeps the existing dequant + F.linear fallback verbatim. Other guards mirror the fused gemv (K % 32 == 0, power-of-two blocksize >= 32, packed-size check, hasattr guard for stale dylibs); nested absmax is still unpacked to plain fp32 absmax before routing, unchanged. Parity: tests/test_mps_parity.py -k "gemm_4bit or gemv" = 69 passed under BNB_MPS_REQUIRE_NATIVE=1. Native path asserted via spy (fp32/fp16 x nf4/fp4 x +/-bias x +/-nested absmax); bf16 and K%32!=0 fallbacks asserted; graceful fallback with the native handle off covered. fp32 vs MPSMatMul accumulation stays within the documented 1e-5 atol at the calibrated K <= 256 (the one trip found was in the PURE-TORCH fallback composition at K=256/M=4; the fallback test pins K=64 like the main gemm test -- the documented tolerance is unchanged). Speedup vs the dequant+F.linear fallback (nf4/bs64, N=K=4096, 30 iters, benchmarks_wip/bench_gemm_baseline.py): fp16 2.5x (M=8), 1.5x (M=64/512), 1.08x (M=2048); fp32 1.6x/1.5x (M=8/64), 1.1x (M=512), ~1.0x (M=2048). Win = one sync + a much faster chunked dequant at small/medium M; flat at M=2048 where the GEMM dominates and MPSMatMul ~= F.linear. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… semantics verified, honest docs Task 1 (per-call sync overhead): measured and decomposed, NOT optimized away -- torch exposes no safe handle to its MPS queue/stream, so the private-queue + sync design stays, now with the cost precisely documented (MPS_STATUS.md 11.3). - BNB_MPS_PROFILE=1 now decomposes each native matmul call into sched (commit -> GPU start), gpu (kernel), and done (GPU end -> wait return), on the mach_absolute_time timebase GPUStartTime uses. Steady-state fp16: the fixed tax is ~0.15ms/call (~0.02 encode + ~0.07 sched + ~0.07 done); the pre-call torch.mps.synchronize() itself is ~1us on an idle queue. That tax is ~52% of a 4096x4096 gemv call and ~10% of an M=512 gemm. - Queue sharing rejected with evidence: torch 2.12.1 exposes no stream handle from Python; libtorch's at::mps::getCurrentMPSStream()/MPSStream internals are reachable only via dlsym'd C++ mangled symbols + header-derived ivar offsets (commandQueue()/queue() are inline, commit()/flush() private, encodes must run on torch's private _serialQueue) -- an unguardable ABI trap. Sharing only the queue would not even remove the pre-sync (command buffers execute in commit order; torch batches into an uncommitted buffer). Future directions (libtorch-linked extension, torch.mps.compile_shader) recorded with their costs in 11.3. - New test_sync_discipline_interleave_stress pins the discipline: dependent torch writes interleaved with native matmuls on the same buffers, parity every iteration. Verified to have teeth: fails 30/30 with the pre-sync no-op'd, passes 100% with it. Task 2 (offset-0 input copy): verified on this build, clone stays. - A view's data_ptr() IS base + storage_offset*itemsize (raw arithmetic); objc-probing such an interior pointer SIGSEGVs uncatchably, so casting it would be silent corruption or a crash -- the clone in _ensure_native_buffer is load-bearing. - Better than NEXT_MATMUL_PLAN 5 feared: untyped_storage().data_ptr() DOES recover the base MTLBuffer (passes the contract check), so offset binding is possible; not implemented because steady-state matmul inputs are offset-0 (the clone ~never fires) and it would touch every ABI entry point for no measured benefit. Recipe recorded in 11.4; semantics pinned by test_view_data_ptr_is_base_plus_offset. Task 3 (docs): - MPS_STATUS.md: M2/M3 short notes consolidated into a proper 11 (fused gemv, native gemm, the M4 sync/offset findings); op matrix rows for gemv/gemm updated to native; header records the M4 harness result (328 passed + the known unrelated lion strict-xfail XPASS). - docs/apple_silicon/README.md: gemv_4bit -> native (fused, 3.4-6.2x); gemm_4bit -> native for fp16/fp32 with the honest caveats (bf16 falls back -- MPSMatrixMultiplication has no bf16 on macOS 26.4.1; large-M ~ parity with F.linear); Known limitations rewritten around the real remaining costs. - README.md accelerator table, mps QLoRA 4-bit: 🐢 -> ✅ with a footnote. Justification: "slow implementation" is no longer accurate (inference gemv is fused native Metal at 3.4-6.2x; gemm is native fp16/fp32), but a bare ✅ would overclaim, so the footnote carries the bf16-fallback and large-M-parity caveats inline in the README. - "mis-cast" -> "miscast" (typos hook); _typos.toml learns the mach timebase field `numer` so the hook stops rewriting it to `number`. Parity gate: BNB_MPS_REQUIRE_NATIVE=1 pytest tests/test_mps_parity.py -> 328 passed; the only failure is the pre-existing, unrelated lion strict-xfail XPASS that fails identically on clean HEAD (separate optimizer track). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Brings the native Metal 4-bit matmul path onto main, following the earlier quantize/dequantize phases. - M1: benchmarked the unfused baseline; decided the design fork with numbers (gemv is dequant-bound -> fuse; large-M gemm -> MPSMatrixMultiplication). - M2: fused native gemv_4bit MSL kernel (dequant in registers, no B_dq), 3.5-7x. - M3: native gemm_4bit = chunked dequant + MPSMatrixMultiplication on one command buffer (bf16 falls back -- MPSMatMul has no bf16 on macOS 26.4.1). - M4: measured the per-call sync tax (~0.15ms, architectural -- documented, not removed; queue-sharing rejected with evidence), verified the offset-0 clone is load-bearing, and updated docs honestly (README QLoRA-4bit row -> checkmark with a footnote on the bf16/large-M caveats). 328 parity tests green under BNB_MPS_REQUIRE_NATIVE=1; native paths asserted via spy tests; graceful fallback preserved. LLM.int8 and 8-bit optimizers remain out of scope on mps. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…fixed test_lion_weight_decay_backend_divergence was a strict xfail pinning a real cross-backend bug: the `default` optimizer kernel (used on mps) applied COUPLED weight decay for Lion while the cpu/CUDA kernels applied DECOUPLED decay. That bug was fixed upstream (bitsandbytes-foundation#1992 / bitsandbytes-foundation#1993) -- the default backend now excludes LION from the coupled fold and shrinks the param directly -- so mps and the cpu oracle agree and the strict xfail was XPASSing (counting as a suite failure). - Remove the xfail; rename to test_lion_weight_decay_decoupled_parity and document it as the regression guard for the fix. - Update the stale "backends disagree" note in test_optimizer_update_32bit. - MPS_STATUS.md: mark the §5 finding resolved, fix the §2 matrix row, update the harness header (329 passed, 0 xfailed). Full parity suite: 329 passed, 0 xfailed under BNB_MPS_REQUIRE_NATIVE=1 (native lib built in this checkout). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Follow-up to bitsandbytes-foundation#1993, which fixed the same coupled-vs-decoupled weight-decay bug for Lion in the default backend and the CUDA 32-bit kernel. matthewdouglas invited a Triton follow-up when approving bitsandbytes-foundation#1993. _optimizer_update_1state_32bit_triton_kernel folded weight decay into the gradient (coupled L2) for every 1-state optimizer, Lion included. Lion requires *decoupled* (AdamW-style) decay applied to the param directly (p *= 1 - lr*wd), outside the sign update (Chen et al. 2023); folding decay into the gradient corrupts the input to sign(), changing the update direction, not just its magnitude. This mirrors the CUDA 32-bit fix exactly: exclude LION (id 4) from the gradient fold, and shrink the param in the LION branch before the sign update. The existing test_lion32bit_weight_decay (added in bitsandbytes-foundation#1993) is parametrized over get_available_devices(), so it already exercises this kernel on XPU hardware in CI. Like bitsandbytes-foundation#1993's CUDA change, the Triton kernel is XPU-only and cannot be built or run on the available hardware. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Follow-up to bitsandbytes-foundation#1993 and the Triton 32-bit fix. _optimizer_update_1state_8bit_ blockwise_triton_kernel gated its decoupled-decay branch on OPTIMIZER_ID == 2, which is ADAGRAD, not LION (id 4) -- a copy/paste error on the id, right under a `# LION` comment. The consequences: - Lion (id 4) fell through to the coupled (L2) `elif` and folded weight decay into the gradient before sign(), corrupting the update direction. Lion needs *decoupled* decay applied to the param directly (Chen et al. 2023). - Adagrad (id 2) wrongly got Lion's decoupled decay instead of the coupled L2 fold it expects (and that the cpu backend applies for it). Correcting the id to 4 fixes both optimizers with a one-character change: Lion gets decoupled decay, Adagrad falls through to the coupled fold. Adds test_lion8bit_blockwise_weight_decay, the 8-bit companion to test_lion32bit_weight_decay (bitsandbytes-foundation#1993). It runs Lion8bit against the lion-pytorch reference with weight_decay=0.1 and no per-step resync, so the coupled-vs- decoupled divergence accumulates into a reliable signal. Validated by injecting the coupled bug into the CPU 8-bit kernel: the test fails under the bug and passes with decoupled decay (fp32 separates ~160x, fp16 ~15x). bf16 is excluded -- its mantissa rounds the bug's per-element update difference away, so it cannot resolve coupled vs decoupled regardless of correctness. MPS is skipped (the 8-bit blockwise op is unimplemented there). Like bitsandbytes-foundation#1993's CUDA change and the Triton 32-bit fix, this kernel is XPU-only and cannot be built or run on the available hardware. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Mirror of the merged upstream fix (bitsandbytes-foundation#1993): kOptimizer32bit1State folded weight decay into the gradient (coupled L2) for every optimizer including Lion, which corrupts Lion's sign update. Exclude LION from the coupled fold and apply decoupled decay (p *= 1 - lr*wd) in the LION branch, matching the 8-bit kernel and the cpu/default backends. The default backend + test_lion32bit_weight_decay are already on this fork's main; this adds the remaining CUDA kernel piece. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fix Lion decoupled weight decay in the Triton 32-bit kernel
Fix Lion (and Adagrad) decoupled weight decay in the Triton 8-bit blockwise kernel
Fix Lion decoupled weight decay in the CUDA 32-bit kernel
Brings in 6 upstream commits (through bitsandbytes-foundation#2004). Conflict resolutions: - csrc/mps_kernels.metal, csrc/mps_ops.mm: upstream bitsandbytes-foundation#2004 DELETED these as 'unused metal build code' (they were stubs upstream); the fork built them into the real MPS port. Kept the fork's versions (rejected the deletion). - CMakeLists.txt: bitsandbytes-foundation#2004 removed upstream's stub metal build. Restored the fork's full MPS build support that the auto-merge dropped: the COMPUTE_BACKEND=mps -> BUILD_MPS selection branch, set(MPS_FILES)/set(METAL_FILES), and the late if(BUILD_MPS) framework-linking + add_dependencies(metallib) block. Verified a clean MPS build produces the metallib + _mps.dylib. - bitsandbytes/backends/default/ops.py: dropped the stale Lion NOTE (claimed CUDA and Triton still had the coupled-decay bug -- all three are now fixed on this fork); took upstream's cleaner version. The decoupled-decay code is unchanged. - tests/test_optim.py: kept the fork's test_lion8bit_blockwise_weight_decay alongside upstream's test_lion32bit_weight_decay. - csrc/pythonInterface.cpp: took upstream's removal of a commented-out dead #include (fork's MPS loads via dladdr, doesn't use it). Verified on the merged tree: MPS builds, tests/test_mps_parity.py 329 passed under BNB_MPS_REQUIRE_NATIVE=1 (bit-exact quant/dequant + native matmul intact). Pre-existing paged_lion-on-mps test_optim failures are unchanged (paging is CUDA-only; identical on fork main pre-merge). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # README.md # docs/source/installation.mdx
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.
No description provided.