modelopt.patch 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. --- a/vllm/modelopt.py
  2. +++ b/vllm/modelopt.py
  3. @@ -2356,6 +2356,22 @@
  4. def _quantized_layer_prefix_candidates(prefix: str) -> tuple[str, ...]:
  5. candidates = [prefix]
  6. + # MTP draft naming (see nvfp4 stand-up notes): strip mtp_block,
  7. + # then re-root over every naming the shared quantized_layers dict
  8. + # can be in - checkpoint form (model.language_model.*) or the
  9. + # target-tree form its mapper rewrites to (language_model.model.*).
  10. + _bases = {prefix, prefix.replace(".mtp_block.", ".")}
  11. + for _b in list(_bases):
  12. + for _root in ("model.", "model.language_model.", "language_model.model."):
  13. + if _b.startswith(_root):
  14. + _bases.add(_b[len(_root) :])
  15. + for _b in list(_bases):
  16. + if _b.startswith("layers."):
  17. + _bases.add("model." + _b)
  18. + _bases.add("model.language_model." + _b)
  19. + _bases.add("language_model.model." + _b)
  20. + candidates.extend(_bases - {prefix})
  21. +
  22. if prefix.endswith(".lm_head"):
  23. candidates.append("lm_head")
  24. @@ -2387,6 +2403,13 @@
  25. return None
  26. quant_algo = self._resolve_quant_algo(prefix)
  27. + import sys as _sys
  28. + print(
  29. + "[quantprobe] layer=%s prefix=%s algo=%s"
  30. + % (type(layer).__name__, prefix, quant_algo),
  31. + file=_sys.stderr,
  32. + flush=True,
  33. + )
  34. if isinstance(layer, (LinearBase, ParallelLMHead)):
  35. if quant_algo == "FP8":