File size: 5,588 Bytes
a64dea9 | 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 | """Find a (quant_method, format) pair that lets transformers parse this
modelopt MIXED_PRECISION checkpoint, and check that per-layer schemes resolve
to the correct precision for BOTH the FP8 and the NVFP4 tensors."""
import gc
import json
import os
import shutil
import traceback
SRC = os.environ.get("DWCT_NEMOTRON", "/models/nemotron")
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".qprobe"))
print("=" * 72)
print("environment")
print("=" * 72)
try:
import transformers
print("transformers", transformers.__version__)
except Exception as exc:
print("transformers ERR", exc)
try:
import compressed_tensors as ct
print("compressed_tensors", getattr(ct, "__version__", "?"))
except Exception as exc:
print("compressed_tensors IMPORT FAILED:", type(exc).__name__, exc)
try:
from compressed_tensors.config import QuantizationFormat
print("QuantizationFormat values:", [f.value for f in QuantizationFormat])
except Exception as exc:
print("QuantizationFormat ERR", exc)
for mod in ("quantizer_compressed_tensors", "quantizer_nvfp4", "quantizer_finegrained_fp8"):
try:
m = __import__("transformers.quantizers." + mod, fromlist=["*"])
print(mod, "->", [n for n in dir(m) if not n.startswith("_") and ("Quant" in n or "Scheme" in n)])
except Exception as exc:
print(mod, "ERR", exc)
from transformers import PreTrainedModel
if not isinstance(getattr(PreTrainedModel, "all_tied_weights_keys", None), property):
def _g(self):
return getattr(self, "_all_tied_weights_keys", None) or getattr(self, "_tied_weights_keys", None) or []
def _s(self, v):
self._all_tied_weights_keys = v
PreTrainedModel.all_tied_weights_keys = property(_g, _s)
PROBE_LAYERS = [
"language_model.backbone.layers.1.mixer.experts.0.up_proj", # NVFP4
"language_model.backbone.layers.1.mixer.shared_experts.up_proj", # FP8
"language_model.backbone.layers.0.mixer.in_proj", # FP8
"language_model.backbone.embeddings", # dense
]
CANDIDATES = [
("compressed-tensors", "nvfp4-pack-quantized"),
("compressed-tensors", "mixed-precision"),
("nvfp4", "nvfp4-pack-quantized"),
("nvfp4", "mixed-precision"),
("compressed-tensors", "float-quantized"),
]
def make_overlay(tag, method, fmt):
dst = os.path.join(ROOT, tag)
os.makedirs(dst, exist_ok=True)
for name in os.listdir(SRC):
s, d = os.path.join(SRC, name), os.path.join(dst, name)
if os.path.lexists(d):
continue
try:
os.symlink(s, d)
except OSError:
shutil.copytree(s, d) if os.path.isdir(s) else shutil.copy2(s, d)
cfg = json.load(open(os.path.join(SRC, "config.json")))
q = cfg["quantization_config"]
q["quant_method"] = method
q["format"] = fmt
q["quantization_status"] = "compressed"
cfg["quantization_config"] = q
p = os.path.join(dst, "config.json")
if os.path.islink(p):
os.unlink(p)
json.dump(cfg, open(p, "w"), indent=2)
return dst
def scheme_of(qz, name):
for obj in (qz, getattr(qz, "quantization_config", None)):
fn = getattr(obj, "get_scheme", None)
if fn is None:
continue
try:
return fn(name)
except Exception as exc:
return "ERR " + repr(exc)[:140]
return "no get_scheme"
def try_candidate(method, fmt):
tag = ("%s__%s" % (method, fmt)).replace("-", "_")
print("-" * 72)
print("CANDIDATE quant_method=%r format=%r" % (method, fmt))
try:
d = make_overlay(tag, method, fmt)
from transformers import AutoConfig
from transformers.quantizers import auto as qauto
cfg = AutoConfig.from_pretrained(d, trust_remote_code=True)
qc = getattr(cfg, "quantization_config", None)
print(" parsed config class:", type(qc).__name__)
get_q = getattr(qauto, "get_quantizer", None)
qz = get_q(qc) if get_q else qauto.AutoQuantizer.from_config(qc)
print(" quantizer:", type(qz).__name__)
for name in PROBE_LAYERS:
sch = scheme_of(qz, name)
label = type(sch).__name__ if not isinstance(sch, str) else sch
print(" scheme(%-52s) -> %s" % (name.split("backbone.")[-1], label))
if isinstance(sch, str) and sch.startswith("ERR"):
print(" RESULT: scheme resolution failed")
return False
print(" building on meta ...")
import torch
from transformers import AutoModelForCausalLM
with torch.device("meta"):
model = AutoModelForCausalLM.from_pretrained(
d, trust_remote_code=True, dtype=torch.bfloat16,
attn_implementation="eager",
)
try:
mod = model.language_model.backbone.layers[1].mixer.experts[0].up_proj
print(" expert module:", type(mod).__name__)
except Exception as exc:
print(" expert module lookup ERR", exc)
del model
gc.collect()
print(" RESULT: OK")
return True
except Exception as exc:
print(" RESULT: FAIL", type(exc).__name__, repr(exc)[:300])
traceback.print_exc(limit=2)
return False
winners = []
for method, fmt in CANDIDATES:
try:
if try_candidate(method, fmt):
winners.append((method, fmt))
except Exception:
traceback.print_exc(limit=3)
print("=" * 72)
print("WORKING COMBINATIONS:", winners)
|