RedHatAI

/

gemma-4-12B-it-NVFP4

Like 6

Follow

Red Hat AI 3.32k

import torch
from datasets import concatenate_datasets, load_dataset
from transformers import AutoModelForImageTextToText, AutoProcessor
from llmcompressor import oneshot
from llmcompressor.modifiers.gptq import GPTQModifier

MODEL_ID = "TrevorJS/gemma-4-12B-it-uncensored"
SAVE_DIR = "gemma-4-12B-it-heretic-v2-NVFP4"
MAX_SEQUENCE_LENGTH = 2048
NUM_TEXT, NUM_IMAGE, NUM_AUDIO = 256, 128, 128
MAX_AUDIO_MS = 750 * 40  # audio_seq_length x audio_ms_per_token from processor_config.json

model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype="auto")
processor = AutoProcessor.from_pretrained(MODEL_ID)

# The chat template only reads content-block *type* and emits a placeholder token;
# it never reads the media payload. Media must be passed to the processor separately.
def render(messages):
    return processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

def build_text(ex):
    messages = [{"role": m["role"], "content": [{"type": "text", "text": m["content"]}]} for m in ex["messages"]]
    return processor.apply_chat_template(
        messages, return_tensors="pt", padding=False, truncation=True, max_length=MAX_SEQUENCE_LENGTH,
        tokenize=True, add_special_tokens=False, return_dict=True, add_generation_prompt=False,
    )

def build_image(ex):
    caption = ex["caption"][0] if isinstance(ex["caption"], list) else ex["caption"]
    messages = [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]},
        {"role": "assistant", "content": [{"type": "text", "text": caption}]},
    ]
    return processor(
        text=render(messages), images=ex["image"].convert("RGB"), return_tensors="pt",
        padding=False, truncation=True, max_length=MAX_SEQUENCE_LENGTH, add_special_tokens=False,
    )

def build_audio(ex):
    audio = ex["audio"]
    messages = [
        {"role": "user", "content": [{"type": "audio"}, {"type": "text", "text": "Transcribe this audio."}]},
        {"role": "assistant", "content": [{"type": "text", "text": ex["text"]}]},
    ]
    return processor(
        text=render(messages), audio=audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt",
        padding=False, truncation=True, max_length=MAX_SEQUENCE_LENGTH, add_special_tokens=False,
    )

def data_collator(batch):
    assert len(batch) == 1
    out = {}
    for k, v in batch[0].items():
        if v is None:  # concatenate_datasets unifies schemas; text rows carry pixel_values=None, etc.
            continue
        t = torch.as_tensor(v)
        out[k] = t.to(torch.bfloat16) if t.is_floating_point() else t
    return out

ds_text = load_dataset("neuralmagic/calibration", name="LLM", split=f"train[:{NUM_TEXT}]")
ds_text = ds_text.map(build_text, remove_columns=ds_text.column_names)

ds_image = load_dataset("lmms-lab/flickr30k", split=f"test[:{NUM_IMAGE}]")
ds_image = ds_image.map(build_image, remove_columns=ds_image.column_names)

ds_audio = load_dataset("MLCommons/peoples_speech", "validation", split=f"validation[:{NUM_AUDIO * 4}]")
ds_audio = ds_audio.filter(lambda x: x["duration_ms"] <= MAX_AUDIO_MS).select(range(NUM_AUDIO))
ds_audio = ds_audio.map(build_audio, remove_columns=ds_audio.column_names)

# Shuffle so GPTQ's sequential calibration interleaves the three distributions.
ds = concatenate_datasets([ds_text, ds_image, ds_audio]).shuffle(seed=42)

recipe = GPTQModifier(
    targets="Linear",
    scheme="NVFP4",
    ignore=["re:.*vision.*", "re:.*audio.*", "lm_head", "re:.*embed.*"],
    kv_cache_scheme={"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": False, "symmetric": True},
)

oneshot(
    model=model,
    recipe=recipe,
    dataset=ds,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=len(ds),
    data_collator=data_collator,
    sequential_targets=["Gemma4UnifiedTextDecoderLayer"],
    tracing_ignore=["Gemma4UnifiedVisionEmbedder", "Gemma4UnifiedMultimodalEmbedder"],
)

model.save_pretrained(SAVE_DIR, save_compressed=True)
processor.save_pretrained(SAVE_DIR)

Run with llm-compressor 0.13.0, transformers 5.14.1, torch 2.13.0+cu130 on an NVIDIA GB10. Audio decoding requires soundfile, librosa, torchcodec, and a system ffmpeg.