From 312ae4d98bb6c688e1175fc9dd6160248cb72d5f Mon Sep 17 00:00:00 2001 From: krizaltang Date: Wed, 1 Jul 2026 19:02:28 +0800 Subject: [PATCH] [feat]:Support skip layers for int8 quantization. --- tools/int8_channel_quant.py | 253 +++++++++++++++++++++++++++++++----- 1 file changed, 221 insertions(+), 32 deletions(-) diff --git a/tools/int8_channel_quant.py b/tools/int8_channel_quant.py index 7200821e..6c9ca186 100644 --- a/tools/int8_channel_quant.py +++ b/tools/int8_channel_quant.py @@ -16,6 +16,7 @@ import json import os +import re import shutil from argparse import ArgumentParser @@ -68,7 +69,132 @@ def get_suffix_to_quant(model_type): return SUFFIX_TO_QUANT -def build_ignored_layers(input_path): +def parse_skip_layers(skip_specs): + """ + Parse ``--skip-layers`` CLI values into a list of matcher rules. + + Each rule is a tuple ``(kind, pattern)`` where ``kind`` is one of: + + * ``"prefix"`` – ``pattern`` is a dotted path prefix; a weight name matches + if it equals ``pattern`` (as a layer name plus ``.weight``) or if it + starts with ``pattern + "."``. + * ``"regex"`` – ``pattern`` is a compiled regex object; matches a weight + name via :meth:`re.Pattern.search`. + + Accepted input forms (each spec may also contain comma-separated items): + + * Bare integer, e.g. ``77`` -> prefix ``model.layers.77`` + * ``layers.77`` / ``layer.77`` -> prefix ``model.layers.77`` + * Explicit dotted path, e.g. ``model.layers.0.self_attn.q_proj`` + -> prefix that exact module (its ``.weight`` and any children). + * Trailing ``*``, e.g. ``model.layers.5.*`` -> prefix ``model.layers.5`` + * ``re:`` -> regex over the full weight name. + """ + rules = [] + if not skip_specs: + return rules + raw_items = [] + for spec in skip_specs: + if spec is None: + continue + for part in str(spec).split(","): + part = part.strip() + if part: + raw_items.append(part) + + for item in raw_items: + if item.startswith("re:"): + rules.append(("regex", re.compile(item[3:]))) + continue + # Trailing '*' -> prefix match; strip the star (and optional trailing dot). + if item.endswith("*"): + prefix = item[:-1].rstrip(".") + if prefix: + rules.append(("prefix", prefix)) + continue + # Pure integer (or "layers.N" / "layer.N") -> model.layers.N.* + m = re.fullmatch(r"(?:layers?\.)?(\d+)", item) + if m: + rules.append(("prefix", f"model.layers.{m.group(1)}")) + continue + rules.append(("prefix", item.rstrip("."))) + return rules + + +def make_skip_matcher(skip_rules): + """Return a picklable callable ``matcher(weight_name) -> bool``. + + ``skip_rules`` is the output of :func:`parse_skip_layers`. Regex rules + stored inside are re-compiled inside the callable so the result is safe + to pass across ``multiprocessing`` process boundaries (compiled regex + objects pickle fine, but we normalize to string form for robustness). + """ + normalized = [] + for kind, pat in skip_rules: + if kind == "regex": + normalized.append(("regex", pat.pattern)) + else: + normalized.append((kind, pat)) + + return _SkipMatcher(normalized) + + +class _SkipMatcher: + """Picklable matcher used by worker processes.""" + + def __init__(self, rules): + self._rules = rules + self._compiled = None + + def _ensure_compiled(self): + if self._compiled is not None: + return + compiled = [] + for kind, pat in self._rules: + if kind == "regex": + compiled.append(("regex", re.compile(pat))) + else: + compiled.append(("prefix", pat)) + self._compiled = compiled + + def __bool__(self): + return bool(self._rules) + + def __call__(self, weight_name): + if not self._rules: + return False + self._ensure_compiled() + for kind, pat in self._compiled: + if kind == "regex": + if pat.search(weight_name): + return True + else: + # prefix: exact ``pat.weight`` OR anything under ``pat.`` + if weight_name == f"{pat}.weight": + return True + if weight_name.startswith(pat + "."): + return True + return False + + +def expand_skip_layer_names(skip_rules, all_layer_names): + """Expand skip rules to concrete layer names present in the model. + + Used for populating ``ignored_layers`` in ``config.json`` so the runtime + (vLLM / compressed-tensors) also treats these layers as un-quantized. + """ + if not skip_rules: + return [] + matcher = make_skip_matcher(skip_rules) + hits = [] + for name in all_layer_names: + # Test as if it were a real weight tensor. + if matcher(f"{name}.weight"): + hits.append(name) + return hits + + +def build_ignored_layers(input_path, skip_rules=None): """Build ignored layers from the model structure, following fp8_quant_blockwise.py.""" hf_config = AutoConfig.from_pretrained(input_path) model_type = hf_config.model_type @@ -85,6 +211,12 @@ def build_ignored_layers(input_path): print(f"Found {len(layers)} linear layers") ignored_layers = [] + # User-specified skip layers -> add them to ignored_layers so runtime + # loaders (vLLM / compressed-tensors) also treat them as un-quantized. + user_skipped = expand_skip_layer_names(skip_rules, list(layers.keys())) + if user_skipped: + print(f"User-specified skip layers ({len(user_skipped)}): {user_skipped}") + ignored_layers.extend(user_skipped) if model_type in ("qwen3_5_moe", "qwen3_5"): for name, module in model.named_modules(): if not hasattr(module, "weight") or module.weight is None: @@ -110,7 +242,14 @@ def build_ignored_layers(input_path): ignored_layers.append(name) del model - return ignored_layers, model_type + # De-duplicate while preserving order. + seen = set() + deduped = [] + for name in ignored_layers: + if name not in seen: + seen.add(name) + deduped.append(name) + return deduped, model_type def _quant_and_record_int8(weight_name, weight_bf16, new_state_dict, new_weight_map, file_name): @@ -132,6 +271,7 @@ def process_worker( return_dict, suffix_to_quant, input_type="bf16", + skip_matcher=None, ): """ Process worker. @@ -145,6 +285,7 @@ def process_worker( rank = worker_id % num_gpus torch.cuda.set_device(rank) quant_count = 0 + skipped_count = 0 new_weight_map = {} for safetensor_file in safetensor_files: file_name = os.path.basename(safetensor_file) @@ -155,6 +296,13 @@ def process_worker( for weight_name in keys: weight = f.get_tensor(weight_name) if any(weight_name.endswith(suffix) for suffix in suffix_to_quant): + if skip_matcher is not None and skip_matcher(weight_name): + # User asked to skip this layer -> keep original weight, + # do NOT emit a scale, and drop any incoming scale_inv. + skipped_count += 1 + new_state_dict[weight_name] = weight + new_weight_map[weight_name] = file_name + continue quant_count += 1 if input_type == "fp8": scale_inv_name = f"{weight_name}_scale_inv" @@ -179,7 +327,7 @@ def process_worker( new_safetensor_file = os.path.join(int8_path, file_name) save_file(new_state_dict, new_safetensor_file) - return_dict[worker_id] = (quant_count, new_weight_map) + return_dict[worker_id] = (quant_count, new_weight_map, skipped_count) def process_worker_qwen35( @@ -191,6 +339,7 @@ def process_worker_qwen35( return_dict, suffix_to_quant, input_type="bf16", + skip_matcher=None, ): """ Qwen3.5-specific worker. @@ -207,6 +356,7 @@ def process_worker_qwen35( rank = worker_id % num_gpus torch.cuda.set_device(rank) quant_count = 0 + skipped_count = 0 new_weight_map = {} for safetensor_file in safetensor_files: file_name = os.path.basename(safetensor_file) @@ -229,21 +379,20 @@ def process_worker_qwen35( for i in range(num_experts): gate_name = f"{prefix}.experts.{i}.gate_proj.weight" up_name = f"{prefix}.experts.{i}.up_proj.weight" - _quant_and_record_int8( - gate_name, - gate_w[i].contiguous(), - new_state_dict, - new_weight_map, - file_name, - ) - _quant_and_record_int8( - up_name, - up_w[i].contiguous(), - new_state_dict, - new_weight_map, - file_name, - ) - quant_count += 2 + for tname, tval in ((gate_name, gate_w[i]), (up_name, up_w[i])): + if skip_matcher is not None and skip_matcher(tname): + skipped_count += 1 + new_state_dict[tname] = tval.contiguous() + new_weight_map[tname] = file_name + else: + _quant_and_record_int8( + tname, + tval.contiguous(), + new_state_dict, + new_weight_map, + file_name, + ) + quant_count += 1 del weight, weight_bf16, gate_w, up_w torch.cuda.empty_cache() continue @@ -254,14 +403,19 @@ def process_worker_qwen35( prefix = weight_name[: -len(".experts.down_proj")] for i in range(num_experts): down_name = f"{prefix}.experts.{i}.down_proj.weight" - _quant_and_record_int8( - down_name, - weight_bf16[i].contiguous(), - new_state_dict, - new_weight_map, - file_name, - ) - quant_count += 1 + if skip_matcher is not None and skip_matcher(down_name): + skipped_count += 1 + new_state_dict[down_name] = weight_bf16[i].contiguous() + new_weight_map[down_name] = file_name + else: + _quant_and_record_int8( + down_name, + weight_bf16[i].contiguous(), + new_state_dict, + new_weight_map, + file_name, + ) + quant_count += 1 del weight, weight_bf16 torch.cuda.empty_cache() continue @@ -270,6 +424,11 @@ def process_worker_qwen35( # Regular tensors # ------------------------------------------------------------ if any(weight_name.endswith(suffix) for suffix in suffix_to_quant): + if skip_matcher is not None and skip_matcher(weight_name): + skipped_count += 1 + new_state_dict[weight_name] = weight + new_weight_map[weight_name] = file_name + continue quant_count += 1 weight_bf16 = weight _quant_and_record_int8( @@ -287,7 +446,7 @@ def process_worker_qwen35( new_safetensor_file = os.path.join(int8_path, file_name) save_file(new_state_dict, new_safetensor_file) - return_dict[worker_id] = (quant_count, new_weight_map) + return_dict[worker_id] = (quant_count, new_weight_map, skipped_count) # Helper function to get tensor from the correct file @@ -332,7 +491,7 @@ def weight_quant(tensor: torch.Tensor): return quantized.to(torch.int8), scale.to(torch.float32) -def main(input_path, int8_path, num_workers): +def main(input_path, int8_path, num_workers, skip_layers=None): """ Run the FP8-to-INT8 per-channel quantization pipeline. @@ -395,8 +554,13 @@ def main(input_path, int8_path, num_workers): weight_map = {name: "model.safetensors" for name in f.keys()} print(f"Found {len(safetensor_files)} safetensor files") - ignored_layers, model_type = build_ignored_layers(input_path) - print(f"Ignored layers: {ignored_layers}") + skip_rules = parse_skip_layers(skip_layers) + if skip_rules: + print(f"Skip rules ({len(skip_rules)}): {skip_rules}") + skip_matcher = make_skip_matcher(skip_rules) if skip_rules else None + + ignored_layers, model_type = build_ignored_layers(input_path, skip_rules=skip_rules) + print(f"Ignored layers ({len(ignored_layers)}): {ignored_layers}") is_qwen35 = model_type in ("qwen3_5_moe", "qwen3_5") config["quantization_config"] = { @@ -443,6 +607,7 @@ def main(input_path, int8_path, num_workers): print(f"config.json modified and saved to {config_file}") quant_count = 0 + skipped_count = 0 new_weight_map = {} file_subsets = [safetensor_files[i::num_workers] for i in range(num_workers)] @@ -464,6 +629,7 @@ def main(input_path, int8_path, num_workers): return_dict, suffix_to_quant, input_type, + skip_matcher, ), ) p.start() @@ -472,10 +638,13 @@ def main(input_path, int8_path, num_workers): p.join() for i in range(num_workers): - qc, wm = return_dict[i] + qc, wm, sc = return_dict[i] quant_count += qc + skipped_count += sc new_weight_map.update(wm) print(f"{quant_count} weights are quantized.") + if skip_matcher is not None: + print(f"{skipped_count} weights are skipped (kept original) by --skip-layers.") if has_index: # modify model.safetensors.index.json @@ -494,7 +663,27 @@ def main(input_path, int8_path, num_workers): parser.add_argument("--input-path", type=str, required=True) parser.add_argument("--output-int8-path", type=str, required=True) parser.add_argument("--num-workers", type=int, default=32) + parser.add_argument( + "--skip-layers", + type=str, + nargs="*", + default=None, + help=( + "Layer specs to skip (keep original weights, no INT8 quantization). " + "Accepts multiple values and/or comma-separated lists. Forms: " + "'77' -> model.layers.77.*; " + "'model.layers.0.self_attn.q_proj' -> that exact module and its children; " + "'model.layers.5.*' -> prefix match; " + "'re:' -> regex over the full weight name. " + "Example: --skip-layers 77 78 79" + ), + ) args = parser.parse_args() - main(args.input_path, args.output_int8_path, args.num_workers) + main( + args.input_path, + args.output_int8_path, + args.num_workers, + skip_layers=args.skip_layers, + ) print("done")