feat: load ConvRot int4 and per-channel int8 safetensors

[mosaic-devel]

Problem

Checkpoints quantized by an external toolchain store the payload as raw bytes plus sidecar tensors describing how to dequantize it. Today they cannot be loaded at all: I8 is rejected as an unsupported dtype, and U8 tensors are silently dropped before they reach tensor_storages.

That includes the ConvRot int4 exports now appearing for Krea 2 (~4.5 bits/weight, roughly a third the size of Q8_0), which matters most on small cards — on an 8 GB RX 6600 XT the int4 model runs where the larger GGUF quants did not.

Solution

Repack once, at load time, into native ggml types, so inference uses the ordinary kernels and no backend-specific support is needed:

On-disk	Repacked to	Lossless?
ConvRot convrot_w4a4 (int4)	Q4_0	yes, apart from F32→F16 on the scale
int8_tensorwise / ConvRot linear_dtype: int8	Q8_0	yes, apart from F32→F16 on the scale
bitsandbytes NF4	F32	yes (non-uniform codebook, no 4-bit ggml equivalent)

Both integer paths are exact: int4 values occupy [-7,7], so the +8 bias lands inside Q4_0's nibble range, and the per-output-channel scale becomes the block scale. Note Q4_0 interleaves element j with j+16 while the source packs adjacent pairs, so the repack is a permutation rather than a copy.

The rotation

ConvRot stores W·H rather than W, where H is a block-diagonal regular-Hadamard rotation over convrot_groupsize input channels. So the activation has to be rotated to match:

y = x·Wᵀ = x·(W'·H)ᵀ = (x·H)·W'ᵀ

Ignoring this does not degrade gracefully — it produces noise. H is the radix-4 butterfly core Kronecker-composed log4(groupsize) times and scaled by 1/sqrt(groupsize), which makes it symmetric as well as orthogonal, so there is no transpose bookkeeping. It is applied as one matmul against a shared constant, costing groupsize/out_features extra FLOPs — under 5% at 6144 wide.

Activations stay F32, so this is W4A16: slower than a dedicated 4-bit-activation kernel, but more accurate than the paper's W4A4.

Because the hook lives in Linear, every model in the repo gets this automatically — no per-architecture changes.

Two deliberate refusals

• LoRA is skipped on rotated layers, with a warning. Adapter deltas are trained against the unrotated weight and cannot be mixed into a rotated activation.
• convert refuses rotated weights. No output container carries the rotation marker, so an export would load cleanly and generate noise.

Verification

Format was reverse-engineered and checked numerically against a reference GGUF of the same base model (correlation of dequantized weights vs. reference):

	correlation
dequantized as-is	0.06 — noise
dequantized + un-rotated by H	0.954 – 0.993

The repack itself is bit-exact against a direct reference dequantization (dequantize_row_q4_0 / _q8_0), comparing against the F16-rounded scale:

int4->q4_0 max abs err = 0.000e+00
int8->q8_0 max abs err = 0.000e+00

End-to-end on Vulkan (RX 6600 XT, RADV), Krea 2 Turbo ConvRot int4 + Wan2.1 VAE:

quantized weights: 262 int4, 0 int8, 0 nf4 (262 convrot-rotated), 262 sidecars consumed

generates a correct image with legible rendered sign text — which is a fairly sensitive check, since a subtly wrong rotation, nibble order or scale garbles text first. Happy to attach samples.

Caveats

• NF4 is implemented but unverified — I had no bitsandbytes checkpoint to test against. It follows the documented layout (codebook + per-block absmax, including the double-quantized absmax path), but I would rather flag it than imply it is tested. Glad to drop it from this PR if you'd prefer the int4/int8 paths alone.
• int8 shares the tested code path but was not exercised on a real int8 checkpoint either, for the same reason.
• Only the Vulkan backend was built and tested here.

Note

This was written with AI assistance. Errors are mine; happy to rework or drop any part of it.

[leejet]

Thanks for the PR. But I don't think converting ConvRot/int4 weights into existing ggml quant types in the safetensors loader is the right abstraction. These weights are not semantically Q4_0/Q8_0; they require additional runtime behavior (activation rotation and special dequantization). Repacking them into existing ggml types hides the actual tensor representation and prevents backend-specific optimized kernels in the future.

IMO this should be implemented as a native ggml tensor type/operator instead. The loader should preserve the original quantization format and let ggml handle execution.