experts4bit-qlora 0.1.0__tar.gz

experts4bit_qlora-0.1.0/LICENSE

@@ -0,0 +1,29 @@
+MIT License
+
+Copyright (c) 2026 Jordan Anderson
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+experts4bit_qlora/_vendor/experts.py is vendored from bitsandbytes
+(https://github.com/bitsandbytes-foundation/bitsandbytes), which is likewise
+distributed under the MIT License. It is included here only while the upstream
+Experts4bit proposal (bitsandbytes#1965) is unmerged, and will be removed once
+the class ships in a bitsandbytes release.

experts4bit_qlora-0.1.0/PKG-INFO

@@ -0,0 +1,147 @@
+Metadata-Version: 2.4
+Name: experts4bit-qlora
+Version: 0.1.0
+Summary: QLoRA fine-tuning of fused 4-bit Mixture-of-Experts on a single small GPU, on stock bitsandbytes.
+Author-email: Jordan Anderson <paul.jordan.anderson@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/pjordanandrsn/experts4bit-qlora
+Project-URL: Upstream (bitsandbytes#1965), https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1965
+Project-URL: Tracking (bitsandbytes#1849), https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1849
+Keywords: bitsandbytes,qlora,moe,mixture-of-experts,quantization,nf4,lora,olmoe
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.2
+Requires-Dist: bitsandbytes>=0.43
+Provides-Extra: train
+Requires-Dist: transformers>=5.0; extra == "train"
+Requires-Dist: datasets>=2.14; extra == "train"
+Requires-Dist: accelerate>=0.30; extra == "train"
+Requires-Dist: safetensors>=0.4; extra == "train"
+Requires-Dist: huggingface_hub>=0.23; extra == "train"
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Requires-Dist: transformers>=5.0; extra == "test"
+Requires-Dist: accelerate>=0.30; extra == "test"
+Requires-Dist: safetensors>=0.4; extra == "test"
+Dynamic: license-file
+
+# experts4bit-qlora
+
+[![CI](https://github.com/pjordanandrsn/experts4bit-qlora/actions/workflows/ci.yml/badge.svg)](https://github.com/pjordanandrsn/experts4bit-qlora/actions/workflows/ci.yml)
+
+QLoRA fine-tuning of **fused Mixture-of-Experts** weights on a single small GPU — the part that
+doesn't fit anywhere else yet.
+
+## The problem
+
+transformers v5 stores MoE experts as one fused 3-D `nn.Parameter` per layer
+(`OlmoeExperts`, `Qwen3MoeExperts`, …). bitsandbytes' 4-bit walker only replaces `nn.Linear`
+modules, so it **silently skips the experts** — which are the overwhelming majority of a MoE's
+weights. `load_in_4bit` "shrinks" the model but the experts stay in full precision
+([bitsandbytes#1849](https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1849)).
+
+`Experts4bit` is the primitive that 4-bit-quantizes exactly that fused stack. This package pairs
+it with a **streaming loader** and **per-expert LoRA**, so you can actually *fine-tune* a real
+sparse-MoE on reasonable hardware.
+
+## What it buys you (measured on an RTX A2000 12 GB)
+
+- **It fits at all.** Full bf16 OLMoE-1B-7B is ~13.9 GB — it **OOMs** on a 12 GB card. In 4-bit
+  it loads at **4.69 GB** and trains in <8 GB. The streaming loader never materializes the bf16
+  model in CPU *or* GPU RAM (verified under a 3 GB container RAM cap).
+- **It trains.** QLoRA on the frozen NF4 experts improves a held-out Alpaca eval from
+  **1.4813 → 1.0290** (see [`docs/METHODOLOGY.md`](docs/METHODOLOGY.md)).
+- **Honest caveat — this is a memory technology, not an energy one.** On a GPU that *already*
+  fits the model, 4-bit is a **1.2–2.3× energy penalty** (NF4 is storage-only; the GEMM runs in
+  bf16 either way, plus dequant). The energy win only shows up when memory is the binding
+  constraint — then it's the difference between running and not, and up to **4.4× lower
+  energy/token** from the batch that freed memory unlocks. Numbers and method in the docs.
+
+## Install
+
+```bash
+# primitive + adapters + benchmarks (torch + bitsandbytes):
+pip install "git+https://github.com/pjordanandrsn/experts4bit-qlora"
+# + the streaming MoE trainer (transformers>=5.0, datasets, ...):
+pip install "experts4bit-qlora[train] @ git+https://github.com/pjordanandrsn/experts4bit-qlora"
+```
+
+Runs on a **stock** `pip install bitsandbytes` today — see "Relationship to bitsandbytes" below.
+
+## Quickstart
+
+```python
+import torch
+from experts4bit_qlora import Experts4bit, ExpertsLoRA
+
+# Freeze a fused expert stack in 4-bit, attach trainable per-expert LoRA.
+gate_up = torch.randn(8, 2 * 256, 128)          # [num_experts, 2*intermediate, hidden]
+down    = torch.randn(8, 128, 256)              # [num_experts, hidden, intermediate]
+base    = Experts4bit.from_float(gate_up, down, quant_type="nf4", compute_dtype=torch.float32)
+model   = ExpertsLoRA(base, r=8, alpha=16)      # only the LoRA adapters train
+```
+
+End-to-end OLMoE QLoRA fine-tune (needs a CUDA GPU + `[train]` extras):
+
+```bash
+STEPS=150 R=8 TRAIN_EXPERTS=1 TRAIN_ATTENTION=0 OUT=./out \
+  python -m experts4bit_qlora.train
+```
+
+## Scope
+
+The `Experts4bit` primitive and `ExpertsLoRA` adapters are **model-agnostic**. The **streaming loader /
+trainer** (`python -m experts4bit_qlora.train`) supports fused-MoE architectures that store experts
+per-expert on disk under `model.layers.{i}.mlp.experts.{e}.{gate,up,down}_proj.weight` with a SwiGLU gate:
+
+- **OLMoE** (OLMoE-1B-7B) — convergence-tested end-to-end; fits a 12 GB card at ~4.7 GB.
+- **Qwen3-MoE / Qwen3.5-MoE** — same checkpoint + module layout (verified byte-for-byte identical to
+  OLMoE's on-disk format); structurally tested in `tests/test_loader_architectures.py`. The real weights
+  (30–35B) need a ≥24 GB card — or the CPU-offloading path (tracked separately) — to fit 12 GB.
+
+Anything else **fails fast with a clear error**. **Gemma 4** is a genuinely different design (experts at
+`layers.{i}.experts` beside a parallel dense MLP, with a custom router) and needs its own loader
+adaptation — not yet supported. PRs welcome.
+
+## Benchmarks
+
+```bash
+# Runs on stock bitsandbytes (uses the portable dequantize forward):
+python bench/bench_energy_excluded.py                    # memory wall + tokens-per-joule vs batch
+
+# Require bitsandbytes >= 0.50 — measure the upstream matmul_4bit optimization (#1965):
+python bench/_upstream/bench_matmul4bit.py --mode both   # equivalence + latency/memory
+python bench/_upstream/bench_energy.py                   # joules/op: bf16 vs dequant vs matmul_4bit
+```
+
+The LoRA-placement ablation (which of experts / attention / router to train) and full energy
+analysis are written up in [`docs/METHODOLOGY.md`](docs/METHODOLOGY.md). Short version: on Alpaca
+the placements are largely **redundant**, attention-only is the efficiency pick, and training the
+router **hurts**.
+
+## Relationship to bitsandbytes
+
+`Experts4bit` is a bitsandbytes primitive, proposed upstream in
+[bitsandbytes#1965](https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1965). Until it
+ships in a release, this package **vendors** a copy (`experts4bit_qlora/_vendor/experts.py`) so it
+runs on stock bitsandbytes today. The import shim prefers the upstream class when present:
+
+```python
+try:
+    from bitsandbytes.nn import Experts4bit      # once bitsandbytes#1965 releases
+except ImportError:
+    from ._vendor.experts import Experts4bit     # vendored fallback (stock bnb)
+```
+
+The vendored forward also **auto-detects** whether `matmul_4bit` is correct on your installed
+bitsandbytes — it only handles this weight layout correctly on **bnb ≥ 0.50**, so on older releases
+the primitive uses the portable dequantize path, and the `matmul_4bit` memory optimization engages
+automatically once you upgrade. Results are correct on any supported bnb either way.
+
+When it lands upstream: bump the `bitsandbytes` floor and delete `_vendor/` — no API change.
+
+## License
+
+MIT (see [LICENSE](LICENSE)). `experts4bit_qlora/_vendor/experts.py` is vendored from
+bitsandbytes (also MIT) pending upstream merge.

experts4bit_qlora-0.1.0/README.md

@@ -0,0 +1,119 @@
+# experts4bit-qlora
+
+[![CI](https://github.com/pjordanandrsn/experts4bit-qlora/actions/workflows/ci.yml/badge.svg)](https://github.com/pjordanandrsn/experts4bit-qlora/actions/workflows/ci.yml)
+
+QLoRA fine-tuning of **fused Mixture-of-Experts** weights on a single small GPU — the part that
+doesn't fit anywhere else yet.
+
+## The problem
+
+transformers v5 stores MoE experts as one fused 3-D `nn.Parameter` per layer
+(`OlmoeExperts`, `Qwen3MoeExperts`, …). bitsandbytes' 4-bit walker only replaces `nn.Linear`
+modules, so it **silently skips the experts** — which are the overwhelming majority of a MoE's
+weights. `load_in_4bit` "shrinks" the model but the experts stay in full precision
+([bitsandbytes#1849](https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1849)).
+
+`Experts4bit` is the primitive that 4-bit-quantizes exactly that fused stack. This package pairs
+it with a **streaming loader** and **per-expert LoRA**, so you can actually *fine-tune* a real
+sparse-MoE on reasonable hardware.
+
+## What it buys you (measured on an RTX A2000 12 GB)
+
+- **It fits at all.** Full bf16 OLMoE-1B-7B is ~13.9 GB — it **OOMs** on a 12 GB card. In 4-bit
+  it loads at **4.69 GB** and trains in <8 GB. The streaming loader never materializes the bf16
+  model in CPU *or* GPU RAM (verified under a 3 GB container RAM cap).
+- **It trains.** QLoRA on the frozen NF4 experts improves a held-out Alpaca eval from
+  **1.4813 → 1.0290** (see [`docs/METHODOLOGY.md`](docs/METHODOLOGY.md)).
+- **Honest caveat — this is a memory technology, not an energy one.** On a GPU that *already*
+  fits the model, 4-bit is a **1.2–2.3× energy penalty** (NF4 is storage-only; the GEMM runs in
+  bf16 either way, plus dequant). The energy win only shows up when memory is the binding
+  constraint — then it's the difference between running and not, and up to **4.4× lower
+  energy/token** from the batch that freed memory unlocks. Numbers and method in the docs.
+
+## Install
+
+```bash
+# primitive + adapters + benchmarks (torch + bitsandbytes):
+pip install "git+https://github.com/pjordanandrsn/experts4bit-qlora"
+# + the streaming MoE trainer (transformers>=5.0, datasets, ...):
+pip install "experts4bit-qlora[train] @ git+https://github.com/pjordanandrsn/experts4bit-qlora"
+```
+
+Runs on a **stock** `pip install bitsandbytes` today — see "Relationship to bitsandbytes" below.
+
+## Quickstart
+
+```python
+import torch
+from experts4bit_qlora import Experts4bit, ExpertsLoRA
+
+# Freeze a fused expert stack in 4-bit, attach trainable per-expert LoRA.
+gate_up = torch.randn(8, 2 * 256, 128)          # [num_experts, 2*intermediate, hidden]
+down    = torch.randn(8, 128, 256)              # [num_experts, hidden, intermediate]
+base    = Experts4bit.from_float(gate_up, down, quant_type="nf4", compute_dtype=torch.float32)
+model   = ExpertsLoRA(base, r=8, alpha=16)      # only the LoRA adapters train
+```
+
+End-to-end OLMoE QLoRA fine-tune (needs a CUDA GPU + `[train]` extras):
+
+```bash
+STEPS=150 R=8 TRAIN_EXPERTS=1 TRAIN_ATTENTION=0 OUT=./out \
+  python -m experts4bit_qlora.train
+```
+
+## Scope
+
+The `Experts4bit` primitive and `ExpertsLoRA` adapters are **model-agnostic**. The **streaming loader /
+trainer** (`python -m experts4bit_qlora.train`) supports fused-MoE architectures that store experts
+per-expert on disk under `model.layers.{i}.mlp.experts.{e}.{gate,up,down}_proj.weight` with a SwiGLU gate:
+
+- **OLMoE** (OLMoE-1B-7B) — convergence-tested end-to-end; fits a 12 GB card at ~4.7 GB.
+- **Qwen3-MoE / Qwen3.5-MoE** — same checkpoint + module layout (verified byte-for-byte identical to
+  OLMoE's on-disk format); structurally tested in `tests/test_loader_architectures.py`. The real weights
+  (30–35B) need a ≥24 GB card — or the CPU-offloading path (tracked separately) — to fit 12 GB.
+
+Anything else **fails fast with a clear error**. **Gemma 4** is a genuinely different design (experts at
+`layers.{i}.experts` beside a parallel dense MLP, with a custom router) and needs its own loader
+adaptation — not yet supported. PRs welcome.
+
+## Benchmarks
+
+```bash
+# Runs on stock bitsandbytes (uses the portable dequantize forward):
+python bench/bench_energy_excluded.py                    # memory wall + tokens-per-joule vs batch
+
+# Require bitsandbytes >= 0.50 — measure the upstream matmul_4bit optimization (#1965):
+python bench/_upstream/bench_matmul4bit.py --mode both   # equivalence + latency/memory
+python bench/_upstream/bench_energy.py                   # joules/op: bf16 vs dequant vs matmul_4bit
+```
+
+The LoRA-placement ablation (which of experts / attention / router to train) and full energy
+analysis are written up in [`docs/METHODOLOGY.md`](docs/METHODOLOGY.md). Short version: on Alpaca
+the placements are largely **redundant**, attention-only is the efficiency pick, and training the
+router **hurts**.
+
+## Relationship to bitsandbytes
+
+`Experts4bit` is a bitsandbytes primitive, proposed upstream in
+[bitsandbytes#1965](https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1965). Until it
+ships in a release, this package **vendors** a copy (`experts4bit_qlora/_vendor/experts.py`) so it
+runs on stock bitsandbytes today. The import shim prefers the upstream class when present:
+
+```python
+try:
+    from bitsandbytes.nn import Experts4bit      # once bitsandbytes#1965 releases
+except ImportError:
+    from ._vendor.experts import Experts4bit     # vendored fallback (stock bnb)
+```
+
+The vendored forward also **auto-detects** whether `matmul_4bit` is correct on your installed
+bitsandbytes — it only handles this weight layout correctly on **bnb ≥ 0.50**, so on older releases
+the primitive uses the portable dequantize path, and the `matmul_4bit` memory optimization engages
+automatically once you upgrade. Results are correct on any supported bnb either way.
+
+When it lands upstream: bump the `bitsandbytes` floor and delete `_vendor/` — no API change.
+
+## License
+
+MIT (see [LICENSE](LICENSE)). `experts4bit_qlora/_vendor/experts.py` is vendored from
+bitsandbytes (also MIT) pending upstream merge.

experts4bit_qlora-0.1.0/experts4bit_qlora/__init__.py

@@ -0,0 +1,18 @@
+"""experts4bit-qlora — QLoRA fine-tuning of fused 4-bit Mixture-of-Experts on a single small GPU.
+
+``Experts4bit`` resolves to the upstream bitsandbytes class once it ships in a release
+(bitsandbytes#1965); until then it falls back to a vendored copy, so this package works on a
+stock ``pip install bitsandbytes`` today. When upstream releases it, bump the bitsandbytes floor
+and delete ``experts4bit_qlora/_vendor/`` — no API change for callers.
+"""
+
+try:
+    # Upstream (bitsandbytes#1965) once released; else the vendored copy on stock bitsandbytes.
+    from bitsandbytes.nn import Experts4bit
+except ImportError:
+    from ._vendor.experts import Experts4bit
+
+from .lora import ExpertsLoRA, LoRALinear, add_attention_lora
+
+__all__ = ["Experts4bit", "ExpertsLoRA", "LoRALinear", "add_attention_lora"]
+__version__ = "0.1.0"

experts4bit_qlora-0.1.0/experts4bit_qlora/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Vendored bitsandbytes primitives (temporary — see experts.py header for provenance)."""

experts4bit_qlora-0.1.0/experts4bit_qlora/_vendor/experts.py

@@ -0,0 +1,355 @@
+# Vendored from bitsandbytes (github.com/bitsandbytes-foundation/bitsandbytes), pending the
+# upstream merge of Experts4bit (bitsandbytes#1965). While unmerged, this copy lets
+# experts4bit_qlora run on a stock `pip install bitsandbytes`; it depends only on released
+# bitsandbytes public APIs (quantize_4bit, dequantize_4bit, QuantState, get_4bit_type,
+# matmul_4bit). Once Experts4bit ships in a bitsandbytes release, experts4bit_qlora imports it
+# from bitsandbytes.nn directly (see the shim in experts4bit_qlora/__init__.py) and this file
+# is deleted.
+#
+# Original author: Jordan Anderson <paul.jordan.anderson@gmail.com>. Licensed under MIT.
+from collections.abc import Callable
+from typing import Optional
+
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F_nn
+
+import bitsandbytes as bnb
+import bitsandbytes.functional as F
+from bitsandbytes.functional import QuantState
+
+
+@functools.lru_cache(maxsize=1)
+def _matmul_4bit_matches_dequant() -> bool:
+    """Whether ``bnb.matmul_4bit`` agrees with dequantize-then-linear for the ``[packed, 1]`` layout.
+
+    bitsandbytes < 0.50 mishandles the transpose back-compat shim (keyed on ``A.shape[0] == 1``) for a
+    ``[packed, 1]``-shaped 4-bit weight, so ``matmul_4bit`` returns wrong results there. On those
+    releases :meth:`Experts4bit.forward` falls back to the (bit-exact) dequantize path. Probed once,
+    lazily, on the first forward — GPU-only, so returns ``False`` without CUDA.
+    """
+    try:
+        if not torch.cuda.is_available():
+            return False
+        dev = "cuda"
+        w = torch.randn(64, 64, dtype=torch.bfloat16, device=dev)
+        x = torch.randn(2, 64, dtype=torch.bfloat16, device=dev)
+        q, st = F.quantize_4bit(w.contiguous(), blocksize=64, compress_statistics=False, quant_type="nf4")
+        qs = QuantState(
+            absmax=st.absmax.reshape(-1),
+            shape=torch.Size((64, 64)),
+            code=F.get_4bit_type("nf4", device=dev),
+            blocksize=64,
+            quant_type="nf4",
+            dtype=torch.bfloat16,
+        )
+        ref = F_nn.linear(x, F.dequantize_4bit(q.reshape(-1, 1), quant_state=qs))
+        got = bnb.matmul_4bit(x, q.reshape(-1, 1), quant_state=qs)
+        # Accept bf16 gemv rounding vs dequantize-then-linear (a few %); reject the catastrophic
+        # transpose-shim bug on bitsandbytes<0.50 (relative error >> 1).
+        return bool((got - ref).abs().max() / ref.abs().max().clamp_min(1e-6) < 0.05)
+    except Exception:
+        return False
+
+
+class Experts4bit(nn.Module):
+    """4-bit quantized storage for fused Mixture-of-Experts expert weights.
+
+    A growing number of models in the Hugging Face ecosystem store their MoE expert
+    weights as a single 3D ``nn.Parameter`` of shape ``[num_experts, out_features,
+    in_features]`` (e.g. ``OlmoeExperts``, ``Qwen3MoeExperts``) rather than as a
+    collection of ``nn.Linear`` layers. The default 4-bit quantization walker only
+    replaces ``nn.Linear`` modules, so these fused experts are silently skipped and
+    stay in full precision — the dominant contribution to the model's memory footprint
+    (see https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1849).
+
+    ``Experts4bit`` holds the two expert projections (``gate_up_proj`` and ``down_proj``)
+    in 4-bit NF4/FP4 precision. Unlike :class:`Linear4bit`, the packed weights are kept
+    as plain ``nn.Parameter`` buffers and the per-expert quantization statistics
+    (``absmax``) live on the module as ordinary buffers. This avoids bending
+    :class:`Params4bit`'s tensor-subclass and device-movement machinery around a 3D
+    stack, and it means the module serializes through the standard ``state_dict``
+    mechanism with no custom save/load hooks.
+
+    The forward pass dequantizes a single expert at a time (a per-expert loop), mirroring
+    the reference fused-experts forward. Grouped-GEMM is intentionally left for future
+    work.
+
+    <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>
+
+    Args:
+        num_experts (`int`): Number of experts in the layer.
+        hidden_dim (`int`): Model hidden size (the ``in_features`` of ``gate_up_proj`` and
+            the ``out_features`` of ``down_proj``).
+        intermediate_dim (`int`): Expert intermediate size (the ``in_features`` of
+            ``down_proj``).
+        has_gate (`bool`, *optional*, defaults to `True`): Whether ``gate_up_proj`` packs a
+            gate and an up projection (SwiGLU-style). When `False`, the projection is a
+            plain up projection of size ``intermediate_dim``.
+        activation (`Callable`, *optional*): The activation applied to the gate. Defaults
+            to ``torch.nn.functional.silu`` (SwiGLU), matching OLMoE / Qwen3-MoE.
+        compute_dtype (`torch.dtype`, *optional*): The dtype expert weights are
+            dequantized to for the matmul. When `None`, the input's dtype is used.
+        quant_type (`str`, *optional*, defaults to `"nf4"`): The 4-bit data type, ``nf4``
+            or ``fp4``.
+        blocksize (`int`, *optional*, defaults to `64`): The quantization block size.
+        device (*optional*): The device for the (empty) packed buffers.
+
+    Raises:
+        ValueError: If ``quant_type`` is invalid, or if ``hidden_dim`` / ``intermediate_dim``
+            is not divisible by ``blocksize`` (required so per-expert quantization blocks
+            never straddle an expert boundary).
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        hidden_dim: int,
+        intermediate_dim: int,
+        has_gate: bool = True,
+        activation: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        compute_dtype: Optional[torch.dtype] = None,
+        quant_type: str = "nf4",
+        blocksize: int = 64,
+        device=None,
+    ):
+        super().__init__()
+
+        if quant_type not in ("nf4", "fp4"):
+            raise ValueError(f"quant_type must be 'nf4' or 'fp4', got {quant_type!r}")
+
+        # Each expert is quantized independently, so an expert occupies a contiguous
+        # `out_features * in_features` run of elements. Requiring the in_features dim to
+        # be a multiple of the blocksize guarantees `out_features * in_features` is too,
+        # so blocks tile each expert exactly and absmax reshapes cleanly to
+        # [num_experts, blocks_per_expert]. (gate_up in_features is hidden_dim; down_proj
+        # in_features is intermediate_dim.)
+        for name, in_features in (
+            ("hidden_dim", hidden_dim),
+            ("intermediate_dim", intermediate_dim),
+        ):
+            if in_features % blocksize != 0:
+                raise ValueError(
+                    f"{name} ({in_features}) must be divisible by blocksize ({blocksize}) "
+                    "so per-expert quantization blocks align with expert boundaries"
+                )
+
+        self.num_experts = num_experts
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.has_gate = has_gate
+        self.act_fn = activation if activation is not None else F_nn.silu
+        self.compute_dtype = compute_dtype
+        self.quant_type = quant_type
+        self.blocksize = blocksize
+
+        gate_up_out = 2 * intermediate_dim if has_gate else intermediate_dim
+        self._gate_up_shape = (gate_up_out, hidden_dim)
+        self._down_shape = (hidden_dim, intermediate_dim)
+
+        gate_up_numel = gate_up_out * hidden_dim
+        down_numel = hidden_dim * intermediate_dim
+
+        # Packed 4-bit weights as plain (frozen) parameters: two 4-bit values per byte.
+        self.gate_up_proj = nn.Parameter(
+            torch.empty(num_experts, gate_up_numel // 2, dtype=torch.uint8, device=device),
+            requires_grad=False,
+        )
+        self.down_proj = nn.Parameter(
+            torch.empty(num_experts, down_numel // 2, dtype=torch.uint8, device=device),
+            requires_grad=False,
+        )
+
+        # Per-expert quantization scales.
+        self.register_buffer(
+            "gate_up_absmax",
+            torch.empty(
+                num_experts,
+                gate_up_numel // blocksize,
+                dtype=torch.float32,
+                device=device,
+            ),
+        )
+        self.register_buffer(
+            "down_absmax",
+            torch.empty(num_experts, down_numel // blocksize, dtype=torch.float32, device=device),
+        )
+
+        # The 4-bit codebook is identical for every expert and fully determined by
+        # quant_type, so it is reconstructed at init rather than serialized.
+        self.register_buffer("code", F.get_4bit_type(quant_type, device=device), persistent=False)
+
+    @classmethod
+    def from_float(
+        cls,
+        gate_up_proj: torch.Tensor,
+        down_proj: torch.Tensor,
+        has_gate: bool = True,
+        activation: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        compute_dtype: Optional[torch.dtype] = None,
+        quant_type: str = "nf4",
+        blocksize: int = 64,
+    ) -> "Experts4bit":
+        """Build an :class:`Experts4bit` by quantizing full-precision expert weights.
+
+        Args:
+            gate_up_proj (`torch.Tensor`): Shape ``[num_experts, gate_up_out, hidden_dim]``,
+                where ``gate_up_out`` is ``2 * intermediate_dim`` when ``has_gate`` else
+                ``intermediate_dim``.
+            down_proj (`torch.Tensor`): Shape ``[num_experts, hidden_dim, intermediate_dim]``.
+
+        Returns:
+            `Experts4bit`: A module holding the quantized weights on the inputs' device.
+        """
+        if gate_up_proj.dim() != 3 or down_proj.dim() != 3:
+            raise ValueError("gate_up_proj and down_proj must be 3D [num_experts, out, in] tensors")
+
+        num_experts, _, hidden_dim = gate_up_proj.shape
+        intermediate_dim = down_proj.shape[2]
+
+        module = cls(
+            num_experts,
+            hidden_dim,
+            intermediate_dim,
+            has_gate=has_gate,
+            activation=activation,
+            compute_dtype=compute_dtype if compute_dtype is not None else gate_up_proj.dtype,
+            quant_type=quant_type,
+            blocksize=blocksize,
+            device=gate_up_proj.device,
+        )
+
+        gate_up_packed, gate_up_absmax = module._quantize_stack(gate_up_proj)
+        down_packed, down_absmax = module._quantize_stack(down_proj)
+
+        module.gate_up_proj = nn.Parameter(gate_up_packed, requires_grad=False)
+        module.down_proj = nn.Parameter(down_packed, requires_grad=False)
+        module.gate_up_absmax = gate_up_absmax
+        module.down_absmax = down_absmax
+        return module
+
+    def _quantize_stack(self, weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Quantize a ``[num_experts, out, in]`` stack to packed bytes + per-expert absmax."""
+        packed = []
+        absmax = []
+        for e in range(weights.shape[0]):
+            q, state = F.quantize_4bit(
+                weights[e].contiguous(),
+                blocksize=self.blocksize,
+                compress_statistics=False,
+                quant_type=self.quant_type,
+            )
+            packed.append(q.reshape(-1))
+            absmax.append(state.absmax.reshape(-1))
+        return torch.stack(packed), torch.stack(absmax)
+
+    def _dequantize_expert(
+        self,
+        packed: torch.Tensor,
+        absmax: torch.Tensor,
+        shape: tuple[int, int],
+        expert_idx: int,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Dequantize a single expert's 2D weight ``[out, in]`` for the matmul."""
+        quant_state = QuantState(
+            absmax=absmax[expert_idx],
+            shape=torch.Size(shape),
+            code=self.code,
+            blocksize=self.blocksize,
+            quant_type=self.quant_type,
+            dtype=dtype,
+        )
+        # Restore the [packed, 1] layout quantize_4bit emits (and which keeps the
+        # transpose back-compat shim — keyed on A.shape[0] == 1 — from firing).
+        return F.dequantize_4bit(packed[expert_idx].reshape(-1, 1), quant_state=quant_state)
+
+    def _expert_matmul(
+        self,
+        packed: torch.Tensor,
+        absmax: torch.Tensor,
+        shape: tuple[int, int],
+        expert_idx: int,
+        x: torch.Tensor,
+        compute_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Apply one expert's frozen 4-bit ``[out, in]`` weight to ``x``.
+
+        Routes through :func:`bnb.matmul_4bit` — the same path :class:`Linear4bit` uses — so that
+        in training the autograd ``Function`` re-dequantizes the weight in the backward pass rather
+        than stashing the full-precision expert weight as a saved activation, and never produces a
+        gradient for the frozen packed weights. Numerically identical to dequantize-then-``linear``;
+        this keeps memory low when training a LoRA/adapter on top of a frozen ``Experts4bit`` base.
+        """
+        quant_state = QuantState(
+            absmax=absmax[expert_idx],
+            shape=torch.Size(shape),
+            code=self.code,
+            blocksize=self.blocksize,
+            quant_type=self.quant_type,
+            dtype=compute_dtype,
+        )
+        return bnb.matmul_4bit(x, packed[expert_idx].reshape(-1, 1), quant_state=quant_state)
+
+    def _project(self, packed, absmax, shape, expert_idx, x, compute_dtype, use_matmul_4bit):
+        """One expert projection: matmul_4bit when supported (bnb>=0.50), else dequantize+linear."""
+        if use_matmul_4bit:
+            return self._expert_matmul(packed, absmax, shape, expert_idx, x, compute_dtype)
+        weight = self._dequantize_expert(packed, absmax, shape, expert_idx, compute_dtype)
+        return F_nn.linear(x, weight)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        compute_dtype = self.compute_dtype if self.compute_dtype is not None else hidden_states.dtype
+        hidden_states = hidden_states.to(compute_dtype)
+
+        # matmul_4bit lowers training memory but is CUDA-only here and only correct on
+        # bitsandbytes>=0.50 for the [packed, 1] layout; otherwise use the bit-exact dequantize path.
+        use_matmul_4bit = hidden_states.is_cuda and _matmul_4bit_matches_dequant()
+
+        # Accumulate in float32 for numerical stability with bf16/fp16 routing weights.
+        final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32)
+
+        with torch.no_grad():
+            expert_mask = F_nn.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1)
+
+        for expert_idx in expert_hit:
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+
+            proj = self._project(
+                self.gate_up_proj,
+                self.gate_up_absmax,
+                self._gate_up_shape,
+                expert_idx,
+                current_state,
+                compute_dtype,
+                use_matmul_4bit,
+            )
+            if self.has_gate:
+                gate, up = proj.chunk(2, dim=-1)
+                current_hidden = self.act_fn(gate) * up
+            else:
+                current_hidden = self.act_fn(proj)
+
+            current_hidden = self._project(
+                self.down_proj,
+                self.down_absmax,
+                self._down_shape,
+                expert_idx,
+                current_hidden,
+                compute_dtype,
+                use_matmul_4bit,
+            )
+            current_hidden = current_hidden * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden.to(final_hidden_states.dtype))
+
+        return final_hidden_states.to(hidden_states.dtype)