File size: 1,642 Bytes
62274c6 | 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 57 | """Give Qwen3.8 Flash-Next output heads their configured quantization method."""
from __future__ import annotations
import argparse
import re
from pathlib import Path
MARKER = "qwen38-flash-dgx: support a quantized lm_head checkpoint"
CONSTRUCTOR_HOOK = re.compile(
r'^(?P<indent>[ \t]+)config\.hidden_size,\n'
r'(?P=indent)prefix=maybe_prefix\(prefix, "lm_head"\),$',
re.MULTILINE,
)
def patch_source(source: str) -> str:
if MARKER in source:
return source
def replacement(match: re.Match) -> str:
indent = match.group("indent")
return (
f"{indent}config.hidden_size,\n"
f"{indent}# {MARKER}\n"
f"{indent}quant_config=self.quant_config,\n"
f'{indent}prefix="lm_head",'
)
patched, occurrences = CONSTRUCTOR_HOOK.subn(replacement, source)
if occurrences != 1:
raise RuntimeError(
"expected exactly one Qwen output-head constructor hook, "
f"found {occurrences}"
)
return patched
def main(argv=None) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("model_paths", type=Path, nargs="+")
args = parser.parse_args(argv)
for path in args.model_paths:
source = path.read_text(encoding="utf-8")
patched = patch_source(source)
if patched != source:
path.write_text(patched, encoding="utf-8")
print(f"patched quantized Qwen output head: {path}")
else:
print(f"quantized Qwen output-head support already present: {path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|