keras-team/keras · error · NotImplementedError

Currently, `_float8_call` doesn't support LoRA

Error message

Currently, `_float8_call` doesn't support LoRA

What it means

The float8 quantized forward path (_float8_call) of EinsumDense does not implement LoRA merging, so a layer with both lora_enabled=True and quantization_mode 'float8' raises NotImplementedError at runtime on the first forward pass.

Source

Thrown at keras/src/layers/core/einsum_dense.py:1265

                ops.convert_to_tensor(self.kernel_zero),
                ops.convert_to_tensor(self.g_idx),
            )

        if self.lora_enabled:
            lora_x = ops.einsum(self.equation, inputs, self.lora_kernel_a)
            lora_x = ops.matmul(lora_x, self.lora_kernel_b)
            x = ops.add(x, (self.lora_alpha / self.lora_rank) * lora_x)
            x = ops.cast(x, dtype=self.compute_dtype)
        # Bias & activation
        if self.bias is not None:
            x = ops.add(x, self.bias)
        if self.activation is not None:
            x = self.activation(x)
        return x

    def _float8_call(self, inputs, training=None):
        if self.lora_enabled:
            raise NotImplementedError(
                "Currently, `_float8_call` doesn't support LoRA"
            )

        @ops.custom_gradient
        def quantized_dequantize_inputs(inputs, scale, amax_history):
            if training:
                new_scale = quantizers.compute_float8_scale(
                    ops.max(amax_history, axis=0),
                    scale,
                    ops.cast(
                        float(ml_dtypes.finfo("float8_e4m3fn").max), "float32"
                    ),
                )
                new_amax_history = quantizers.compute_float8_amax_history(
                    inputs, amax_history
                )
            else:
                new_scale = None

View on GitHub (pinned to 7a34a03db6)

Solutions

  1. Exclude float8-quantized layers from LoRA enabling (check layer.quantization_mode before enable_lora).
  2. Or exclude this layer from float8 quantization via filters if it must be LoRA-trained.
  3. Use a supported combination, e.g. LoRA on a non-quantized or int8 path for this layer.

Example fix

# before
layer.enable_lora(rank=8)
model.quantize(float8_config)  # covers layer
model(x)  # NotImplementedError

# after
if layer.quantization_mode != 'float8':
    layer.enable_lora(rank=8)
Defensive patterns

Strategy: validation

Validate before calling

for l in model.layers:
    if getattr(l, 'lora_enabled', False) and getattr(l, 'quantization_mode', None) == 'float8':
        raise RuntimeError(f'{l.name}: float8 + LoRA unsupported')

Type guard

def lora_forward_safe(layer) -> bool:
    return not (layer.lora_enabled and layer.quantization_mode == 'float8')

Try / catch

try:
    out = model(x)
except NotImplementedError as e:
    if 'float8' in str(e):
        disable_lora_on_float8_layers(model)
    else:
        raise

Prevention

When it happens

Trigger: layer.enable_lora(...) on a layer that is (or later becomes) float8-quantized, then running model(x) or model.fit; a float8 quantization config covering an EinsumDense combined with LoRA fine-tuning.

Common situations: PEFT pipelines on float8 training configs; enabling LoRA broadly across a model that also applies float8 quantization to projections; upgrading configs to float8 while existing LoRA setup code remains.

Related errors

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/f979eb7f53754380. Report an issue: GitHub.