from datasets import Dataset
from transformers import AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.utils import load_context
model_id = "/path/to/glm-5.3-tiny-mtp-source"
dataset = Dataset.from_dict(
{
"text": [
"Explain why speculative decoding can reduce inference latency.",
"Write a short Python function that adds two integers.",
"Summarize the benefits of mixture-of-experts language models.",
"What is the capital of France?",
"List three considerations when deploying a language model.",
"Describe the difference between weights and activations.",
"Complete the sequence: one, two, three, four.",
"Give a concise definition of quantization calibration.",
]
}
)
recipe = QuantizationModifier(
targets="Linear",
scheme="NVFP4",
ignore=[
"lm_head",
r"re:.*mlp\.gate$",
r"re:.*self_attn\.indexer\.(?:weights_proj|wk)$",
],
)
with load_context(AutoModelForCausalLM):
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="bfloat16")
model.generation_config.top_p = None
oneshot(
model=model,
recipe=recipe,
dataset=dataset,
output_dir="GLM-5.3-MTP-NVFP4-Test-Fixture",
max_seq_length=128,
num_calibration_samples=len(dataset),
moe_calibrate_all_experts=True,
mtp_scheme="NVFP4",
)