Views
No views yet
1from compressed_tensors.offload import dispatch_model
2from datasets import load_dataset
3from transformers import (
4 AutoProcessor,
5 MuseGlimmerForConditionalGeneration,
6)
7
8from llmcompressor import oneshot
9from llmcompressor.modifiers.gptq import GPTQModifier
10from llmcompressor.utils import load_context
11
12MODEL_ID = "meta-models/Muse-Glimmer-30B"
13
14# Load model.
15with load_context(MuseGlimmerForConditionalGeneration):
16 model = MuseGlimmerForConditionalGeneration.from_pretrained(MODEL_ID)
17processor = AutoProcessor.from_pretrained(MODEL_ID)
18
19DATASET_ID = "mlabonne/open-perfectblend"
20DATASET_SPLIT = "train"
21
22NUM_CALIBRATION_SAMPLES = 1024
23MAX_SEQUENCE_LENGTH = 2048
24
25# Load dataset and preprocess.
26ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}")
27ds = ds.shuffle(seed=42)
28
29ROLE_MAP = {"human": "user", "gpt": "assistant"}
30
31
32def preprocess(example):
33 messages = [
34 {"role": ROLE_MAP.get(msg["from"], msg["from"]), "content": msg["value"]}
35 for msg in example["conversations"]
36 ]
37 return {
38 "text": processor.apply_chat_template(
39 messages,
40 tokenize=False,
41 )
42 }
43
44
45ds = ds.map(preprocess)
46
47
48# Tokenize inputs.
49def tokenize(sample):
50 return processor.tokenizer(
51 sample["text"],
52 padding=False,
53 max_length=MAX_SEQUENCE_LENGTH,
54 truncation=True,
55 add_special_tokens=False,
56 )
57
58
59ds = ds.map(tokenize, remove_columns=ds.column_names)
60
61# Configure the quantization algorithm and scheme.
62recipe = GPTQModifier(
63 targets="Linear",
64 scheme="NVFP4",
65 ignore=["re:.*vision.*", "lm_head", "re:.*embed_tokens.*"],
66)
67
68# Apply quantization.
69oneshot(
70 model=model,
71 dataset=ds,
72 recipe=recipe,
73 max_seq_length=MAX_SEQUENCE_LENGTH,
74 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
75)
76
77print("\n\n")
78print("========== SAMPLE GENERATION ==============")
79dispatch_model(model)
80input_ids = processor.tokenizer(
81 "Hello my name is", return_tensors="pt"
82).input_ids.to(model.device)
83output = model.generate(input_ids, max_new_tokens=100)
84print(processor.tokenizer.decode(output[0]))
85print("==========================================\n\n")
86
87# Save to disk in compressed-tensors format.
88SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
89model.save_pretrained(SAVE_DIR)
90processor.save_pretrained(SAVE_DIR)
91docker run --gpus all \
--privileged --ipc=host -p 8000:8000 \
-v ~/.cache/huggingface:/root/.cache/huggingface \
vllm/vllm-openai:muse-glimmer RedHatAI/Muse-Glimmer-30B-NVFP4 \
--generation-config auto \
--tensor-parallel-size 1 \
--enable-auto-tool-choice \
--tool-call-parser muse_glimmer \
--reasoning-parser muse_glimmer