Views
No views yet
RangerX/Qwen3.6-35B-REAP-Pruned-ratio-0.5 with llm-compressor.
The model has both weights and activations quantized to NVFP4 format in
compressed-tensors.llm-compressor version 0.10.1.dev131+g22ebb057.
The compression run was performed on a machine with 64 GB system RAM and an
NVIDIA GeForce RTX 5070 Ti 16 GiB GPU.uv pip install flashinfer-python flashinfer-cubin flashinfer-jit-cacheflashinfer-jit-cache is important because it avoids running the FlashInfer
JIT compilation phase at serve startup, which can run out of memory even on a
machine with 64 GiB system RAM.1export CUDA_VISIBLE_DEVICES=0
2export OMP_NUM_THREADS=1
3export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
4
5vllm serve sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 \
6 --host 0.0.0.0 \
7 --port 8000 \
8 --language-model-only \
9 --reasoning-parser qwen3 \
10 --moe_backend flashinfer_cutlass \
11 --max-model-len auto \
12 --max-num-seqs 1 \
13 --max-num-batched-tokens 1024 \
14 --kv-cache-dtype fp8 \
15 --gpu-memory-utilization 0.94 \
16 --mm-processor-cache-gb 035,623 FP8 KV-cache tokens.1uvx llama-benchy --base-url "http://0.0.0.0:8000/v1" \
2 --depth 0 2048 4096 8096 \
3 --tg 128 \
4 --latency-mode generation| model | test | t/s | peak t/s | ttfr (ms) | est_ppt (ms) | e2e_ttft (ms) |
|---|---|---|---|---|---|---|
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | pp2048 | 28408.75 +/- 37.69 | 132.43 +/- 0.10 | 72.13 +/- 0.10 | 132.43 +/- 0.10 | |
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | tg128 | 133.27 +/- 0.02 | 134.33 +/- 0.02 | |||
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | pp2048 @ d2048 | 20669.58 +/- 430.28 | 258.59 +/- 4.14 | 198.28 +/- 4.14 | 258.59 +/- 4.14 | |
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | tg128 @ d2048 | 132.89 +/- 0.03 | 133.94 +/- 0.03 | |||
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | pp2048 @ d4096 | 18389.86 +/- 185.67 | 394.47 +/- 3.38 | 334.17 +/- 3.38 | 394.47 +/- 3.38 | |
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | tg128 @ d4096 | 131.84 +/- 0.02 | 132.88 +/- 0.03 | |||
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | pp2048 @ d8096 | 16152.89 +/- 1.75 | 688.37 +/- 0.02 | 628.06 +/- 0.02 | 688.37 +/- 0.02 | |
| sroecker/Qwen3.6-35B-REAP-Pruned-ratio-0.5-NVFP4 | tg128 @ d8096 | 131.05 +/- 0.01 | 132.08 +/- 0.01 |
uv run examples/quantization_w4a4_fp4/rangerx_qwen3_6_reap_pruned_nvfp4_optimized_bucketed_batch8.py1import torch
2from compressed_tensors.utils import save_mtp_tensors_to_checkpoint
3from datasets import load_dataset
4from torch.nn.utils.rnn import pad_sequence
5from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration
6from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
7 Qwen3_5MoeDecoderLayer,
8)
9
10from llmcompressor import oneshot
11from llmcompressor.modifiers.quantization import QuantizationModifier
12
13# NOTE: This example requires transformers >= v5
14
15MODEL_ID = "RangerX/Qwen3.6-35B-REAP-Pruned-ratio-0.5"
16SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
17
18
19def qwen3_5_moe_decoder_forward_for_calibration(
20 self,
21 hidden_states,
22 position_embeddings,
23 attention_mask=None,
24 position_ids=None,
25 past_key_values=None,
26 **kwargs,
27):
28 residual = hidden_states
29 hidden_states = self.input_layernorm(hidden_states)
30
31 if self.layer_type == "linear_attention":
32 hidden_states = self.linear_attn(
33 hidden_states=hidden_states,
34 cache_params=past_key_values,
35 attention_mask=attention_mask,
36 )
37 elif self.layer_type == "full_attention":
38 hidden_states, _ = self.self_attn(
39 hidden_states=hidden_states,
40 attention_mask=attention_mask,
41 position_ids=position_ids,
42 past_key_values=past_key_values,
43 position_embeddings=position_embeddings,
44 **kwargs,
45 )
46
47 hidden_states = residual + hidden_states
48
49 residual = hidden_states
50 hidden_states = self.post_attention_layernorm(hidden_states)
51 hidden_states = self.mlp(hidden_states)
52 hidden_states = residual + hidden_states
53
54 return hidden_states
55
56
57# The upstream decoder contains a defensive tuple-unpack branch after the MoE MLP.
58# During llm-compressor sequential tracing, that branch is autowrapped into a helper
59# that references "_" before assignment. The Qwen3.5 MoE MLP returns a tensor for this
60# model, so removing the branch keeps the calibration forward equivalent and traceable.
61Qwen3_5MoeDecoderLayer.forward = qwen3_5_moe_decoder_forward_for_calibration
62
63# Load model.
64model = Qwen3_5MoeForConditionalGeneration.from_pretrained(MODEL_ID, dtype="auto")
65processor = AutoProcessor.from_pretrained(MODEL_ID)
66tokenizer = getattr(processor, "tokenizer", processor)
67pad_token_id = tokenizer.pad_token_id
68if pad_token_id is None or pad_token_id < 0:
69 pad_token_id = tokenizer.eos_token_id
70
71# No need to include mtp layers as they are not loaded
72# through Qwen3_5MoeForConditionalGeneration
73recipe = QuantizationModifier(
74 targets="Linear",
75 scheme="NVFP4",
76 ignore=[
77 "re:.*lm_head",
78 "re:visual.*",
79 "re:model.visual.*",
80 "re:.*mlp.gate$",
81 "re:.*embed_tokens$",
82 "re:.*shared_expert_gate$",
83 "re:.*linear_attn.*",
84 ],
85)
86
87NUM_CALIBRATION_SAMPLES = 256
88MAX_SEQUENCE_LENGTH = 4096
89BATCH_SIZE = 8
90SEQUENTIAL_TARGETS_PER_SUBGRAPH = 2
91PADDING_STATS_INTERVAL = 25
92TRACING_IGNORE = [
93 "_update_causal_mask",
94 "create_causal_mask",
95 "_update_mamba_mask",
96 "make_causal_mask",
97 "get_causal_mask",
98 "mask_interface",
99 "mask_function",
100 "_prepare_4d_causal_attention_mask",
101 "_prepare_fsmt_decoder_inputs",
102 "_prepare_4d_causal_attention_mask_with_cache_position",
103 "_update_linear_attn_mask",
104 "project_per_layer_inputs",
105 "apply_mask_to_padding_states",
106]
107
108ds = load_dataset(
109 "HuggingFaceH4/ultrachat_200k",
110 split=f"train_sft[:{NUM_CALIBRATION_SAMPLES}]",
111)
112ds = ds.select_columns(["messages"])
113ds = ds.shuffle(seed=42)
114
115
116def _to_sequence_tensor(value):
117 tensor = torch.as_tensor(value)
118 if tensor.ndim == 2 and tensor.shape[0] == 1:
119 tensor = tensor.squeeze(0)
120 return tensor
121
122
123def _percentile(sorted_values, percentile):
124 if not sorted_values:
125 return 0
126
127 index = round((len(sorted_values) - 1) * percentile)
128 return sorted_values[index]
129
130
131def _padding_waste(lengths, batch_size):
132 real_tokens = sum(lengths)
133 padded_tokens = 0
134
135 for start in range(0, len(lengths), batch_size):
136 batch_lengths = lengths[start : start + batch_size]
137 padded_tokens += max(batch_lengths) * len(batch_lengths)
138
139 waste = 1 - (real_tokens / padded_tokens) if padded_tokens else 0
140 return real_tokens, padded_tokens, waste
141
142
143def _print_length_report(label, lengths):
144 sorted_lengths = sorted(lengths)
145 real_tokens, padded_tokens, waste = _padding_waste(lengths, BATCH_SIZE)
146 print(
147 f"[lengths] {label}: "
148 f"samples={len(lengths)} "
149 f"min={sorted_lengths[0]} "
150 f"p50={_percentile(sorted_lengths, 0.50)} "
151 f"p90={_percentile(sorted_lengths, 0.90)} "
152 f"max={sorted_lengths[-1]} "
153 f"padding_waste={waste:.1%} "
154 f"tokens={real_tokens}/{padded_tokens}"
155 )
156
157
158def preprocess_function(example):
159 messages = [
160 {"role": m["role"], "content": [{"type": "text", "text": m["content"]}]}
161 for m in example["messages"]
162 ]
163 encoded = processor.apply_chat_template(
164 messages,
165 tokenize=True,
166 return_dict=True,
167 add_generation_prompt=False,
168 processor_kwargs={
169 "return_tensors": "pt",
170 "padding": False,
171 "truncation": True,
172 "max_length": MAX_SEQUENCE_LENGTH,
173 "add_special_tokens": False,
174 },
175 )
176 return {key: _to_sequence_tensor(value).tolist() for key, value in encoded.items()}
177
178
179def add_length(example):
180 return {"length": len(example["input_ids"])}
181
182
183ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
184ds = ds.map(add_length, batched=False)
185
186shuffled_lengths = ds["length"]
187_print_length_report("shuffled", shuffled_lengths)
188
189ds = ds.sort("length")
190bucketed_lengths = ds["length"]
191_print_length_report("length-bucketed", bucketed_lengths)
192ds = ds.remove_columns(["length"])
193
194padding_stats = {
195 "batches": 0,
196 "real_tokens": 0,
197 "padded_tokens": 0,
198}
199
200
201def data_collator(batch):
202 features = [
203 {key: _to_sequence_tensor(value) for key, value in example.items()}
204 for example in batch
205 ]
206
207 input_lengths = [feature["input_ids"].numel() for feature in features]
208 batch_real_tokens = sum(input_lengths)
209 batch_padded_tokens = max(input_lengths) * len(input_lengths)
210
211 padding_stats["batches"] += 1
212 padding_stats["real_tokens"] += batch_real_tokens
213 padding_stats["padded_tokens"] += batch_padded_tokens
214
215 batch_index = padding_stats["batches"]
216 if batch_index <= 5 or batch_index % PADDING_STATS_INTERVAL == 0:
217 batch_waste = 1 - (batch_real_tokens / batch_padded_tokens)
218 total_waste = 1 - (
219 padding_stats["real_tokens"] / padding_stats["padded_tokens"]
220 )
221 print(
222 f"[padding] batch={batch_index} "
223 f"size={len(features)} "
224 f"max_len={max(input_lengths)} "
225 f"waste={batch_waste:.1%} "
226 f"running_waste={total_waste:.1%} "
227 f"tokens={batch_real_tokens}/{batch_padded_tokens}"
228 )
229
230 collated = {}
231 for key in features[0]:
232 padding_value = pad_token_id if key == "input_ids" else 0
233 collated[key] = pad_sequence(
234 [feature[key] for feature in features],
235 batch_first=True,
236 padding_value=padding_value,
237 )
238
239 return collated
240
241
242# Apply quantization.
243oneshot(
244 model=model,
245 recipe=recipe,
246 dataset=ds,
247 batch_size=BATCH_SIZE,
248 max_seq_length=MAX_SEQUENCE_LENGTH,
249 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
250 shuffle_calibration_samples=False,
251 moe_calibrate_all_experts=True,
252 data_collator=data_collator,
253 # Optimized layerwise mode: batch calibration samples, keep the next cached
254 # activation batch prefetched, and group two decoder layers into each subgraph.
255 sequential_targets=["Qwen3_5MoeDecoderLayer"],
256 sequential_targets_per_subgraph=SEQUENTIAL_TARGETS_PER_SUBGRAPH,
257 sequential_prefetch=True,
258 tracing_ignore=TRACING_IGNORE,
259)
260
261# Save to disk in compressed-tensors format.
262model.save_pretrained(SAVE_DIR)
263processor.save_pretrained(SAVE_DIR)
264
265# MTP layers are excluded from the model through Qwen3_5MoeForConditionalGeneration
266# Save them as-is from the original checkpoint into the quantized output.
267try:
268 save_mtp_tensors_to_checkpoint(source_model=MODEL_ID, dest_dir=SAVE_DIR)
269except ValueError as exc:
270 if "No tensors with prefix 'mtp'" not in str(exc):
271 raise
272 print(f"Skipping MTP tensor copy: {exc}")43.3% in shuffled order to 3.8% in bucketed order at batch
size 8.6: stable, lower VRAM, slower than batch 88: best observed balance12: fit in VRAM but was slower and approached the 16 GiB limit