File size: 1,797 Bytes
8e8ea9a | 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 | """Remap compressed-tensors ignore paths for Qwen3.8's standalone MTP layer."""
from __future__ import annotations
import argparse
from pathlib import Path
MARKER = "qwen38-flash-dgx: remap compressed-tensors MTP ignore paths"
_METHOD_ANCHOR = "def _make_draft_vllm_config(\n"
_INSERTION_ANCHOR = " draft_vllm_config = replace(\n"
def patch_source(source: str) -> str:
if MARKER in source:
return source
for label, anchor in (
("Qwen3.8 draft-config method", _METHOD_ANCHOR),
("Qwen3.8 draft-config replacement", _INSERTION_ANCHOR),
):
occurrences = source.count(anchor)
if occurrences != 1:
raise RuntimeError(f"expected exactly one {label} hook, found {occurrences}")
block = f''' # {MARKER}
compressed_tensors_ignore = getattr(draft_quant_config, "ignore", None)
if compressed_tensors_ignore:
setattr( # noqa: B010
draft_quant_config,
"ignore",
_remap_ignored_layers(
compressed_tensors_ignore, mtp_start_layer_idx
),
)
'''
return source.replace(_INSERTION_ANCHOR, block + _INSERTION_ANCHOR, 1)
def main(argv=None) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("model_path", type=Path)
args = parser.parse_args(argv)
source = args.model_path.read_text(encoding="utf-8")
patched = patch_source(source)
if patched != source:
args.model_path.write_text(patched, encoding="utf-8")
print(f"patched compressed-tensors MTP ignore remap: {args.model_path}")
else:
print(f"compressed-tensors MTP ignore remap already present: {args.model_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|