Support ternary models

[diimdeep]

ggml-quants : ternary packing for TriLMs and BitNet b1.58 #8151 is on the horizon

There are TerDiT: Ternary Diffusion Models

4.2B of ternary weights packed lossless:

• 2,000 bpw - 1.05GB
• 1,625 bpw - 914 MB

[kevinlaw91]

Bonsai Image Ternary dropped 🎉. Please add support

[Green-Sky]

iirc bonsai's ternary is different to what ggml offers.

However, q1_0 is supported by ggml.

https://huggingface.co/prism-ml/bonsai-image-binary-4B-unpacked should be quantizable.
All *transformer_blocks.*.weight should be q1_0.

Arch looks to be flux klein 4B, so maybe this is just a question of quantizing it.

[Green-Sky]

[INFO ] model.cpp:1182 - unknown tensor 'model.diffusion_model.double_stream_modulation_img.linear.weight | bf16 | 2 [3072, 18432, 1, 1, 1]' in model file
[INFO ] model.cpp:1182 - unknown tensor 'model.diffusion_model.double_stream_modulation_txt.linear.weight | bf16 | 2 [3072, 18432, 1, 1, 1]' in model file
[ERROR] model.cpp:1191 - tensor 'model.diffusion_model.final_layer.linear.weight' has wrong shape in model file: got [3072, 128, 1, 1], expected [3072, 64, 1, 1]
[WARN ] model.cpp:991  - process tensor failed: 'model.diffusion_model.final_layer.linear.weight'

Hmm.

edit: someone else has to convert this model from diffusers, I'm no expert in model surgery.

[diimdeep]

bonsai so far is just quantized qwen/flux packaged as a branded binary/ternary models and somehow novel, but they are not, not worth imho

[Green-Sky]

Did some model surgery and its mostly there, but not sure where the last issue(s) are.

script

import torch
import safetensors
import safetensors.torch

tensors = {}

replace_keys = {
    "double_stream_modulation_img.linear.weight": "double_stream_modulation_img.lin.weight",
    "double_stream_modulation_txt.linear.weight": "double_stream_modulation_txt.lin.weight",
    "single_stream_modulation.linear.weight": "single_stream_modulation.lin.weight",
    "context_embedder.weight": "txt_in.weight",
    "norm_out.linear.weight": "final_layer.adaLN_modulation.1.weight",
    "proj_out.weight": "final_layer.linear.weight",
    "x_embedder.weight": "img_in.weight",
    "time_guidance_embed.timestep_embedder.linear_1.weight": "time_in.in_layer.weight",
    "time_guidance_embed.timestep_embedder.linear_2.weight": "time_in.out_layer.weight",
    "single_transformer_blocks.": "single_blocks.",
    "transformer_blocks.": "double_blocks.",
}

single_replace_keys = {
    ".attn.norm_k.weight": ".norm.key_norm.scale",
    ".attn.norm_q.weight": ".norm.query_norm.scale",
    ".attn.to_qkv_mlp_proj.": ".linear1.",
    ".attn.to_out.": ".linear2.",
}

double_replace_keys = {
    ".ff_context.linear_in.": ".txt_mlp.0.",
    ".ff_context.linear_out.": ".txt_mlp.2.",
    ".ff.linear_in.": ".img_mlp.0.",
    ".ff.linear_out.": ".img_mlp.2.",
    ".attn.norm_q.weight": ".img_attn.norm.query_norm.scale",
    ".attn.norm_k.weight": ".img_attn.norm.key_norm.scale",
    ".attn.norm_added_q.weight": ".txt_attn.norm.query_norm.scale",
    ".attn.norm_added_k.weight": ".txt_attn.norm.key_norm.scale",
    ".attn.to_out.0.": ".img_attn.proj.",
    ".attn.to_add_out.": ".txt_attn.proj.",
    ".attn.to_q.": ".img_attn.to_q.", # temp, reshaped later
    ".attn.add_q_proj.": ".txt_attn.to_q.", # temp, reshaped later
}

with safetensors.safe_open("path_to_bonsai_binary/diffusion_pytorch_model.safetensors", framework="pt") as f:
    for k_orig in f.keys():
        k = k_orig
        if k.startswith("model.diffusion_model"):
            k.removeprefix("model.diffusion_model")

        for r, rr in replace_keys.items():
            k = k.replace(r, rr)

        if k.startswith("single_blocks."):
            for r, rr in single_replace_keys.items():
                k = k.replace(r, rr)

        if k.startswith("double_blocks."):
            for r, rr in double_replace_keys.items():
                k = k.replace(r, rr)

        if k.startswith("double_blocks.") and ".to_q" in k:
            q_weight = f.get_tensor(k_orig)
            k_weight = f.get_tensor(k_orig.replace("_q", "_k"))
            v_weight = f.get_tensor(k_orig.replace("_q", "_v"))
            tensors[k.replace("to_q", "qkv")] = torch.cat([q_weight, k_weight, v_weight], dim=0)
            continue

        if not ".attn." in k:
            tensors[k] = f.get_tensor(k_orig)
        else:
            print("left over attn tensor: ", k, " og: ", k_orig)

if "final_layer.adaLN_modulation.1.weight" in tensors:
    w = tensors["final_layer.adaLN_modulation.1.weight"]
    tensors["final_layer.adaLN_modulation.1.weight"] = torch.cat([w[w.shape[0]//2:], w[:w.shape[0]//2]], dim=0)

for k in tensors.keys():
    print(k)

# saving time
print("saving to disk...")
safetensors.torch.save_file(tensors, "bonsai-image-unpacked.safetensors")

print("done.")

Maybe someone can spot the error.

I tested quantizing some of the layers they used to q1_0 and the result is identical, which means its working.

edit: sd.cpp quantization command that should cover all q1_0
result/bin/sd-cli -M convert -m bonsai-image-unpacked.safetensors -o bonsai_image_4b-q1_0.gguf --tensor-type-rules "^single_blocks.[0-9]+.linear1.weight=q1_0,^single_blocks.[0-9]+.linear2.weight=q1_0,^double_blocks.[0-9]+.img_attn.qkv.weight=q1_0,^double_blocks.[0-9]+.img_attn.proj.weight=q1_0,^double_blocks.[0-9]+.txt_attn.qkv.weight=q1_0,^double_blocks.[0-9]+.txt_attn.proj.weight=q1_0,^double_blocks.[0-9]+.img_mlp.[0-9].weight=q1_0,^double_blocks.[0-9]+.txt_mlp.[0-9].weight=q1_0" --convert-name

edit: uploaded the broken files here https://huggingface.co/Green-Sky/bonsai-image-binary-4B-GGUF/tree/main/broken

edit: llm suggested a fix and it now works. uploading to https://huggingface.co/Green-Sky/bonsai-image-binary-4B-GGUF/tree/main/

[Green-Sky]

bonsai so far is just quantized qwen/flux packaged as a branded binary/ternary models and somehow novel, but they are not, not worth imho

They did QAT during a finetuning step, since low bit quants are way too rough. But yea, it is "just" flux 2 klein 4b.

[sandeshrajbhandari]

@Green-Sky Can you shrae the command you used to run this?

[Green-Sky]

@Green-Sky Can you share the command you used to run this?

I posted them above. Also the hf repo contains them too.

The image command is based on the example prompt in the klein 4b sd.cpp docs.

[Green-Sky]

I shrank the model by another ~11% down to 765MiB by quantizing the doublestream modulation to q8_0.

866MiB	765MiB

---

I really wish they would add the 128 block size ternary quant support to ggml though, as the binary model really is skirting the bottom end of possible compression. (ggml ternary quants tq1_0 and tq2_0 use a 256 block size)

[Green-Sky]

@khosravipasha q1_0 working well. 😉

[khosravipasha]

Nice this is exciting. Yeah I am trying to add Q2_0 for ternary for llama.cpp they wanted group size 64 so need to make some changes in our side. Probably need to keep the group size 128 version in our fork.

Let me know if can help in any other way too

We have Q2_0 with group 128 in our llama.cpo fork that supports ternary but it's 2 bit I guess.