Laya, int8 for Burn and WebGPU

Laya's three checkpoints (english at the root, multilingual/, typed-decisions/) with every linear weight quantized to 8 bits, in Burn's own burnpack format. The files are laid out as the upstream repo's, and each model.safetensors is replaced by a model.bpk. They run in a browser, on the user's GPU (WebGPU), through taconite-laya-web; natively they run through taconite-laya's webgpu feature.

Quantization

• Linear weights: every one is Burn Q8S: symmetric int8 in [-127, 127], one float32 scale per 32 consecutive inputs, four values packed in a u32. This is the integer layout Burn's WGSL backend runs: WGSL has no i8, and no f16 without shader-f16.
  • Burn's matmul reads the weights as stored and dequantizes inside the kernel.
  • On the GPU they take a third of the float32 memory, at float32's speed (Radeon 890M).
• Token embedding (gathered on the host): float16 in english and typed-decisions, Q8S in multilingual.
• Everything else (norms, biases, the act head, the type embedding): the checkpoint's own float16.
• Quantizer: Burn itself (Tensor::quantize_dynamic), by taconite-laya's examples/export_q8.rs.

Each file's metadata records the scheme, the source and the writer.

Accuracy

Every reference case of the three checkpoints was run and compared with the float32 model's (PyTorch). The cases are 7 English, 3 multilingual and 1 typed-decisions request, 83 questions in all, requests of 236 to 4509 tokens.

• Every answer is unchanged.
• Probabilities are within 0.019 (worst case; most within 0.01).
• Option logits are within 1.04.

Coarser schemes measured worse:

The embedding result is why it stays float16 here.

Usage

In the browser:

import { Laya, Router, HUB_Q8 } from "./js/laya.js";   // taconite-laya-web

const laya = await Laya.load("english", { base: HUB_Q8 });
const res = await laya.predict(
  { subject: "Duplicate charge", body: "We were billed twice. Refund it or we cancel." },
  { churn_risk: { type: "noul", instructions: "Does the user threaten to cancel or leave?" } },
);

const router = new Router({ base: HUB_Q8 });   // laya's routing between english / multilingual

In Rust (taconite-laya, feature webgpu), after hf download brishen/laya-burn-q8 --local-dir laya-q8:

use taconite_laya::webgpu::{Laya, WebGpu};
let laya = Laya::<WebGpu>::load_dir("laya-q8/multilingual".as_ref(), Default::default())?;
let answers = laya.predict(&state, &questions).await?;

Provenance

• Upstream: convaiinnovations/laya at revision 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851, Apache-2.0. Copyright Convai Innovations. The configs and tokenizers are copied unchanged.
• Quantized with: taconite-laya examples/export_q8.rs (Burn 0.21).