File size: 6,846 Bytes
000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 7484310 000c932 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """
NVFP4 PTQ of llm-jp-4-8b-instruct with NVIDIA Model Optimizer — the recipe
used to produce this checkpoint.
Mixed precision:
- MLP linears (gate/up/down): NVFP4 (FP4 block-16, FP8 block scales)
- Attention projections (q/k/v/o): FP8 (E4M3 per-tensor)
- First 2 and last 2 decoder layers, embeddings, lm_head, norms: BF16
Calibration: calib_full.jsonl (bundled) — chat-templated conversations plus
raw-text strict-format samples. Conversation samples go through
tokenizer.apply_chat_template so activation ranges see deployment-realistic
Harmony token streams; {"text": ...} samples are encoded as-is (raw
zero-shot completion form).
Requires nvidia-modelopt==0.45.0. Run with the GPU otherwise idle (needs
~20GB+ headroom for the BF16 model plus calibration activations):
python3 03_ptq_modelopt.py --calib calib_full.jsonl --export-dir <export-dir>
"""
import argparse
import copy
import json
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.quantization as mtq
from modelopt.torch.export import export_hf_checkpoint
DEFAULT_MODEL = "llm-jp/llm-jp-4-8b-instruct"
DEFAULT_REVISION = "098f2b2cf33021eba19a6d3582aa3d071ccc0aff" # the revision this checkpoint was built from
BF16_LAYERS = (0, 1, 30, 31) # Nemotron-style: first/last 2 layers stay BF16
def build_calib_batches(calib_path: str, tokenizer, n_samples: int, max_seq_len: int):
"""Tokenize calibration samples.
NOTE: apply_chat_template(tokenize=True)'s return type varies across
transformers versions (list vs BatchEncoding — the latter silently broke
len() checks on 5.5.x). Always template to a string, then encode.
"""
batches = []
with open(calib_path) as f:
for i, line in enumerate(f):
if i >= n_samples:
break
rec = json.loads(line)
if "messages" in rec:
text = tokenizer.apply_chat_template(
rec["messages"], add_generation_prompt=False, tokenize=False
)
else:
# Raw-completion sample (strict-format slice): encode as-is —
# these prompts are served with no chat template.
text = rec["text"]
ids = tokenizer.encode(text, add_special_tokens=False)[:max_seq_len]
batches.append(torch.tensor([ids], dtype=torch.long))
return batches
def make_forward_loop(batches):
def forward_loop(model):
with torch.no_grad():
for i, input_ids in enumerate(batches):
model(input_ids=input_ids.to(model.device))
if (i + 1) % 64 == 0:
print(f" calib {i + 1}/{len(batches)}", flush=True)
return forward_loop
def build_quant_config():
# modelopt 0.45's quant_cfg is an ORDERED LIST of {quantizer_name, ...}
# entries where later entries override earlier pattern matches.
# NVFP4_DEFAULT_CFG already disables lm_head, routers, MoE gates, and
# norm/BN layers.
cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG)
# Attention projections to per-tensor FP8 (E4M3); MLP linears stay NVFP4
# from the base config. Appended entries win over the base patterns.
fp8_w = {"num_bits": (4, 3), "axis": None}
for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
cfg["quant_cfg"].append(
{"quantizer_name": f"*self_attn.{proj}*weight_quantizer", "cfg": dict(fp8_w)}
)
cfg["quant_cfg"].append(
{"quantizer_name": f"*self_attn.{proj}*input_quantizer", "cfg": dict(fp8_w)}
)
# Most quantization-sensitive layers stay BF16.
for layer in BF16_LAYERS:
cfg["quant_cfg"].append({"quantizer_name": f"*layers.{layer}.*", "enable": False})
# Embeddings are not in the default ignore list — exclude explicitly.
cfg["quant_cfg"].append({"quantizer_name": "*embed_tokens*", "enable": False})
return cfg
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default=DEFAULT_MODEL,
help="HF id or local path of the BF16 base model")
parser.add_argument("--revision", default=DEFAULT_REVISION,
help="HF revision of the base model (pinned to the "
"revision this checkpoint was built from; "
"ignored for local paths)")
parser.add_argument("--calib", default="calib_full.jsonl")
parser.add_argument("--n-samples", type=int, default=588)
parser.add_argument("--max-seq-len", type=int, default=4096)
parser.add_argument("--export-dir", required=True)
args = parser.parse_args()
print(f"Loading tokenizer + model from {args.model}...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(
args.model, revision=args.revision, trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
args.model, revision=args.revision, torch_dtype=torch.bfloat16,
device_map="cuda", trust_remote_code=True,
)
model.eval()
calib_path = Path(args.calib)
if not calib_path.exists():
# default calib_full.jsonl ships next to this script — resolve
# relative to the script dir so the command works from anywhere
calib_path = Path(__file__).parent / args.calib
print(f"Building calibration batches from {calib_path}...", flush=True)
batches = build_calib_batches(str(calib_path), tokenizer, args.n_samples, args.max_seq_len)
print(f" {len(batches)} samples, total tokens "
f"{sum(b.numel() for b in batches)}", flush=True)
print("Quantizing...", flush=True)
model = mtq.quantize(model, build_quant_config(), make_forward_loop(batches))
# Mandatory gate: confirm what actually got quantized before any
# export/eval investment. Saved alongside the checkpoint for review.
export_dir = Path(args.export_dir)
export_dir.mkdir(parents=True, exist_ok=True)
import contextlib, io
buf = io.StringIO()
with contextlib.redirect_stdout(buf):
mtq.print_quant_summary(model)
summary = buf.getvalue()
(export_dir / "quant_summary.txt").write_text(summary)
print(summary[:4000], flush=True)
for banned in ("lm_head", "embed_tokens"):
for line in summary.splitlines():
if banned in line and "TensorQuantizer" in line and "disabled" not in line:
raise SystemExit(f"GATE FAILED: {banned} appears quantized: {line}")
print("Quant summary gate passed (lm_head/embed_tokens not quantized).", flush=True)
print(f"Exporting HF checkpoint to {export_dir}...", flush=True)
export_hf_checkpoint(model, export_dir=str(export_dir))
print("Done. Now run 04_export_and_patch.py on the export dir.", flush=True)
if __name__ == "__main__":
main()
|