import torch
from compressed_tensors.offload import init_dist
from compressed_tensors.quantization.quant_scheme import (
FP8_BLOCK,
NVFP4,
QuantizationScheme,
)
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.datasets.utils import get_rank_partition
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.utils import load_context
init_dist()
model_id = "zai-org/GLM-5.2"
with load_context():
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto_offload",
max_memory={},
offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048
ds = load_dataset(
DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES)
)
ds = ds.shuffle(seed=42)
def preprocess(example):
return {
"text": tokenizer.apply_chat_template(
example["messages"],
tokenize=False,
)
}
ds = ds.map(preprocess)
def tokenize(sample):
return tokenizer(
sample["text"],
padding=False,
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
add_special_tokens=False,
)
ds = ds.map(tokenize, remove_columns=ds.column_names)
recipe = QuantizationModifier(
config_groups={
"attention_shared_experts": QuantizationScheme(
targets=[r"re:.*self_attn\..*"],
**FP8_BLOCK,
),
"mlp": QuantizationScheme(
targets=[r"re:.*mlp\..*"],
**NVFP4,
),
},
ignore=[
r"re:^model\.layers\.[0-2]\..*"
r"re:.*mlp\.gate.*",
r"re:.*indexer\.weights_proj$",
r"lm_head",
],
)
oneshot(
model=model,
dataset=ds,
batch_size=4,
recipe=recipe,
shuffle_calibration_samples=False,
)
model.generation_config.top_p = None
SAVE_DIR = (
"/mnt/nvme-data/engine/kylesayrs/"
+ model_id.rstrip("/").split("/")[-1]
+ "-NVFP4-FP8"
)
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)
torch.distributed.destroy_process_group()