from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshotfrom llmcompressor.modifiers.quantization import QuantizationModifierfrom llmcompressor.utils import dispatch_for_generationfrom llmcompressor.modeling import replace_modules_for_calibrationfrom llmcompressor.modeling.granite4 import pack_3d_experts
MODEL_ID = "ibm-granite/granite-4.0-h-tiny"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = replace_modules_for_calibration(model)
ignore_lay = ["lm_head", "re:.*block_sparse_moe.router"]
recipe = QuantizationModifier( targets=["Linear"], scheme="FP8_DYNAMIC", ignore=ignore_lay,)
oneshot(model=model, recipe=recipe)
print("========== SAMPLE GENERATION ==============")dispatch_for_generation(model)input_ids = tokenizer( "Describe Large Language Model", return_tensors="pt").input_ids.to(model.device)output = model.generate(input_ids, max_new_tokens=35)print(tokenizer.decode(output[0]))print("==========================================")
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-dynamic"print(f"Saving to {SAVE_DIR}")
model.save_pretrained(SAVE_DIR)tokenizer.save_pretrained(SAVE_DIR)pack_3d_experts(SAVE_DIR)