Documentation

LeanMlir.E4M3Quant

Pure-Lean E4M3 (fp8) fake-quant over raw-f32 ByteArrays #

The quantizer the Lean fp8 trainer (MainMnistLinearE4M3Verified) wraps around the verified linear train-step kernel. It implements the same E4M3 (1-4-3, bias 7) round-to-nearest-even grid as the numpy oracle (scripts/demos/mnist_e4m3_demo.py:to_e4m3, whose np.round rounds ties to even, as OCP FP8 E4M3 does): subnormals on the e = −6 grid (step 2⁻⁹), saturating at ±448.

This is the host-side "operand byte preparation" half of the fp8 render tie (LeanMlir/Proofs/Float/E4M3Fold.lean): fp8 = fp32 arithmetic on operands projected onto the E4M3 grid, with fp32 accumulate inside the kernel. No fp8 hardware or fp8 StableHLO type is needed — q runs here, in Lean, before the verified matmul sees the bytes.

Values move through F32.read (extern f32→Float) and are re-encoded f32→Float32 → 4 little-endian bytes; the master weights stay fp32 across the gradient update.

References #

source
def F32E4M3.quantPerTensor (ba : ByteArray) :
ByteArray

Per-tensor E4M3 quant (one scale s = max|·|/448). Round-trips through the grid: returns s · q(vᵢ/s) as f32 bytes (the dequantized operand).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    def F32E4M3.quantPerColumn (ba : ByteArray) (d0 d1 : Nat) :
    ByteArray

    Per-output-column E4M3 quant for a row-major [d0 × d1] matrix (the "block scale" sWⱼ of the fp8 render tie: each output column scaled independently).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      def F32E4M3.addDelta (master wOut wq : ByteArray) :
      ByteArray

      fp32-master gradient-delta update. The fused kernel returns wOut = wq − lr·∇ (it updated the quantized operand). Applying the same gradient to the fp32 master is master + (wOut − wq) = master − lr·∇.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        def F32E4M3.quantPackedParams (params : ByteArray) (specs : Array (Array Nat × Nat)) :
        ByteArray

        Quantize a packed param buffer slot-by-slot, per the (dims, initKind) layout (VerifiedNetSpec.toSpecs, func-arg order). initKind == 0 weights: dense [ic, oc] → quantPerColumn (per output column), conv [oc, ic, k, k] → quantPerLeadingBlock (per output channel). Biases / γ / β (initKind 1 or 2) are copied (kept fp32). This is the packed analogue of the per-tensor/per-column quant the linear trainer does — the weight operand prep for the verified MLP/CNN train step.

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For