Views
No views yet
1from vllm import LLM, SamplingParams
2from transformers import AutoTokenizer
3
4model_id = "Saktsant/Qwen3-14B-NVFP4"
5number_gpus = 1
6
7sampling_params = SamplingParams(temperature=0.6, top_p=0.9, max_tokens=256)
8
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10
11messages = [
12 {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
13 {"role": "user", "content": "Who are you?"},
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)1from datasets import load_dataset
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from llmcompressor import oneshot
4from llmcompressor.modifiers.quantization import QuantizationModifier
5from llmcompressor.utils import dispatch_for_generation
6
7
8MODEL_ID = "Qwen/Qwen3-14B"
9
10
11# Load model and tokenizer
12model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
13tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
14DATASET_ID = "HuggingFaceH4/ultrachat_200k"
15DATASET_SPLIT = "train_sft"
16NUM_CALIBRATION_SAMPLES = 512
17MAX_SEQUENCE_LENGTH = 2048
18# Load dataset and preprocess
19ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
20ds = ds.shuffle(seed=42)
21def preprocess(example):
22 return {
23 "text": tokenizer.apply_chat_template(
24 example["messages"],
25 tokenize=False,
26 )
27 }
28ds = ds.map(preprocess)
29# Tokenize inputs
30def tokenize(sample):
31 return tokenizer(
32 sample["text"],
33 padding=False,
34 max_length=MAX_SEQUENCE_LENGTH,
35 truncation=True,
36 add_special_tokens=False,
37 )
38ds = ds.map(tokenize, remove_columns=ds.column_names)
39# Configure quantization
40recipe = [
41 QuantizationModifier(
42 ignore=[
43 "re:.*lm_head.*",
44 "re:.*q_proj.*",
45 "re:.*k_proj.*",
46 "re:.*v_proj.*",
47 "re:.*o_proj.*",
48 "re:.*gate_proj.*",
49 "re:.*up_proj.*",
50 "re:.*down_proj.*",
51 ],
52 config_groups={
53 "group_0": {
54 "targets": ["Linear"],
55 "weights": {
56 "num_bits": 4,
57 "type": "float",
58 "strategy": "tensor_group",
59 "group_size": 16,
60 "symmetric": True,
61 "observer": "minmax",
62 },
63 "input_activations": {
64 "num_bits": 4,
65 "type": "float",
66 "strategy": "tensor_group",
67 "group_size": 16,
68 "symmetric": True,
69 "dynamic": "local",
70 "observer": "minmax",
71 },
72 }
73 },
74 )
75]
76# Save directory
77SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
78# Apply quantization
79oneshot(
80 model=model,
81 dataset=ds,
82 recipe=recipe,
83 max_seq_length=MAX_SEQUENCE_LENGTH,
84 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
85 output_dir=SAVE_DIR,
86)
87# Re-dispatch for generation (Accelerate handles device placement)
88model = dispatch_for_generation(model)
89print("\n\n")
90print("========== SAMPLE GENERATION ==============")
91# Prepare inputs with attention_mask
92inputs = tokenizer("Hello my name is", return_tensors="pt")
93inputs = {k: v.to("cuda") for k, v in inputs.items()}
94# Generate
95output = model.generate(**inputs, max_new_tokens=100)
96print(tokenizer.decode(output[0]))
97print("==========================================\n\n")
98# Save compressed model and tokenizer
99model.save_pretrained(SAVE_DIR, save_compressed=True)
100tokenizer.save_pretrained(SAVE_DIR)
101