gemma-4-26B-A4B-it-FP8-Block

# Mixed FP8 quantization for google/gemma-4-26B-A4B-it
#
# Gemma-4-26B-A4B-it has intermediate_size=2112 (shared MLP) and
# moe_intermediate_size=704 (experts), neither divisible by 128.
# Attention layers have dims divisible by 128 (hidden_size=2816, head_dim=256),
# so they can use FP8_BLOCK (128x128 block scaling).
# MLP/expert layers fall back to FP8 static per-tensor.

import os

os.environ["HF_HOME"] = os.path.expanduser("~/hf_hub")
os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_datasets_cache"

import torch
import datasets
from compressed_tensors.quantization.quant_scheme import FP8, FP8_BLOCK
from datasets import load_dataset

datasets.disable_caching()
from transformers import AutoProcessor, Gemma4ForConditionalGeneration

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.utils import load_context

MODEL_ID = "google/gemma-4-26B-A4B-it"

# Load model.
with load_context(Gemma4ForConditionalGeneration):
    model = Gemma4ForConditionalGeneration.from_pretrained(MODEL_ID)
processor = AutoProcessor.from_pretrained(MODEL_ID)

# MoE expert handling is applied automatically by the pipeline.

# FP8_BLOCK for attention layers (all dims divisible by 128).
scheme_block = FP8_BLOCK
scheme_block["targets"] = [
    "re:.*self_attn.q_proj",
    "re:.*self_attn.k_proj",
    "re:.*self_attn.v_proj",
    "re:.*self_attn.o_proj",
]

# FP8 static per-tensor for MLP/expert layers (dims not divisible by 128).
scheme_static = FP8
scheme_static["targets"] = [
    "re:.*gate_proj",
    "re:.*up_proj",
    "re:.*down_proj",
]

recipe = QuantizationModifier(
    config_groups={"group_block": scheme_block, "group_static": scheme_static},
    ignore=[
        "lm_head",
        "re:.*embed.*",
        "re:.*router",
        "re:.*vision_tower.*",
    ],
)

# Calibration dataset (needed for FP8 static activation scales).
DATASET_ID = "neuralmagic/calibration"
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 8192

ds = load_dataset(DATASET_ID, name="LLM", split=f"train[:{NUM_CALIBRATION_SAMPLES}]")


def preprocess_function(example):
    messages = []
    for message in example["messages"]:
        messages.append(
            {
                "role": message["role"],
                "content": [{"type": "text", "text": message["content"]}],
            }
        )

    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,
    )


ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)


def data_collator(batch):
    assert len(batch) == 1
    return {
        key: (
            torch.tensor(value)
            if key != "pixel_values"
            else torch.tensor(value, dtype=torch.bfloat16).squeeze(0)
        )
        for key, value in batch[0].items()
    }


# Apply quantization.
oneshot(
    model=model,
    recipe=recipe,
    dataset=ds,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    data_collator=data_collator,
)

# Save to disk in compressed-tensors format.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-Block"
model.save_pretrained(SAVE_DIR, save_compressed=True)
processor.save_pretrained(SAVE_DIR)