Views
No views yet
llm-compressor.linear_attn layers have been quantized as well in this model to save memory for longer context lengths on RTX 5090 GPUs. Click the dropdown to see the full quantization script.1import torch
2from compressed_tensors.utils import save_mtp_tensors_to_checkpoint
3from datasets import load_dataset
4from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration, AutoModel
5
6from llmcompressor import oneshot
7from llmcompressor.modifiers.quantization import QuantizationModifier
8
9# NOTE: This example requires transformers >= v5
10
11MODEL_ID = "Qwen/Qwen3.6-27B"
12
13# Load model.
14# model = AutoModel.from_pretrained(MODEL_ID, dtype="auto")
15model = Qwen3_5ForConditionalGeneration.from_pretrained(MODEL_ID, dtype="auto")
16processor = AutoProcessor.from_pretrained(MODEL_ID)
17
18# No need to include mtp layers as they are not loaded
19# through Qwen3_5MoeForConditionalGeneration
20recipe = QuantizationModifier(
21 targets="Linear",
22 scheme="NVFP4",
23 ignore=[
24 "re:.*lm_head",
25 "re:visual.*",
26 "re:model.visual.*",
27 "re:.*mlp.gate$",
28 "re:.*embed_tokens$",
29 "re:.*shared_expert_gate$",
30 # "re:.*linear_attn.*",
31 ],
32)
33
34NUM_CALIBRATION_SAMPLES = 256
35MAX_SEQUENCE_LENGTH = 4096
36
37ds = load_dataset(
38 "HuggingFaceH4/ultrachat_200k",
39 split=f"train_sft[:{NUM_CALIBRATION_SAMPLES}]",
40)
41ds = ds.select_columns(["messages"])
42ds = ds.shuffle(seed=42)
43
44
45def preprocess_function(example):
46 messages = [
47 {"role": m["role"], "content": [{"type": "text", "text": m["content"]}]}
48 for m in example["messages"]
49 ]
50 return processor.apply_chat_template(
51 messages,
52 tokenize=True,
53 return_dict=True,
54 add_generation_prompt=False,
55 processor_kwargs={
56 "return_tensors": "pt",
57 "padding": False,
58 "truncation": True,
59 "max_length": MAX_SEQUENCE_LENGTH,
60 "add_special_tokens": False,
61 },
62 )
63
64
65ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
66
67
68def data_collator(batch):
69 assert len(batch) == 1
70 return {key: torch.tensor(value) for key, value in batch[0].items()}
71
72
73# Apply quantization.
74oneshot(
75 model=model,
76 recipe=recipe,
77 dataset=ds,
78 max_seq_length=MAX_SEQUENCE_LENGTH,
79 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
80 moe_calibrate_all_experts=True,
81 data_collator=data_collator,
82)
83
84# Save to disk in compressed-tensors format.
85SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
86model.save_pretrained(SAVE_DIR)
87processor.save_pretrained(SAVE_DIR)
88
89# MTP layers are excluded from the model through Qwen3_5MoeForConditionalGeneration
90# Save them as-is from the original checkpoint into the quantized output.
91save_mtp_tensors_to_checkpoint(source_model=MODEL_ID, dest_dir=SAVE_DIR)