| """Make the modelopt MIXED_PRECISION Nemotron checkpoint loadable. |
| |
| /models/nemotron declares quant_method="modelopt": 98 static-FP8 tensors |
| (group_0) plus 5888 NVFP4 group_size-16 expert tensors (group_1). transformers |
| 5.20.1 has no modelopt quantizer, so it builds plain bf16 layers -- every packed |
| NVFP4 expert comes back at half width, and 30B params in bf16 is ~60 GB against |
| a 50.8 GB slice. |
| |
| Fix: rewrite quant_method/format into the compressed-tensors schema transformers |
| does understand, keeping config_groups exactly as they are (group_0 really is |
| 8-bit, group_1 really is 4-bit), and serve it from a writable overlay of |
| symlinks so /models stays read-only. |
| |
| The (quant_method, format) pair is not hardcoded: we ask transformers to build |
| the quantizer for each candidate and keep the first one where a known NVFP4 |
| layer resolves to a 4-bit scheme AND a known FP8 layer resolves to an 8-bit one. |
| install() monkeypatches the Auto* entry points, so no other file needs editing. |
| """ |
| import copy |
| import json |
| import os |
| import shutil |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| WORK = os.environ.get("DWCT_WORK", os.path.join(os.path.dirname(_HERE), ".quantfix")) |
| SRC_DEFAULT = os.environ.get("DWCT_NEMOTRON", "/models/nemotron") |
|
|
| EXPERT_LAYER = "language_model.backbone.layers.1.mixer.experts.0.up_proj" |
| FP8_LAYER = "language_model.backbone.layers.1.mixer.shared_experts.up_proj" |
|
|
| CANDIDATES = [ |
| ("compressed-tensors", "mixed-precision"), |
| ("compressed-tensors", "nvfp4-pack-quantized"), |
| ("compressed-tensors", "pack-quantized"), |
| ("nvfp4", "nvfp4-pack-quantized"), |
| ("nvfp4", "mixed-precision"), |
| ] |
|
|
| _choice = {} |
| _installed = False |
|
|
|
|
| def _log(*a): |
| print("[quantfix]", *a) |
|
|
|
|
| def _build_quantizer(qdict): |
| from transformers.quantizers import auto as qauto |
| cls = getattr(qauto, "AutoHfQuantizer", None) |
| if cls is not None and hasattr(cls, "from_config"): |
| return cls.from_config(qdict) |
| fn = getattr(qauto, "get_quantizer", None) |
| if fn is not None: |
| return fn(qdict) |
| raise RuntimeError("no quantizer factory in transformers.quantizers.auto") |
|
|
|
|
| def _scheme(qz, name): |
| cands = [qz, getattr(qz, "quantization_config", None)] |
| try: |
| from compressed_tensors.config import CompressedTensorsConfig as _C |
| raw = getattr(qz, "quantization_config", None) |
| if isinstance(raw, dict): |
| cands.append(_C.from_dict(raw)) |
| except Exception: |
| pass |
| for obj in cands: |
| fn = getattr(obj, "get_scheme", None) |
| if fn is None: |
| continue |
| for extra in range(0, 4): |
| try: |
| return fn(*([name] + [None] * extra)) |
| except TypeError: |
| continue |
| except Exception as exc: |
| return "ERR " + repr(exc)[:140] |
| return None |
|
|
|
|
| def _verdict(qz): |
| exp = _scheme(qz, EXPERT_LAYER) |
| fp8 = _scheme(qz, FP8_LAYER) |
| def _label(x): |
| if x is None: |
| return None |
| return x if isinstance(x, str) else type(x).__name__ |
|
|
| names = {"expert": _label(exp), "fp8": _label(fp8)} |
| if exp is None or fp8 is None: |
| return False, names |
| e = (names["expert"] or "").lower() |
| f = (names["fp8"] or "").lower() |
| ok = ("fp4" in e or "nvfp4" in e or "w4" in e) and ("fp8" in f or "w8" in f) |
| return ok, names |
|
|
|
|
| def choose(src): |
| """Return ((method, fmt), patched_qc_dict, log_lines).""" |
| base = json.load(open(os.path.join(src, "config.json"))) |
| qc = base.get("quantization_config") |
| if not qc: |
| return None, None, ["no quantization_config"] |
| if "modelopt" in (qc.get("quant_method") or ""): |
| _log("checkpoint is modelopt; rewriting into compressed-tensors schema") |
|
|
| env_m = os.environ.get("DWCT_QUANT_METHOD") |
| env_f = os.environ.get("DWCT_QUANT_FORMAT") |
| cands = [(env_m, env_f)] if env_m and env_f else list(CANDIDATES) |
|
|
| log, validated, built, last = [], None, None, None |
| for method, fmt in cands: |
| q = copy.deepcopy(qc) |
| q["quant_method"] = method |
| q["format"] = fmt |
| q["quantization_status"] = "compressed" |
| q.setdefault("ignore", []) |
| if os.environ.get("DWCT_KV_SCHEME", "0") != "1": |
| q.pop("kv_cache_scheme", None) |
| last = ((method, fmt), q) |
| try: |
| qz = _build_quantizer(q) |
| except Exception as exc: |
| log.append("%s/%s: quantizer build failed (%s: %s)" |
| % (method, fmt, type(exc).__name__, str(exc)[:140])) |
| continue |
| if built is None: |
| built = ((method, fmt), q) |
| ok, names = _verdict(qz) |
| got = getattr(getattr(qz, "quantization_config", None), "format", None) |
| got = getattr(got, "value", got) |
| log.append("%s/%s: expert=%s fp8=%s fmt_readback=%r" |
| % (method, fmt, names["expert"], names["fp8"], got)) |
| if ok: |
| validated = ((method, fmt), q) |
| break |
|
|
| for line in log: |
| _log(" ", line) |
| if validated is not None: |
| return validated[0], validated[1], log |
| if built is not None: |
| _log("no candidate validated schemes; using first that built") |
| return built[0], built[1], log |
| _log("no candidate built; using first rewritten config") |
| return last[0], last[1], log |
|
|
|
|
| def _hoist_heads(cfg): |
| """compressed-tensors' KV-cache path reads num_attention_heads off the |
| top-level config, but the composite wrapper keeps it inside llm_config.""" |
| llm = cfg.get("llm_config") or {} |
| for key in ("num_attention_heads", "head_dim"): |
| if cfg.get(key) is None and llm.get(key) is not None: |
| cfg[key] = llm[key] |
| _log("hoisted %s=%s from llm_config" % (key, llm[key])) |
|
|
|
|
| def build(src, dst): |
| src, dst = os.path.abspath(src), os.path.abspath(dst) |
| 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) |
|
|
| pair, q, log = choose(src) |
| cfg = json.load(open(os.path.join(src, "config.json"))) |
| if q is not None: |
| cfg["quantization_config"] = q |
| _hoist_heads(cfg) |
| cp = os.path.join(dst, "config.json") |
| if os.path.islink(cp): |
| os.unlink(cp) |
| json.dump(cfg, open(cp, "w"), indent=2) |
| return dst, pair |
|
|
|
|
| def load_dir(src=None): |
| src = os.path.abspath(src or SRC_DEFAULT) |
| if src in _choice: |
| return _choice[src] |
| if not os.path.isdir(src) or not os.path.exists(os.path.join(src, "config.json")): |
| _choice[src] = src |
| return src |
| dst = os.path.join(WORK, "cfg-" + os.path.basename(src)) |
| try: |
| dst, pair = build(src, dst) |
| _log("overlay:", dst) |
| _log("chosen:", pair) |
| except Exception as exc: |
| _log("overlay failed (%s: %s); using %s" % (type(exc).__name__, exc, src)) |
| dst = src |
| _choice[src] = dst |
| return dst |
|
|
|
|
| def _maybe(path): |
| if not isinstance(path, str): |
| return path |
| try: |
| ap = os.path.abspath(path) |
| except Exception: |
| return path |
| if ap in _choice: |
| return _choice[ap] |
| if os.path.basename(ap) == "nemotron" or ap == os.path.abspath(SRC_DEFAULT): |
| return load_dir(ap) |
| return path |
|
|
|
|
| def _patch(cls): |
| raw = getattr(cls.from_pretrained, "__func__", cls.from_pretrained) |
| if getattr(raw, "_quantfix", False): |
| return |
|
|
| def from_pretrained(cls_, pretrained_model_name_or_path=None, *a, **kw): |
| if pretrained_model_name_or_path is None: |
| pretrained_model_name_or_path = kw.pop("pretrained_model_name_or_path", None) |
| return raw(cls_, _maybe(pretrained_model_name_or_path), *a, **kw) |
|
|
| from_pretrained._quantfix = True |
| cls.from_pretrained = classmethod(from_pretrained) |
|
|
|
|
| def install(): |
| global _installed |
| if _installed or os.environ.get("DWCT_QUANTFIX", "1") == "0": |
| return |
| _installed = True |
| try: |
| import transformers |
| except Exception: |
| return |
| try: |
| from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING as M |
| if "modelopt" in M: |
| _log("transformers already knows modelopt; leaving config alone") |
| return |
| except Exception: |
| pass |
| for name in ("AutoConfig", "AutoModel", "AutoModelForCausalLM", |
| "AutoModelForImageTextToText", "AutoProcessor", |
| "AutoTokenizer", "PreTrainedModel"): |
| try: |
| cls = getattr(transformers, name, None) |
| if cls is not None: |
| _patch(cls) |
| except Exception as exc: |
| _log("patch %s failed: %r" % (name, exc)) |
| _log("installed (work=%s)" % WORK) |
|
|
