Views
No views yet
1from compressed_tensors.offload import dispatch_model
2from datasets import load_dataset
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5from llmcompressor import oneshot
6from llmcompressor.modifiers.gptq import GPTQModifier
7from llmcompressor.utils import load_context
8
9MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
10
11with load_context(AutoModelForCausalLM):
12 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
13tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
14
15DATASET_ID = "HuggingFaceH4/ultrachat_200k"
16DATASET_SPLIT = "train_sft"
17
18# Select number of samples. 512 samples is a good place to start.
19# Increasing the number of samples can improve accuracy.
20NUM_CALIBRATION_SAMPLES = 512
21MAX_SEQUENCE_LENGTH = 2048
22
23# Load dataset and preprocess.
24ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
25ds = ds.shuffle(seed=42)
26
27
28def preprocess(example):
29 return {
30 "text": tokenizer.apply_chat_template(
31 example["messages"],
32 tokenize=False,
33 )
34 }
35
36
37ds = ds.map(preprocess)
38
39
40# Tokenize inputs.
41def tokenize(sample):
42 return tokenizer(
43 sample["text"],
44 padding=False,
45 max_length=MAX_SEQUENCE_LENGTH,
46 truncation=True,
47 add_special_tokens=False,
48 )
49
50
51ds = ds.map(tokenize, remove_columns=ds.column_names)
52
53recipe = GPTQModifier(
54 targets="Linear",
55 scheme="FP8",
56 ignore=[
57 r"re:.*conv1d.*",
58 r"backbone\.embeddings",
59 r"re:.*_latent_proj.*",
60 r"re:.*mixer.gate\..*",
61 r"re:mtp.layers.*",
62 "backbone.norm_f",
63 "lm_head",
64 ],
65)
66
67oneshot(
68 model=model,
69 dataset=ds,
70 recipe=recipe,
71 max_seq_length=MAX_SEQUENCE_LENGTH,
72 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
73)
74
75print("========== SAMPLE GENERATION ==============")
76dispatch_model(model)
77input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to(
78 model.device
79)
80output = model.generate(input_ids, max_new_tokens=20)
81print(tokenizer.decode(output[0]))
82print("==========================================")
83
84SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8"
85model.save_pretrained(SAVE_DIR)
86tokenizer.save_pretrained(SAVE_DIR)
87
881vllm serve RedHatAI/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-FP8 \
2 --mamba-backend flashinfer \
3 --mamba-cache-mode align \
4 --enable-prefix-caching \
5 --max-num-batched-tokens 16384 \
6 --moe-backend flashinfer_cutlass \
7 --tensor-parallel-size 1 \
8 --reasoning-parser nemotron_v3 \
9 --enable-auto-tool-choice \
10 --tool-call-parser qwen3_xml \
11 --mamba-ssm-cache-dtype float16 \
12 --enable-mamba-cache-stochastic-rounding \
13 --mamba-cache-philox-rounds 5
14