feat: generalize Experts4bit to ExpertsNbit with 8-bit and 16-bit storage

Refactor the fused-MoE expert storage module into a generic ExpertsNbit
parametrized by quant_type, with Experts4bit kept as a fully back-compatible
fixed-4-bit subclass (same API, same state_dict layout, same tests):

- "nf4" / "fp4": the existing 4-bit blockwise path, unchanged.
- "int8" / "fp8": 8-bit blockwise via quantize_blockwise — one codebook
  index per byte, 2x compression at much higher fidelity than 4-bit
  ("int8" = the blockwise dynamic map, not LLM.int8() vectorwise; "fp8" =
  an e4m3 codebook via create_fp8_map).
- "bf16" / "fp16": unquantized passthrough storage (no codebook/absmax) as
  a reference baseline and per-layer opt-out; skips the blocksize check.

Tests: new tests/test_experts_nbit.py (roundtrip tolerances per scheme,
int8-beats-nf4 fidelity ordering, forward/backward vs reference, strict
state_dict roundtrip incl. absent absmax keys for passthrough, frozen-base
LoRA training on int8, subclass quant_type restriction).

Also folds in a one-line ruff-format fix to the #1849 regression test so
`pre-commit run --all-files` passes on the branch.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: pjordanandrsn

bitsandbytes/nn/__init__.py

@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-from .experts import Experts4bit
+from .experts import Experts4bit, ExpertsNbit
 from .modules import (
     Embedding,
     Embedding4bit,

bitsandbytes/nn/experts.py

@@ -12,25 +12,64 @@
 import bitsandbytes.functional as F
 from bitsandbytes.functional import QuantState
 
-
-class Experts4bit(nn.Module):
-    """4-bit quantized storage for fused Mixture-of-Experts expert weights.
+# Storage bit-width per supported quantization scheme. The 4-bit schemes pack two values per
+# byte via quantize_4bit; the 8-bit schemes store one codebook index per byte via
+# quantize_blockwise ("int8" = the blockwise dynamic map, NOT the LLM.int8() vectorwise
+# scheme; "fp8" = an e4m3 float codebook). The 16-bit entries are unquantized passthrough
+# storage — no codebook, no absmax — useful as a reference baseline or a per-layer opt-out.
+_SCHEME_BITS = {
+    "nf4": 4,
+    "fp4": 4,
+    "int8": 8,
+    "fp8": 8,
+    "bf16": 16,
+    "fp16": 16,
+}
+
+_PASSTHROUGH_DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def _build_code(quant_type: str, device) -> Optional[torch.Tensor]:
+    """The shared per-scheme codebook (`None` for 16-bit passthrough)."""
+    if quant_type in ("nf4", "fp4"):
+        return F.get_4bit_type(quant_type, device=device)
+    if quant_type == "int8":
+        code = F.create_dynamic_map()
+    elif quant_type == "fp8":
+        code = F.create_fp8_map(signed=True, exponent_bits=4, precision_bits=3, total_bits=8)
+    else:
+        return None
+    return code.to(device) if device is not None else code
+
+
+class ExpertsNbit(nn.Module):
+    """Low-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
+    collection of ``nn.Linear`` layers. The default quantization walkers only replace
+    ``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.
+    ``ExpertsNbit`` holds the two expert projections (``gate_up_proj`` and ``down_proj``)
+    at a configurable storage precision selected by ``quant_type``:
+
+      * ``"nf4"`` / ``"fp4"`` — 4-bit blockwise (two values per byte, 4x compression).
+        :class:`Experts4bit` is the fixed-4-bit specialization of this class.
+      * ``"int8"`` / ``"fp8"`` — 8-bit blockwise (one codebook index per byte, 2x
+        compression at substantially higher fidelity than 4-bit). ``"int8"`` uses the
+        blockwise dynamic map (`quantize_blockwise`'s default); it is *not* the
+        LLM.int8() vectorwise scheme. ``"fp8"`` uses an e4m3 float codebook.
+      * ``"bf16"`` / ``"fp16"`` — unquantized passthrough storage (no compression). A
+        reference baseline and a per-layer opt-out inside an otherwise-quantized model.
+
+    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
@@ -51,17 +90,20 @@ class Experts4bit(nn.Module):
             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``.
+        quant_type (`str`, *optional*, defaults to `"nf4"`): The storage scheme — one of
+            ``nf4``, ``fp4``, ``int8``, ``fp8``, ``bf16``, ``fp16``.
         blocksize (`int`, *optional*, defaults to `64`): The quantization block size.
+            Ignored for the 16-bit passthrough schemes.
         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).
+        ValueError: If ``quant_type`` is invalid, or if a quantized scheme is selected and
+            ``hidden_dim`` / ``intermediate_dim`` is not divisible by ``blocksize``
+            (required so per-expert quantization blocks never straddle an expert boundary).
     """
 
+    _ALLOWED_QUANT_TYPES: tuple = tuple(_SCHEME_BITS)
+
     def __init__(
         self,
         num_experts: int,
@@ -76,21 +118,25 @@ def __init__(
     ):
         super().__init__()
 
-        if quant_type not in ("nf4", "fp4"):
-            raise ValueError(f"quant_type must be 'nf4' or 'fp4', got {quant_type!r}")
+        allowed = type(self)._ALLOWED_QUANT_TYPES
+        if quant_type not in allowed:
+            raise ValueError(f"quant_type must be one of {allowed}, got {quant_type!r}")
+
+        self.bits = _SCHEME_BITS[quant_type]
 
         # 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"
-                )
+        # in_features is intermediate_dim.) Passthrough storage has no blocks.
+        if self.bits < 16:
+            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
@@ -108,29 +154,39 @@ def __init__(
         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),
-        )
+        # Packed weights as plain (frozen) parameters. 4-bit: two values per byte; 8-bit:
+        # one codebook index per byte; 16-bit: the weights themselves in the storage dtype.
+        if self.bits == 16:
+            storage_dtype = _PASSTHROUGH_DTYPES[quant_type]
+            gate_up_storage = torch.empty(num_experts, gate_up_numel, dtype=storage_dtype, device=device)
+            down_storage = torch.empty(num_experts, down_numel, dtype=storage_dtype, device=device)
+        else:
+            packed_per_value = 2 if self.bits == 4 else 1
+            gate_up_storage = torch.empty(
+                num_experts, gate_up_numel // packed_per_value, dtype=torch.uint8, device=device
+            )
+            down_storage = torch.empty(num_experts, down_numel // packed_per_value, dtype=torch.uint8, device=device)
+        self.gate_up_proj = nn.Parameter(gate_up_storage, requires_grad=False)
+        self.down_proj = nn.Parameter(down_storage, requires_grad=False)
+
+        # Per-expert quantization scales (absent for passthrough storage — registering None
+        # keeps attribute access uniform while leaving state_dict free of the keys).
+        if self.bits < 16:
+            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),
+            )
+        else:
+            self.register_buffer("gate_up_absmax", None)
+            self.register_buffer("down_absmax", None)
 
-        # 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)
+        # The 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", _build_code(quant_type, device), persistent=False)
 
     @classmethod
     def from_float(
@@ -142,8 +198,8 @@ def from_float(
         compute_dtype: Optional[torch.dtype] = None,
         quant_type: str = "nf4",
         blocksize: int = 64,
-    ) -> "Experts4bit":
-        """Build an :class:`Experts4bit` by quantizing full-precision expert weights.
+    ) -> "ExpertsNbit":
+        """Build an :class:`ExpertsNbit` by quantizing full-precision expert weights.
 
         Args:
             gate_up_proj (`torch.Tensor`): Shape ``[num_experts, gate_up_out, hidden_dim]``,
@@ -152,7 +208,7 @@ def from_float(
             down_proj (`torch.Tensor`): Shape ``[num_experts, hidden_dim, intermediate_dim]``.
 
         Returns:
-            `Experts4bit`: A module holding the quantized weights on the inputs' device.
+            `ExpertsNbit`: 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")
@@ -181,30 +237,55 @@ def from_float(
         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."""
+    def _quantize_stack(self, weights: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Quantize a ``[num_experts, out, in]`` stack to packed storage + per-expert absmax.
+
+        For 16-bit passthrough the stack is stored as-is (flattened per expert) and the
+        returned absmax is `None`.
+        """
+        if self.bits == 16:
+            storage_dtype = _PASSTHROUGH_DTYPES[self.quant_type]
+            return weights.reshape(weights.shape[0], -1).to(storage_dtype).contiguous(), None
+
         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,
-            )
+            if self.bits == 4:
+                q, state = F.quantize_4bit(
+                    weights[e].contiguous(),
+                    blocksize=self.blocksize,
+                    compress_statistics=False,
+                    quant_type=self.quant_type,
+                )
+            else:
+                q, state = F.quantize_blockwise(
+                    weights[e].contiguous(),
+                    code=self.code,
+                    blocksize=self.blocksize,
+                )
             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,
+        absmax: Optional[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."""
+        if self.bits == 16:
+            return packed[expert_idx].reshape(shape).to(dtype)
+        if self.bits == 8:
+            quant_state = QuantState(
+                absmax=absmax[expert_idx],
+                code=self.code,
+                blocksize=self.blocksize,
+                dtype=dtype,
+            )
+            return F.dequantize_blockwise(packed[expert_idx].reshape(shape), quant_state=quant_state)
         quant_state = QuantState(
             absmax=absmax[expert_idx],
             shape=torch.Size(shape),
@@ -217,6 +298,11 @@ def _dequantize_expert(
         # 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 _project(self, packed, absmax, shape, expert_idx, x, compute_dtype):
+        """One expert projection: dequantize + ``linear``."""
+        weight = self._dequantize_expert(packed, absmax, shape, expert_idx, compute_dtype)
+        return F_nn.linear(x, weight)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -238,21 +324,43 @@ def forward(
             top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
             current_state = hidden_states[token_idx]
 
-            gate_up_w = self._dequantize_expert(
-                self.gate_up_proj, self.gate_up_absmax, self._gate_up_shape, expert_idx, compute_dtype
+            proj = self._project(
+                self.gate_up_proj,
+                self.gate_up_absmax,
+                self._gate_up_shape,
+                expert_idx,
+                current_state,
+                compute_dtype,
             )
-            proj = F_nn.linear(current_state, gate_up_w)
             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)
 
-            down_w = self._dequantize_expert(
-                self.down_proj, self.down_absmax, self._down_shape, expert_idx, compute_dtype
+            current_hidden = self._project(
+                self.down_proj,
+                self.down_absmax,
+                self._down_shape,
+                expert_idx,
+                current_hidden,
+                compute_dtype,
             )
-            current_hidden = F_nn.linear(current_hidden, down_w)
             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)
+
+
+class Experts4bit(ExpertsNbit):
+    """4-bit quantized storage for fused Mixture-of-Experts expert weights.
+
+    The fixed-4-bit (``nf4`` / ``fp4``) specialization of :class:`ExpertsNbit` — see the
+    parent class for the storage design, constructor arguments, and forward semantics.
+    Kept as a distinct class so the 4-bit contract (two packed values per byte, 4x
+    compression) has a stable name.
+
+    <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>
+    """
+
+    _ALLOWED_QUANT_TYPES = ("nf4", "fp4")

tests/test_experts4bit.py

@@ -323,7 +323,8 @@ def __init__(self):
     )
     assert q.gate_up_proj.dtype == torch.uint8 and q.down_proj.dtype == torch.uint8
     quantized_bytes = (
-        q.gate_up_proj.numel() + q.down_proj.numel()  # uint8 packed
+        q.gate_up_proj.numel()
+        + q.down_proj.numel()  # uint8 packed
         + (q.gate_up_absmax.numel() + q.down_absmax.numel()) * 4  # fp32 absmax
     )
     assert quantized_bytes < fp16_bytes / 3  # ~4x on the weights, minus small absmax overhead