"""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)