Spaces:
Paused
Paused
| """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) | |