Views
No views yet
1from vllm import LLM, SamplingParams
2from transformers import AutoTokenizer
3
4model_id = "GaleneAI/Qwen3-VL-235B-A22B-Thinking-NVFP4"
5number_gpus = 1
6
7sampling_params = SamplingParams(temperature=0.6, top_p=0.9)
8
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10
11messages = [
12 {"role": "system", "content": "You are a proficient Rust developer."},
13 {"role": "user", "content": "Provide me with a quicksort implementation in Rust."},
14]
15
16prompts = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
17
18llm = LLM(model=model_id, tensor_parallel_size=number_gpus)
19
20outputs = llm.generate(prompts, sampling_params)
21
22generated_text = outputs[0].outputs[0].text
23print(generated_text)1import torch
2from datasets import load_dataset
3from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
4
5from llmcompressor import oneshot
6from llmcompressor.modeling import replace_modules_for_calibration
7from llmcompressor.modifiers.quantization import QuantizationModifier
8from llmcompressor.utils import dispatch_for_generation
9
10# NOTE: Requires a minimum of transformers 4.57.0
11
12MODEL_ID = "Qwen/Qwen3-VL-235B-A22B-Thinking"
13
14
15# Load model.
16model = Qwen3VLMoeForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype="auto")
17processor = AutoProcessor.from_pretrained(MODEL_ID)
18model = replace_modules_for_calibration(model)
19
20DATASET_ID = "neuralmagic/calibration"
21NUM_CALIBRATION_SAMPLES = 20
22MAX_SEQUENCE_LENGTH = 8192
23
24ds = load_dataset(DATASET_ID, name="LLM", split=f"train[:{NUM_CALIBRATION_SAMPLES}]")
25
26
27def preprocess_function(example):
28 messgages = []
29 for message in example["messages"]:
30 messgages.append(
31 {
32 "role": message["role"],
33 "content": [{"type": "text", "text": message["content"]}],
34 }
35 )
36
37 return processor.apply_chat_template(
38 messgages,
39 return_tensors="pt",
40 padding=False,
41 truncation=True,
42 max_length=MAX_SEQUENCE_LENGTH,
43 tokenize=True,
44 add_special_tokens=False,
45 return_dict=True,
46 add_generation_prompt=False,
47 )
48
49
50ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
51
52
53def data_collator(batch):
54 assert len(batch) == 1
55 return {
56 key: (
57 torch.tensor(value)
58 if key != "pixel_values"
59 else torch.tensor(value, dtype=torch.bfloat16).squeeze(0)
60 )
61 for key, value in batch[0].items()
62 }
63
64
65# Configure the quantization algorithm and scheme.
66# In this case, we:
67# * quantize the weights to fp4 with group-wise quantization
68# * quantize the activations to fp4 with dynamic group activations
69recipe = QuantizationModifier(
70 targets="Linear",
71 scheme="NVFP4",
72 ignore=[
73 "re:.*lm_head",
74 "re:visual.*",
75 "re:model.visual.*",
76 "re:.*mlp.gate$",
77 ],
78)
79
80# Apply quantization.
81oneshot(
82 model=model,
83 recipe=recipe,
84 max_seq_length=MAX_SEQUENCE_LENGTH,
85 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
86 dataset=ds,
87 data_collator=data_collator,
88)
89
90# Save to disk in compressed-tensors format.
91SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
92model.save_pretrained(SAVE_DIR)
93processor.save_pretrained(SAVE_DIR)
94
95print("========== SAMPLE GENERATION ==============")
96dispatch_for_generation(model)
97input_ids = processor(text="Hello my name is", return_tensors="pt").input_ids.to("cuda")
98output = model.generate(input_ids, max_new_tokens=20)
99print(processor.decode(output[0]))
100print("==========================================")
101