Add int4 (Q4-style, mixed-precision) RT-J checkpoints + model card

a12ddfd verified 5 days ago

3.76 kB

license: cc-by-nc-sa-4.0
base_model: stanford-star/rt-j
tags:
  - relational-deep-learning
  - relational-databases
  - tabular
  - tabular-classification
  - tabular-regression
  - foundation-model
  - in-context-learning
  - quantized
  - int4
  - relational-transformer
datasets:
  - stanford-star/the-join
  - stanford-star/relbench

RT-J — int4 (Q4-style, mixed-precision) quantized checkpoints

Int4 quantizations of stanford-star/rt-j, the Relational Transformer foundation model for in-context learning over relational databases. Produced and consumed by the RelativeDB native inference engine (cpp/rt_quantize --type q4). Weights stay quantized-resident at inference: the engine's CPU (Accelerate / portable SIMD) and Metal/MPS kernels dequantize inside the GEMM, so DRAM weight traffic is the int4 payload.

Sibling repos: rt-j-int8 · rt-j-fp16

File	Task head	Size
`classification/model.q4.safetensors`	classification / ranking (logits — apply sigmoid)	64 MB
`regression/model.q4.safetensors`	regression / forecasting (normalized values)	64 MB

(vs. 171 MB bf16 upstream, 342 MB fp32 in memory.)

Quantization format

Plain safetensors — no custom container:

Most transformer-block projections (wq/wk/wv/wg, ffn.w1/w3) are uint4 in groups of 32 along the input dim: fp16 (scale, min) pair per group in a <name>.q4_scale companion (U8 payload, two nibbles/byte, low nibble = even index): W[o,i] ≈ nibble * scale + min. Group ranges are chosen by a min-MSE clip search (shrink grid 1.0→0.8).
The residual-writing projections (wo, ffn.w2) stay int8 with per-row fp32 scales (I8 + <name>.q_scale) — the Q4_K_M recipe: their error lands directly on the residual stream every block.
The value/col-name encoders, decoder head, norms, biases and mask embeddings stay fp32 — input-side error would propagate through all 12 blocks.

Accuracy

Measured on the RelativeDB golden batch (B=5, S=16) against the PyTorch reference; identical on CPU and Metal/MPS:

Metric	fp32	int4
`yhat` max abs. error vs. torch	3.9e-3	1.5e-1
`yhat` mean abs. error	5.1e-4	2.4e-2
target-score sign / ranking	preserved	preserved

Int4 is the aggressive end for an 86M-param model: per-score drift is a few hundredths. Use int8 when scores must track fp32 tightly; use int4 when footprint/bandwidth dominates.

Usage (RelativeDB native engine)

# direct path (looser golden gate for int4)
./build/rt_test testdata classification/model.q4.safetensors --tol 100 --device mps

# via the Java / Python / Rust bindings: place the .q4 file next to the fp32
# checkpoint (or point at a directory containing it) and opt in with
export RELATIVEDB_RT_QUANTIZED=q4

The C ABI (rt_model_load) accepts these files directly — the format is auto-detected from the tensor dtypes. Quantized checkpoints currently run on the CPU and Metal/MPS backends (CUDA is fp32-only).

Reproduce

cmake -B build -S cpp && cmake --build build -j
./build/rt_quantize <rt-j>/classification/model.safetensors classification/model.q4.safetensors --type q4
./build/rt_quantize <rt-j>/regression/model.safetensors regression/model.q4.safetensors --type q4

License & attribution

Derivative of stanford-star/rt-j (Stanford STAR lab), redistributed under the same CC-BY-NC-SA-4.0 license. Architecture and training details are described in the upstream model card; only the weight storage format differs here.