已开启
feat: 支持量化结果持久化与恢复第一里程碑 #267
sunshine750创建于 17 天前
feat: 支持量化结果持久化与恢复第一里程碑 #267
已开启
共 14 个文件变更+4925-16
| @@ -28,17 +28,22 @@ class InsertQuantizeModulePass(BaseModuleFusionPass): | |||
| 28 | APIs: match_pattern, do_pass | 28 | APIs: match_pattern, do_pass |
| 29 | """ | 29 | """ |
| 30 | 30 | ||
| 31 | - def __init__(self, quant_config): | 31 | + def __init__(self, quant_config, quant_result_writer=None): |
| 32 | """ | 32 | """ |
| 33 | Function: init object of insert quantize op quant pass | 33 | Function: init object of insert quantize op quant pass |
| 34 | Parameter: | 34 | Parameter: |
| 35 | quant_config: dict, config | 35 | quant_config: dict, config |
| 36 | + quant_result_writer: optional QuantResultWriter injected into every | ||
| 37 | + inserted quantize module so it can dump its quant params at its | ||
| 38 | + own finalization point. None means quant results are not | ||
| 39 | + persisted. | ||
| 36 | Return: None | 40 | Return: None |
| 37 | """ | 41 | """ |
| 38 | super().__init__() | 42 | super().__init__() |
| 39 | self.config = quant_config | 43 | self.config = quant_config |
| 40 | self.quant_layers = list(quant_config.keys()) | 44 | self.quant_layers = list(quant_config.keys()) |
| 41 | self.quantize_ops = dict() | 45 | self.quantize_ops = dict() |
| 46 | + self.quant_result_writer = quant_result_writer | ||
| 42 | 47 | ||
| 43 | def match_pattern(self, module, name): | 48 | def match_pattern(self, module, name): |
| 44 | """ | 49 | """ |
| @@ -80,6 +85,8 @@ class InsertQuantizeModulePass(BaseModuleFusionPass): | |||
| 80 | 85 | ||
| 81 | helper = ModuleHelper(model) | 86 | helper = ModuleHelper(model) |
| 82 | helper.replace_module_by_name(model, object_name, new_module) | 87 | helper.replace_module_by_name(model, object_name, new_module) |
| 88 | + if self.quant_result_writer is not None: | ||
| 89 | + new_module.attach_quant_result_writer(self.quant_result_writer) | ||
| 83 | LOGGER.logd( | 90 | LOGGER.logd( |
| 84 | "Insert quantize op module to '{}' success!".format(object_name), | 91 | "Insert quantize op module to '{}' success!".format(object_name), |
| 85 | 'InsertQuantizeModulePass', | 92 | 'InsertQuantizeModulePass', |
| @@ -0,0 +1,1401 @@ | |||
| 1 | +# -*- coding: UTF-8 -*- | ||
| 2 | +# ---------------------------------------------------------------------------- | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 4 | +# | ||
| 5 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 6 | +# you may not use this file except in compliance with the License. | ||
| 7 | +# You may obtain a copy of the License at | ||
| 8 | +# | ||
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 10 | + | ||
| 11 | +# Unless required by applicable law or agreed to in writing, software | ||
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | +# ---------------------------------------------------------------------------- | ||
| 17 | +"""Quantization result artifact: writer, integrity digests and manifest. | ||
| 18 | + | ||
| 19 | +Implements the fixed quant result wire format. | ||
| 20 | +Directory protocol produced by a successful run:: | ||
| 21 | + | ||
| 22 | + quant_result_dir/ | ||
| 23 | + |-- quant_manifest.json | ||
| 24 | + |-- model_layers_0_self_attn_q_proj.safetensors | ||
| 25 | + `-- ... | ||
| 26 | + | ||
| 27 | +``quant_manifest.json`` is written last and only after every expected unit has | ||
| 28 | +been written and re-verified on disk. A directory without the manifest is not a | ||
| 29 | +valid artifact: it is never partially loaded, never resumed and never | ||
| 30 | +overwritten. The manifest is the single source of truth for the full config; the | ||
| 31 | +safetensors header carries only the minimal cross-check fields, so two copies of | ||
| 32 | +the config can never drift. | ||
| 33 | + | ||
| 34 | +Field names, the ``sha256:`` digest prefix and the header key set here are part | ||
| 35 | +of the cross-stage contract consumed by ``reload_quant_params()``; do not rename | ||
| 36 | +them for style. | ||
| 37 | +""" | ||
| 38 | + | ||
| 39 | +import hashlib | ||
| 40 | +import json | ||
| 41 | +import os | ||
| 42 | +import re | ||
| 43 | +import stat | ||
| 44 | +import uuid | ||
| 45 | + | ||
| 46 | +import copy | ||
| 47 | + | ||
| 48 | +import torch | ||
| 49 | +import torch.nn as nn | ||
| 50 | +from safetensors import safe_open | ||
| 51 | +from safetensors.torch import save_file | ||
| 52 | + | ||
| 53 | +from amct_pytorch.common.utils.log import LOGGER | ||
| 54 | + | ||
| 55 | +MANIFEST_NAME = 'quant_manifest.json' | ||
| 56 | +# Frozen manifest schema constants. Keep every wire-format constant together; | ||
| 57 | +# readers and writers must derive their behavior from this single block. | ||
| 58 | +MANIFEST_FORMAT = 'amct-quant-result-manifest' | ||
| 59 | +SCHEMA_VERSION = 1 | ||
| 60 | +UNIT_SUFFIX = '.safetensors' | ||
| 61 | +LOG_TAG = 'QuantResultWriter' | ||
| 62 | +RELOAD_LOG_TAG = 'reload_quant_params' | ||
| 63 | +DIGEST_PREFIX = 'sha256:' | ||
| 64 | + | ||
| 65 | +_REQUIRED_MANIFEST_FIELDS = ( | ||
| 66 | + 'format', | ||
| 67 | + 'schema_version', | ||
| 68 | + 'artifact_id', | ||
| 69 | + 'source_model', | ||
| 70 | + 'units', | ||
| 71 | +) | ||
| 72 | +_REQUIRED_UNIT_FIELDS = ( | ||
| 73 | + 'layer_name', | ||
| 74 | + 'file', | ||
| 75 | + 'quant_module_type', | ||
| 76 | + 'ori_module_type', | ||
| 77 | + 'quant_config', | ||
| 78 | + 'non_tensor_params', | ||
| 79 | + 'tensors', | ||
| 80 | + 'file_size', | ||
| 81 | + 'sha256', | ||
| 82 | + 'entry_digest', | ||
| 83 | +) | ||
| 84 | +_REQUIRED_SOURCE_MODEL_FIELDS = ('model_type', 'torch_dtype') | ||
| 85 | +_REQUIRED_TENSOR_FIELDS = ('key', 'dtype', 'shape') | ||
| 86 | +_DIGEST_PATTERN = re.compile(r'sha256:[0-9a-f]{64}') | ||
| 87 | + | ||
| 88 | +# 640 for files, matching amct_log conventions. | ||
| 89 | +_FILE_MODE = stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | ||
| 90 | + | ||
| 91 | +_UNSAFE_NAME_CHARS = re.compile(r'[^0-9a-zA-Z_]') | ||
| 92 | +_HASH_CHUNK_SIZE = 1024 * 1024 | ||
| 93 | +# Fields derived from the file itself, therefore excluded from entry_digest. | ||
| 94 | +_DERIVED_ENTRY_FIELDS = ('file_size', 'sha256', 'entry_digest') | ||
| 95 | +_ADDRESS_PATTERN = re.compile(r'0x[0-9a-fA-F]{6,}') | ||
| 96 | + | ||
| 97 | + | ||
| 98 | +def is_quant_result_supported(quant_cls): | ||
| 99 | + """Return the class-level artifact capability declaration. | ||
| 100 | + | ||
| 101 | + The declaration is only an eligibility signal. ``write_unit`` and the | ||
| 102 | + reload validators still verify the actual tensor and metadata payload, so a | ||
| 103 | + class cannot claim a valid artifact merely by setting this flag. | ||
| 104 | + """ | ||
| 105 | + if not isinstance(quant_cls, type): | ||
| 106 | + return False | ||
| 107 | + if not bool(getattr(quant_cls, 'supports_quant_result', False)): | ||
| 108 | + return False | ||
| 109 | + quant_param_keys = getattr(quant_cls, 'quant_param_keys', None) | ||
| 110 | + if not callable(quant_param_keys): | ||
| 111 | + return False | ||
| 112 | + try: | ||
| 113 | + return bool(quant_param_keys()) | ||
| 114 | + except (TypeError, ValueError): | ||
| 115 | + return False | ||
| 116 | + | ||
| 117 | + | ||
| 118 | +def canonical_json(data): | ||
| 119 | + """Serialize to the one canonical JSON form used by every digest.""" | ||
| 120 | + return json.dumps( | ||
| 121 | + data, | ||
| 122 | + ensure_ascii=False, | ||
| 123 | + sort_keys=True, | ||
| 124 | + separators=(",", ":"), | ||
| 125 | + allow_nan=False, | ||
| 126 | + ) | ||
| 127 | + | ||
| 128 | + | ||
| 129 | +def _prefixed(hexdigest): | ||
| 130 | + """Wire form of a digest: ``sha256:<hex>``.""" | ||
| 131 | + return DIGEST_PREFIX + hexdigest | ||
| 132 | + | ||
| 133 | + | ||
| 134 | +def sha256_of_text(text): | ||
| 135 | + """Prefixed SHA-256 of a str, encoded as UTF-8.""" | ||
| 136 | + return _prefixed(hashlib.sha256(text.encode('utf-8')).hexdigest()) | ||
| 137 | + | ||
| 138 | + | ||
| 139 | +def digest_of(data): | ||
| 140 | + """Prefixed SHA-256 over the canonical JSON form of ``data``.""" | ||
| 141 | + return sha256_of_text(canonical_json(data)) | ||
| 142 | + | ||
| 143 | + | ||
| 144 | +def sha256_of_file(file_path): | ||
| 145 | + """Prefixed SHA-256 of a file, read in bounded chunks.""" | ||
| 146 | + hasher = hashlib.sha256() | ||
| 147 | + with open(file_path, 'rb') as file_handle: | ||
| 148 | + for chunk in iter(lambda: file_handle.read(_HASH_CHUNK_SIZE), b''): | ||
| 149 | + hasher.update(chunk) | ||
| 150 | + return _prefixed(hasher.hexdigest()) | ||
| 151 | + | ||
| 152 | + | ||
| 153 | +def _byte_view(tensor): | ||
| 154 | + """Flat contiguous uint8 view of a CPU tensor, dtype-agnostic. | ||
| 155 | + | ||
| 156 | + ``Tensor.numpy()`` refuses bfloat16/float8 dtypes, so reinterpret the raw | ||
| 157 | + storage as uint8 first. Flattening before the view keeps non-contiguous and | ||
| 158 | + 0-d inputs working. | ||
| 159 | + """ | ||
| 160 | + return tensor.detach().reshape(-1).contiguous().view(torch.uint8) | ||
| 161 | + | ||
| 162 | + | ||
| 163 | +def update_hash_with_tensor(hasher, tensor, chunk_elems=_HASH_CHUNK_SIZE): | ||
| 164 | + """Fold a tensor's raw bytes into ``hasher`` without materializing a copy.""" | ||
| 165 | + byte_view = _byte_view(tensor.cpu()) | ||
| 166 | + total = byte_view.numel() | ||
| 167 | + for start in range(0, total, chunk_elems): | ||
| 168 | + hasher.update(byte_view[start : start + chunk_elems].numpy().tobytes()) | ||
| 169 | + | ||
| 170 | + | ||
| 171 | +def dtype_name(dtype): | ||
| 172 | + """Artifact dtype spelling: ``float32``, not ``torch.float32``.""" | ||
| 173 | + text = str(dtype) | ||
| 174 | + return text[len('torch.') :] if text.startswith('torch.') else text | ||
| 175 | + | ||
| 176 | + | ||
| 177 | +def tensor_schema_list(tensors): | ||
| 178 | + """Tensor schema: a key-sorted list of ``{key, dtype, shape}``.""" | ||
| 179 | + return [ | ||
| 180 | + { | ||
| 181 | + 'key': key, | ||
| 182 | + 'dtype': dtype_name(tensors[key].dtype), | ||
| 183 | + 'shape': list(tensors[key].shape), | ||
| 184 | + } | ||
| 185 | + for key in sorted(tensors) | ||
| 186 | + ] | ||
| 187 | + | ||
| 188 | + | ||
| 189 | +def to_jsonable(value, path='quant_config'): | ||
| 190 | + """Strictly coerce a value for the manifest, or raise. | ||
| 191 | + | ||
| 192 | + Only the types the artifact contract admits are converted: str/bool/int/float, | ||
| 193 | + None, ``torch.dtype``, and list/tuple/set/dict of those. Anything else -- | ||
| 194 | + a module, a tensor, an arbitrary object -- is rejected rather than | ||
| 195 | + stringified, because a stringified object cannot reconstruct a module on | ||
| 196 | + reload and would silently produce an unusable artifact. | ||
| 197 | + """ | ||
| 198 | + if value is None or isinstance(value, (bool, int, str)): | ||
| 199 | + return value | ||
| 200 | + if isinstance(value, float): | ||
| 201 | + if value != value or value in (float('inf'), float('-inf')): | ||
| 202 | + raise ValueError( | ||
| 203 | + "{} is the non-finite float {!r}, which canonical JSON cannot " | ||
| 204 | + "represent.".format(path, value) | ||
| 205 | + ) | ||
| 206 | + return value | ||
| 207 | + if isinstance(value, torch.dtype): | ||
| 208 | + return dtype_name(value) | ||
| 209 | + if isinstance(value, dict): | ||
| 210 | + return _strict_mapping(value, path) | ||
| 211 | + if isinstance(value, (set, frozenset)): | ||
| 212 | + items = [to_jsonable(item, '{}[]'.format(path)) for item in value] | ||
| 213 | + return sorted(items, key=canonical_json) | ||
| 214 | + if isinstance(value, (list, tuple)): | ||
| 215 | + return [ | ||
| 216 | + to_jsonable(item, '{}[{}]'.format(path, index)) | ||
| 217 | + for index, item in enumerate(value) | ||
| 218 | + ] | ||
| 219 | + raise TypeError( | ||
| 220 | + "{} holds {} ({!r}), which cannot be serialized into a quant result " | ||
| 221 | + "manifest. Only scalars, torch.dtype and containers of those are " | ||
| 222 | + "allowed, so that reload can reconstruct the module.".format( | ||
| 223 | + path, type(value).__name__, value | ||
| 224 | + ) | ||
| 225 | + ) | ||
| 226 | + | ||
| 227 | + | ||
| 228 | +def _strict_mapping(mapping, path): | ||
| 229 | + """Strictly convert a dict; keys must be str, bool or int.""" | ||
| 230 | + result = {} | ||
| 231 | + for key, item in mapping.items(): | ||
| 232 | + if not isinstance(key, (str, bool, int)): | ||
| 233 | + raise TypeError( | ||
| 234 | + "{} has a {} key ({!r}); quant result config keys must be " | ||
| 235 | + "str, bool or int.".format(path, type(key).__name__, key) | ||
| 236 | + ) | ||
| 237 | + json_key = key if isinstance(key, str) else str(key) | ||
| 238 | + if json_key in result: | ||
| 239 | + raise ValueError( | ||
| 240 | + "{} key {!r} collides with another key after string " | ||
| 241 | + "normalization.".format(path, json_key) | ||
| 242 | + ) | ||
| 243 | + result[json_key] = to_jsonable(item, '{}.{}'.format(path, json_key)) | ||
| 244 | + return result | ||
| 245 | + | ||
| 246 | + | ||
| 247 | +def to_jsonable_lossy(value): | ||
| 248 | + """Coerce anything into a stable digestible form, stringifying unknowns. | ||
| 249 | + | ||
| 250 | + Only for the *fingerprint* of a source model config: a HuggingFace | ||
| 251 | + ``config.to_dict()`` may legitimately carry arbitrary objects, and the | ||
| 252 | + fingerprint only needs to be deterministic and change-detecting, never | ||
| 253 | + reversible. Never use this for ``quant_config`` or ``non_tensor_params`` -- | ||
| 254 | + those must survive a reload, so they go through the strict | ||
| 255 | + :func:`to_jsonable`. | ||
| 256 | + """ | ||
| 257 | + if value is None or isinstance(value, (bool, int, str)): | ||
| 258 | + return value | ||
| 259 | + if isinstance(value, float): | ||
| 260 | + if value != value or value in (float('inf'), float('-inf')): | ||
| 261 | + return repr(value) | ||
| 262 | + return value | ||
| 263 | + if isinstance(value, torch.dtype): | ||
| 264 | + return dtype_name(value) | ||
| 265 | + if isinstance(value, torch.Tensor): | ||
| 266 | + return { | ||
| 267 | + 'dtype': dtype_name(value.dtype), | ||
| 268 | + 'shape': list(value.shape), | ||
| 269 | + 'sha256': tensor_bytes_sha256(value), | ||
| 270 | + } | ||
| 271 | + if isinstance(value, (bytes, bytearray)): | ||
| 272 | + return bytes(value).hex() | ||
| 273 | + if isinstance(value, dict): | ||
| 274 | + return { | ||
| 275 | + (key if isinstance(key, str) else _stringify_unknown(key)): ( | ||
| 276 | + to_jsonable_lossy(item) | ||
| 277 | + ) | ||
| 278 | + for key, item in value.items() | ||
| 279 | + } | ||
| 280 | + if isinstance(value, (set, frozenset)): | ||
| 281 | + return sorted((to_jsonable_lossy(item) for item in value), key=canonical_json) | ||
| 282 | + if isinstance(value, (list, tuple)): | ||
| 283 | + return [to_jsonable_lossy(item) for item in value] | ||
| 284 | + return _stringify_unknown(value) | ||
| 285 | + | ||
| 286 | + | ||
| 287 | +def _stringify_unknown(value): | ||
| 288 | + """Stable string form for a type JSON cannot represent. | ||
| 289 | + | ||
| 290 | + ``str()`` of many objects embeds the memory address, which would make the | ||
| 291 | + fingerprint differ between runs; fall back to the qualified type name then. | ||
| 292 | + """ | ||
| 293 | + text = str(value) | ||
| 294 | + if _ADDRESS_PATTERN.search(text): | ||
| 295 | + value_type = type(value) | ||
| 296 | + return '{}.{}'.format(value_type.__module__, value_type.__qualname__) | ||
| 297 | + return text | ||
| 298 | + | ||
| 299 | + | ||
| 300 | +def tensor_bytes_sha256(tensor): | ||
| 301 | + """Prefixed SHA-256 over the raw bytes of a tensor.""" | ||
| 302 | + hasher = hashlib.sha256() | ||
| 303 | + update_hash_with_tensor(hasher, tensor) | ||
| 304 | + return _prefixed(hasher.hexdigest()) | ||
| 305 | + | ||
| 306 | + | ||
| 307 | +def unit_file_name(layer_name): | ||
| 308 | + """Map a layer name to its unit file name. | ||
| 309 | + | ||
| 310 | + Every character outside ``[0-9a-zA-Z_]`` is replaced, so the result can | ||
| 311 | + never contain a path separator, ``..`` or any other traversal construct. | ||
| 312 | + Distinct layer names can normalize to the same file name (``a.b`` and | ||
| 313 | + ``a-b`` both give ``a_b``); the writer keeps a reverse map and rejects such | ||
| 314 | + collisions instead of silently overwriting. | ||
| 315 | + """ | ||
| 316 | + if not isinstance(layer_name, str) or not layer_name: | ||
| 317 | + raise ValueError( | ||
| 318 | + "layer name for a quant result unit must be a non-empty str, " | ||
| 319 | + "got {!r}.".format(layer_name) | ||
| 320 | + ) | ||
| 321 | + sanitized = _UNSAFE_NAME_CHARS.sub('_', layer_name) | ||
| 322 | + if sanitized.strip('_') == '': | ||
| 323 | + raise ValueError( | ||
| 324 | + "layer name '{}' normalizes to an empty file name.".format(layer_name) | ||
| 325 | + ) | ||
| 326 | + return sanitized + UNIT_SUFFIX | ||
| 327 | + | ||
| 328 | + | ||
| 329 | +def resolve_unit_path(result_dir, file_name): | ||
| 330 | + """Absolute path of a unit file, rejecting anything outside result_dir. | ||
| 331 | + | ||
| 332 | + Also enforced on read, since a manifest is untrusted input: an absolute | ||
| 333 | + path, a ``..`` segment or a nested path must never escape the artifact. | ||
| 334 | + """ | ||
| 335 | + if not isinstance(file_name, str) or not file_name: | ||
| 336 | + raise ValueError("quant result unit file name must be a non-empty str.") | ||
| 337 | + if os.path.isabs(file_name) or os.path.dirname(file_name): | ||
| 338 | + raise ValueError( | ||
| 339 | + "quant result unit file '{}' must be a bare relative file name; " | ||
| 340 | + "absolute paths and directory components are rejected.".format(file_name) | ||
| 341 | + ) | ||
| 342 | + root = os.path.realpath(result_dir) | ||
| 343 | + candidate = os.path.realpath(os.path.join(root, file_name)) | ||
| 344 | + if os.path.dirname(candidate) != root: | ||
| 345 | + raise ValueError( | ||
| 346 | + "quant result unit '{}' resolves outside the result directory '{}'.".format( | ||
| 347 | + file_name, result_dir | ||
| 348 | + ) | ||
| 349 | + ) | ||
| 350 | + return candidate | ||
| 351 | + | ||
| 352 | + | ||
| 353 | +def validate_quant_result_dir(quant_result_dir): | ||
| 354 | + """Validate a caller-created directory before writing quant results. | ||
| 355 | + | ||
| 356 | + Only ``None`` means "do not persist". Everything else must be a usable, | ||
| 357 | + existing directory. Files, missing paths and inaccessible directories are | ||
| 358 | + rejected. Other directory entries may coexist with the artifact; the writer | ||
| 359 | + checks its own manifest and unit paths separately and never overwrites them. | ||
| 360 | + """ | ||
| 361 | + if not isinstance(quant_result_dir, str): | ||
| 362 | + raise TypeError( | ||
| 363 | + "quant_result_dir must be a str path or None, got {}.".format( | ||
| 364 | + type(quant_result_dir).__name__ | ||
| 365 | + ) | ||
| 366 | + ) | ||
| 367 | + if quant_result_dir.strip() == '': | ||
| 368 | + raise ValueError( | ||
| 369 | + "quant_result_dir must be a non-empty path; pass None to disable " | ||
| 370 | + "quant result persistence." | ||
| 371 | + ) | ||
| 372 | + | ||
| 373 | + target = os.path.abspath(os.path.expanduser(quant_result_dir)) | ||
| 374 | + if not os.path.exists(target): | ||
| 375 | + raise ValueError( | ||
| 376 | + "quant_result_dir '{}' must be an existing directory; create it " | ||
| 377 | + "before calling quantize().".format(target) | ||
| 378 | + ) | ||
| 379 | + if not os.path.isdir(target): | ||
| 380 | + raise ValueError( | ||
| 381 | + "quant_result_dir '{}' must be an existing directory, not a file.".format( | ||
| 382 | + target | ||
| 383 | + ) | ||
| 384 | + ) | ||
| 385 | + if not os.access(target, os.R_OK | os.W_OK | os.X_OK): | ||
| 386 | + raise ValueError( | ||
| 387 | + "quant_result_dir '{}' must be readable, writable, and searchable.".format( | ||
| 388 | + target | ||
| 389 | + ) | ||
| 390 | + ) | ||
| 391 | + return target | ||
| 392 | + | ||
| 393 | + | ||
| 394 | +def temp_name_for(file_path): | ||
| 395 | + """Same-directory temp name for an atomic write of ``file_path``.""" | ||
| 396 | + return os.path.join( | ||
| 397 | + os.path.dirname(file_path), | ||
| 398 | + '.{}.{}.tmp'.format(os.path.basename(file_path), os.getpid()), | ||
| 399 | + ) | ||
| 400 | + | ||
| 401 | + | ||
| 402 | +def _atomic_write_via(file_path, write_payload): | ||
| 403 | + """Atomically produce file_path: temp file -> fsync -> os.replace. | ||
| 404 | + | ||
| 405 | + ``write_payload(tmp_path)`` does the actual writing, so a large tensor | ||
| 406 | + payload can stream straight to disk instead of being serialized into a bytes | ||
| 407 | + object first. Temp file and target share a directory, hence a filesystem. | ||
| 408 | + """ | ||
| 409 | + tmp_path = temp_name_for(file_path) | ||
| 410 | + try: | ||
| 411 | + write_payload(tmp_path) | ||
| 412 | + _fsync_file(tmp_path) | ||
| 413 | + os.chmod(tmp_path, _FILE_MODE) | ||
| 414 | + os.replace(tmp_path, file_path) | ||
| 415 | + except BaseException: | ||
| 416 | + # Leave no partial file behind, for any failure including interrupts. | ||
| 417 | + if os.path.exists(tmp_path): | ||
| 418 | + os.remove(tmp_path) | ||
| 419 | + raise | ||
| 420 | + _fsync_dir(os.path.dirname(file_path)) | ||
| 421 | + | ||
| 422 | + | ||
| 423 | +def _fsync_file(file_path): | ||
| 424 | + """Flush a just-written file's contents to stable storage.""" | ||
| 425 | + file_descriptor = os.open(file_path, os.O_RDONLY) | ||
| 426 | + try: | ||
| 427 | + os.fsync(file_descriptor) | ||
| 428 | + finally: | ||
| 429 | + os.close(file_descriptor) | ||
| 430 | + | ||
| 431 | + | ||
| 432 | +def _fsync_dir(directory): | ||
| 433 | + """fsync a directory so a rename is durable; best effort by platform.""" | ||
| 434 | + try: | ||
| 435 | + dir_fd = os.open(directory, os.O_RDONLY) | ||
| 436 | + except OSError: | ||
| 437 | + return | ||
| 438 | + try: | ||
| 439 | + os.fsync(dir_fd) | ||
| 440 | + except OSError: | ||
| 441 | + pass | ||
| 442 | + finally: | ||
| 443 | + os.close(dir_fd) | ||
| 444 | + | ||
| 445 | + | ||
| 446 | +def _write_text_atomically(file_path, text): | ||
| 447 | + """Atomically write UTF-8 text (used for the manifest).""" | ||
| 448 | + | ||
| 449 | + def _write(tmp_path): | ||
| 450 | + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | ||
| 451 | + file_descriptor = os.open(tmp_path, flags, _FILE_MODE) | ||
| 452 | + with os.fdopen(file_descriptor, 'w', encoding='utf-8') as handle: | ||
| 453 | + handle.write(text) | ||
| 454 | + handle.flush() | ||
| 455 | + os.fsync(handle.fileno()) | ||
| 456 | + | ||
| 457 | + _atomic_write_via(file_path, _write) | ||
| 458 | + | ||
| 459 | + | ||
| 460 | +def _source_model_type(model): | ||
| 461 | + """Artifact ``source_model.model_type``: HF model_type, else class name.""" | ||
| 462 | + config = getattr(model, 'config', None) | ||
| 463 | + model_type = getattr(config, 'model_type', None) | ||
| 464 | + if isinstance(model_type, str) and model_type: | ||
| 465 | + return model_type | ||
| 466 | + return type(model).__name__ | ||
| 467 | + | ||
| 468 | + | ||
| 469 | +def _source_model_dtype(model): | ||
| 470 | + """Return the dtype that quantization treats as the source model dtype. | ||
| 471 | + | ||
| 472 | + HuggingFace models declare it on ``config.torch_dtype``. Plain PyTorch | ||
| 473 | + modules do not, so use their first parameter or buffer as the equivalent | ||
| 474 | + declaration. The artifact intentionally records this semantic input, | ||
| 475 | + rather than hashing every weight byte or unrelated configuration field. | ||
| 476 | + """ | ||
| 477 | + config = getattr(model, 'config', None) | ||
| 478 | + configured = getattr(config, 'torch_dtype', None) | ||
| 479 | + if configured is None: | ||
| 480 | + to_dict = getattr(config, 'to_dict', None) | ||
| 481 | + if callable(to_dict): | ||
| 482 | + config_dict = to_dict() | ||
| 483 | + if isinstance(config_dict, dict): | ||
| 484 | + configured = config_dict.get('torch_dtype') | ||
| 485 | + if isinstance(configured, torch.dtype): | ||
| 486 | + return dtype_name(configured) | ||
| 487 | + if isinstance(configured, str) and configured: | ||
| 488 | + return dtype_name(configured) | ||
| 489 | + for tensor in model.parameters(): | ||
| 490 | + return dtype_name(tensor.dtype) | ||
| 491 | + for tensor in model.buffers(): | ||
| 492 | + return dtype_name(tensor.dtype) | ||
| 493 | + return 'unknown' | ||
| 494 | + | ||
| 495 | + | ||
| 496 | +def compute_model_fingerprint(model): | ||
| 497 | + """Build the source-model identity before any quant module is inserted. | ||
| 498 | + | ||
| 499 | + Only model type and source dtype are compared during reload. Quantization | ||
| 500 | + results are intentionally reusable across checkpoints with the same | ||
| 501 | + quantization identity; layer names, module types and tensor shapes remain | ||
| 502 | + independently validated while staging each replacement module. | ||
| 503 | + """ | ||
| 504 | + return { | ||
Y | |||
| 505 | + 'model_type': _source_model_type(model), | ||
| 506 | + 'torch_dtype': _source_model_dtype(model), | ||
| 507 | + } | ||
| 508 | + | ||
| 509 | + | ||
| 510 | +def unit_header(artifact_id, layer_name, entry_digest): | ||
| 511 | + """The minimal safetensors header: three ``str -> str`` cross-check fields. | ||
| 512 | + | ||
| 513 | + The manifest is the only place that declares the artifact format and schema | ||
| 514 | + version. A unit only binds itself to that manifest and its corresponding | ||
| 515 | + entry, so global schema metadata is not duplicated in every file. | ||
| 516 | + """ | ||
| 517 | + return { | ||
| 518 | + 'artifact_id': artifact_id, | ||
| 519 | + 'layer_name': layer_name, | ||
| 520 | + 'entry_digest': entry_digest, | ||
| 521 | + } | ||
| 522 | + | ||
| 523 | + | ||
| 524 | +class QuantResultWriter: | ||
| 525 | + """Collects per-layer quant params and publishes a quant result artifact. | ||
| 526 | + | ||
| 527 | + Lifecycle: one ``write_unit()`` per expected layer at that layer's own | ||
| 528 | + finalization point, then ``publish_if_complete()``. The manifest appears | ||
| 529 | + only when every expected unit has been written and re-verified, so an | ||
| 530 | + interrupted run leaves an unpublished (invalid) directory rather than a | ||
| 531 | + half-usable one. | ||
| 532 | + """ | ||
| 533 | + | ||
| 534 | + def __init__( | ||
| 535 | + self, | ||
| 536 | + result_dir, | ||
| 537 | + expected_units, | ||
| 538 | + model_fingerprint, | ||
| 539 | + artifact_id=None, | ||
| 540 | + ): | ||
| 541 | + self._result_dir = os.path.abspath(result_dir) | ||
| 542 | + self._expected_units = self._check_expected_units(expected_units) | ||
| 543 | + _validate_source_model_fields(model_fingerprint) | ||
| 544 | + self._model_fingerprint = dict(model_fingerprint) | ||
| 545 | + self._artifact_id = artifact_id or str(uuid.uuid4()) | ||
| 546 | + self._entries = {} | ||
| 547 | + self._file_names = {} | ||
| 548 | + self._published = False | ||
| 549 | + | ||
| 550 | + | ||
| 551 | + def _check_expected_units(expected_units): | ||
| 552 | + """Validate the expected unit set: non-empty, unique, all str.""" | ||
| 553 | + units = list(expected_units) | ||
| 554 | + if not units: | ||
| 555 | + raise ValueError( | ||
| 556 | + "no quantizable layer was found for this model and config, so " | ||
| 557 | + "there is no quant result to persist." | ||
| 558 | + ) | ||
| 559 | + for unit in units: | ||
| 560 | + if not isinstance(unit, str) or not unit: | ||
| 561 | + raise ValueError( | ||
| 562 | + "expected quant result unit names must be non-empty str, " | ||
| 563 | + "got {!r}.".format(unit) | ||
| 564 | + ) | ||
| 565 | + if len(set(units)) != len(units): | ||
| 566 | + raise ValueError("expected quant result unit names contain duplicates.") | ||
| 567 | + return tuple(units) | ||
| 568 | + | ||
| 569 | + | ||
| 570 | + def create(cls, quant_result_dir, expected_units, model_fingerprint): | ||
| 571 | + """Validate inputs and return a writer for a caller-created directory.""" | ||
| 572 | + units = cls._check_expected_units(expected_units) | ||
| 573 | + result_dir = validate_quant_result_dir(quant_result_dir) | ||
| 574 | + cls._check_managed_paths_available(result_dir, units) | ||
| 575 | + return cls(result_dir, units, model_fingerprint) | ||
| 576 | + | ||
| 577 | + | ||
| 578 | + def _check_managed_paths_available(result_dir, expected_units): | ||
| 579 | + """Reject only paths this writer owns; leave unrelated entries alone.""" | ||
| 580 | + managed_names = [MANIFEST_NAME] | ||
| 581 | + managed_names.extend(unit_file_name(unit) for unit in expected_units) | ||
| 582 | + existing = sorted( | ||
| 583 | + set( | ||
| 584 | + name | ||
| 585 | + for name in managed_names | ||
| 586 | + if os.path.lexists(os.path.join(result_dir, name)) | ||
| 587 | + ) | ||
| 588 | + ) | ||
| 589 | + if existing: | ||
| 590 | + raise ValueError( | ||
| 591 | + "quant_result_dir '{}' already contains AMCT quant result file(s) " | ||
| 592 | + "{}; refusing to overwrite or resume them.".format(result_dir, existing) | ||
| 593 | + ) | ||
| 594 | + | ||
| 595 | + | ||
| 596 | + def result_dir(self): | ||
| 597 | + return self._result_dir | ||
| 598 | + | ||
| 599 | + | ||
| 600 | + def manifest_path(self): | ||
| 601 | + return os.path.join(self._result_dir, MANIFEST_NAME) | ||
| 602 | + | ||
| 603 | + | ||
| 604 | + def pending_units(self): | ||
| 605 | + """Expected units that still have no entry, in declaration order.""" | ||
| 606 | + return tuple(unit for unit in self._expected_units if unit not in self._entries) | ||
| 607 | + | ||
| 608 | + def write_unit( | ||
| 609 | + self, | ||
| 610 | + layer_name, | ||
| 611 | + tensors, | ||
| 612 | + non_tensor_params=None, | ||
| 613 | + quant_config=None, | ||
| 614 | + quant_module_type=None, | ||
| 615 | + ori_module_type=None, | ||
| 616 | + ): | ||
| 617 | + """Persist one layer's quant params as an atomic safetensors unit. | ||
| 618 | + | ||
| 619 | + Registers the unit in memory only; the manifest is published separately | ||
| 620 | + once every expected unit has landed. | ||
| 621 | + """ | ||
| 622 | + self._check_writable(layer_name) | ||
| 623 | + file_name = self._reserve_file_name(layer_name) | ||
| 624 | + tensors = self._check_tensors(layer_name, tensors) | ||
| 625 | + # Required by the schema the reader enforces: writing None here would | ||
| 626 | + # publish an artifact that cannot be reloaded. | ||
| 627 | + _non_empty_str( | ||
| 628 | + quant_module_type, | ||
| 629 | + "quant_module_type of layer '{}'".format(layer_name), | ||
| 630 | + ) | ||
| 631 | + _non_empty_str( | ||
| 632 | + ori_module_type, "ori_module_type of layer '{}'".format(layer_name) | ||
| 633 | + ) | ||
| 634 | + quant_config = _mapping_or_empty( | ||
| 635 | + quant_config, "quant_config of layer '{}'".format(layer_name) | ||
| 636 | + ) | ||
| 637 | + non_tensor_params = _mapping_or_empty( | ||
| 638 | + non_tensor_params, "non_tensor_params of layer '{}'".format(layer_name) | ||
| 639 | + ) | ||
| 640 | + entry = { | ||
| 641 | + 'layer_name': layer_name, | ||
| 642 | + 'file': file_name, | ||
| 643 | + 'quant_module_type': quant_module_type, | ||
| 644 | + 'ori_module_type': ori_module_type, | ||
| 645 | + 'quant_config': to_jsonable( | ||
| 646 | + quant_config, "quant_config of layer '{}'".format(layer_name) | ||
| 647 | + ), | ||
| 648 | + 'non_tensor_params': to_jsonable( | ||
| 649 | + non_tensor_params, | ||
| 650 | + "non_tensor_params of layer '{}'".format(layer_name), | ||
| 651 | + ), | ||
| 652 | + 'tensors': tensor_schema_list(tensors), | ||
| 653 | + } | ||
| 654 | + self._check_jsonable(layer_name, entry) | ||
| 655 | + entry['entry_digest'] = digest_of(entry) | ||
| 656 | + | ||
| 657 | + file_path = resolve_unit_path(self._result_dir, file_name) | ||
| 658 | + if os.path.lexists(file_path): | ||
| 659 | + raise FileExistsError( | ||
| 660 | + "quant result unit '{}' already exists; refusing to overwrite it.".format( | ||
| 661 | + file_path | ||
| 662 | + ) | ||
| 663 | + ) | ||
| 664 | + header = unit_header(self._artifact_id, layer_name, entry['entry_digest']) | ||
| 665 | + # save_file streams straight to the temp path: a large tensor payload is | ||
| 666 | + # never serialized into an in-memory bytes object first. | ||
| 667 | + _atomic_write_via( | ||
| 668 | + file_path, | ||
| 669 | + lambda tmp_path: save_file(tensors, tmp_path, metadata=header), | ||
| 670 | + ) | ||
| 671 | + # Read the unit back before trusting it, then bind size and digest. | ||
| 672 | + self._verify_unit_file(file_path, entry, header) | ||
| 673 | + entry['file_size'] = os.path.getsize(file_path) | ||
| 674 | + entry['sha256'] = sha256_of_file(file_path) | ||
| 675 | + | ||
| 676 | + self._entries[layer_name] = entry | ||
| 677 | + LOGGER.logd( | ||
| 678 | + "Wrote quant result unit for layer '{}' to '{}'".format( | ||
| 679 | + layer_name, file_name | ||
| 680 | + ), | ||
| 681 | + LOG_TAG, | ||
| 682 | + ) | ||
| 683 | + return entry | ||
| 684 | + | ||
| 685 | + def _check_writable(self, layer_name): | ||
| 686 | + """Reject writes after publish, for unknown layers, or duplicates.""" | ||
| 687 | + if self._published: | ||
| 688 | + raise RuntimeError( | ||
| 689 | + "quant result '{}' is already published; it is immutable.".format( | ||
| 690 | + self._result_dir | ||
| 691 | + ) | ||
| 692 | + ) | ||
| 693 | + if layer_name not in self._expected_units: | ||
| 694 | + raise ValueError( | ||
| 695 | + "layer '{}' is not part of the expected quant result units; " | ||
| 696 | + "refusing to write an unexpected unit.".format(layer_name) | ||
| 697 | + ) | ||
| 698 | + if layer_name in self._entries: | ||
| 699 | + raise ValueError( | ||
| 700 | + "quant result unit for layer '{}' was already written; each " | ||
| 701 | + "layer may be dumped only once.".format(layer_name) | ||
| 702 | + ) | ||
| 703 | + | ||
| 704 | + def _reserve_file_name(self, layer_name): | ||
| 705 | + """Reserve the unit file name, rejecting post-sanitization collisions.""" | ||
| 706 | + file_name = unit_file_name(layer_name) | ||
| 707 | + owner = self._file_names.get(file_name) | ||
| 708 | + if owner is not None and owner != layer_name: | ||
| 709 | + raise ValueError( | ||
| 710 | + "layers '{}' and '{}' both map to quant result file '{}'; " | ||
| 711 | + "refusing to overwrite.".format(owner, layer_name, file_name) | ||
| 712 | + ) | ||
| 713 | + self._file_names[file_name] = layer_name | ||
| 714 | + return file_name | ||
| 715 | + | ||
| 716 | + | ||
| 717 | + def _check_tensors(layer_name, tensors): | ||
| 718 | + """Validate the tensor payload: str keys, CPU, contiguous, detached.""" | ||
| 719 | + if not isinstance(tensors, dict): | ||
| 720 | + raise TypeError( | ||
| 721 | + "quant param tensors of layer '{}' must be a dict, got {}.".format( | ||
| 722 | + layer_name, type(tensors).__name__ | ||
| 723 | + ) | ||
| 724 | + ) | ||
| 725 | + checked = {} | ||
| 726 | + for key, value in tensors.items(): | ||
| 727 | + if not isinstance(key, str) or not key: | ||
| 728 | + raise ValueError( | ||
| 729 | + "quant param tensor keys of layer '{}' must be non-empty " | ||
| 730 | + "str, got {!r}.".format(layer_name, key) | ||
| 731 | + ) | ||
| 732 | + if not isinstance(value, torch.Tensor): | ||
| 733 | + raise TypeError( | ||
| 734 | + "quant param '{}' of layer '{}' must be a torch.Tensor, " | ||
| 735 | + "got {}.".format(key, layer_name, type(value).__name__) | ||
| 736 | + ) | ||
| 737 | + checked[key] = value.detach().cpu().contiguous() | ||
| 738 | + return checked | ||
| 739 | + | ||
| 740 | + | ||
| 741 | + def _check_jsonable(layer_name, entry): | ||
| 742 | + """Fail early if the entry cannot be canonically serialized.""" | ||
| 743 | + try: | ||
| 744 | + canonical_json(entry) | ||
| 745 | + except (TypeError, ValueError) as error: | ||
| 746 | + raise ValueError( | ||
| 747 | + "quant result entry of layer '{}' is not JSON serializable: {}".format( | ||
| 748 | + layer_name, error | ||
| 749 | + ) | ||
| 750 | + ) from error | ||
| 751 | + | ||
| 752 | + | ||
| 753 | + def _verify_unit_file(file_path, entry, header): | ||
| 754 | + """Re-read a just-written unit and check its header and tensor schema.""" | ||
| 755 | + with safe_open(file_path, framework='pt') as file_handle: | ||
| 756 | + actual_header = file_handle.metadata() or {} | ||
| 757 | + actual = {key: file_handle.get_tensor(key) for key in file_handle.keys()} | ||
| 758 | + if actual_header != header: | ||
| 759 | + raise RuntimeError( | ||
| 760 | + "quant result unit '{}' header mismatch after write: expected " | ||
| 761 | + "{}, found {}.".format(file_path, header, actual_header) | ||
| 762 | + ) | ||
| 763 | + actual_schema = tensor_schema_list(actual) | ||
| 764 | + if actual_schema != entry['tensors']: | ||
| 765 | + raise RuntimeError( | ||
| 766 | + "quant result unit '{}' tensor schema mismatch after write: " | ||
| 767 | + "expected {}, found {}.".format( | ||
| 768 | + file_path, entry['tensors'], actual_schema | ||
| 769 | + ) | ||
| 770 | + ) | ||
| 771 | + | ||
| 772 | + def publish_if_complete(self): | ||
| 773 | + """Publish the manifest iff every expected unit has been written. | ||
| 774 | + | ||
| 775 | + Returns True when the artifact is published (including a previous | ||
| 776 | + publish), False when units are still pending. Called at each candidate | ||
| 777 | + finalization point, so the manifest lands as soon as -- and only when -- | ||
| 778 | + the artifact is complete. | ||
| 779 | + """ | ||
| 780 | + if self._published: | ||
| 781 | + return True | ||
| 782 | + pending = self.pending_units | ||
| 783 | + if pending: | ||
| 784 | + LOGGER.logd( | ||
| 785 | + "Quant result not published yet, {} unit(s) pending, e.g. '{}'".format( | ||
| 786 | + len(pending), pending[0] | ||
| 787 | + ), | ||
| 788 | + LOG_TAG, | ||
| 789 | + ) | ||
| 790 | + return False | ||
| 791 | + self._publish() | ||
| 792 | + return True | ||
| 793 | + | ||
| 794 | + def _publish(self): | ||
| 795 | + """Verify the whole file set, then atomically write the manifest last.""" | ||
| 796 | + if os.path.lexists(self.manifest_path): | ||
| 797 | + raise FileExistsError( | ||
| 798 | + "quant result manifest '{}' already exists; refusing to overwrite " | ||
| 799 | + "or resume it.".format(self.manifest_path) | ||
| 800 | + ) | ||
| 801 | + self._verify_result_dir() | ||
| 802 | + manifest = { | ||
| 803 | + 'format': MANIFEST_FORMAT, | ||
| 804 | + 'schema_version': SCHEMA_VERSION, | ||
| 805 | + 'artifact_id': self._artifact_id, | ||
| 806 | + 'source_model': self._model_fingerprint, | ||
| 807 | + 'units': [self._entries[unit] for unit in self._expected_units], | ||
| 808 | + } | ||
| 809 | + _write_text_atomically(self.manifest_path, canonical_json(manifest) + '\n') | ||
| 810 | + self._published = True | ||
| 811 | + LOGGER.logi( | ||
| 812 | + "Published quant result manifest '{}' with {} unit(s).".format( | ||
| 813 | + self.manifest_path, len(manifest['units']) | ||
| 814 | + ), | ||
| 815 | + LOG_TAG, | ||
| 816 | + ) | ||
| 817 | + | ||
| 818 | + def _verify_result_dir(self): | ||
| 819 | + """Check every unit managed by this writer is present and intact.""" | ||
| 820 | + expected_files = {entry['file'] for entry in self._entries.values()} | ||
| 821 | + actual_files = set(os.listdir(self._result_dir)) | ||
| 822 | + missing = sorted(expected_files - actual_files) | ||
| 823 | + if missing: | ||
| 824 | + raise RuntimeError( | ||
| 825 | + "quant result directory '{}' is missing unit file(s) {}; " | ||
| 826 | + "refusing to publish an incomplete artifact.".format( | ||
| 827 | + self._result_dir, missing | ||
| 828 | + ) | ||
| 829 | + ) | ||
| 830 | + for entry in self._entries.values(): | ||
| 831 | + self._verify_entry(entry) | ||
| 832 | + | ||
| 833 | + def _verify_entry(self, entry): | ||
| 834 | + """Re-check one entry against the file on disk before publishing.""" | ||
| 835 | + file_path = resolve_unit_path(self._result_dir, entry['file']) | ||
| 836 | + actual_size = os.path.getsize(file_path) | ||
| 837 | + if actual_size != entry['file_size']: | ||
| 838 | + raise RuntimeError( | ||
| 839 | + "quant result unit '{}' changed size after write: expected {} " | ||
| 840 | + "bytes, found {}.".format(file_path, entry['file_size'], actual_size) | ||
| 841 | + ) | ||
| 842 | + if sha256_of_file(file_path) != entry['sha256']: | ||
| 843 | + raise RuntimeError( | ||
| 844 | + "quant result unit '{}' failed its SHA-256 check before " | ||
| 845 | + "publish.".format(file_path) | ||
| 846 | + ) | ||
| 847 | + digest_input = { | ||
| 848 | + key: value | ||
| 849 | + for key, value in entry.items() | ||
| 850 | + if key not in _DERIVED_ENTRY_FIELDS | ||
| 851 | + } | ||
| 852 | + if digest_of(digest_input) != entry['entry_digest']: | ||
| 853 | + raise RuntimeError( | ||
| 854 | + "quant result entry of layer '{}' failed its entry_digest " | ||
| 855 | + "check before publish.".format(entry['layer_name']) | ||
| 856 | + ) | ||
| 857 | + self._verify_unit_file( | ||
| 858 | + file_path, | ||
| 859 | + entry, | ||
| 860 | + unit_header(self._artifact_id, entry['layer_name'], entry['entry_digest']), | ||
| 861 | + ) | ||
| 862 | + | ||
| 863 | + | ||
| 864 | +def load_quant_manifest(quant_result_dir): | ||
| 865 | + """Load and validate ``quant_manifest.json`` of a published artifact. | ||
| 866 | + | ||
| 867 | + A directory without the manifest is not a valid quant result, so its absence | ||
| 868 | + is an error rather than an empty read. Parsed with ``json`` only -- never | ||
| 869 | + pickle -- and every exact schema field, duplicate-key rule, unit uniqueness | ||
| 870 | + rule and safe relative-path rule are enforced before anything is handed back. | ||
| 871 | + """ | ||
| 872 | + result_dir = os.path.abspath(os.path.expanduser(quant_result_dir)) | ||
| 873 | + manifest_path = os.path.join(result_dir, MANIFEST_NAME) | ||
| 874 | + if not os.path.isfile(manifest_path): | ||
| 875 | + raise FileNotFoundError( | ||
| 876 | + "'{}' is not a valid quant result: {} is missing, so the run that " | ||
| 877 | + "produced it did not complete.".format(quant_result_dir, MANIFEST_NAME) | ||
| 878 | + ) | ||
| 879 | + with open(manifest_path, encoding='utf-8') as file_handle: | ||
| 880 | + manifest = json.load(file_handle, object_pairs_hook=_reject_duplicate_keys) | ||
| 881 | + if not isinstance(manifest, dict): | ||
| 882 | + raise ValueError( | ||
| 883 | + "quant manifest '{}' is not a JSON object.".format(manifest_path) | ||
| 884 | + ) | ||
| 885 | + | ||
| 886 | + _check_required(manifest, _REQUIRED_MANIFEST_FIELDS, 'quant manifest') | ||
| 887 | + if manifest['format'] != MANIFEST_FORMAT: | ||
| 888 | + raise ValueError( | ||
| 889 | + "unsupported quant manifest format '{}', expected '{}'.".format( | ||
| 890 | + manifest['format'], MANIFEST_FORMAT | ||
| 891 | + ) | ||
| 892 | + ) | ||
| 893 | + # bool is an int subclass and 1.0 == 1, so an exact int is demanded rather | ||
| 894 | + # than mere equality: the wire format says this field is an integer. | ||
| 895 | + version = manifest['schema_version'] | ||
| 896 | + if ( | ||
| 897 | + isinstance(version, bool) | ||
| 898 | + or not isinstance(version, int) | ||
| 899 | + or version != SCHEMA_VERSION | ||
| 900 | + ): | ||
| 901 | + raise ValueError( | ||
| 902 | + 'unsupported quant manifest schema_version {!r}, expected the int ' | ||
| 903 | + '{}.'.format(version, SCHEMA_VERSION) | ||
| 904 | + ) | ||
| 905 | + _validate_artifact_id(manifest['artifact_id']) | ||
| 906 | + _validate_source_model_fields(manifest['source_model']) | ||
| 907 | + _validate_units(manifest['units'], result_dir) | ||
| 908 | + return manifest | ||
| 909 | + | ||
| 910 | + | ||
| 911 | +def _non_empty_str(value, what): | ||
| 912 | + """Every identifier in the manifest must be a non-empty string.""" | ||
| 913 | + if not isinstance(value, str) or not value: | ||
| 914 | + raise ValueError('{} must be a non-empty string, got {!r}.'.format(what, value)) | ||
| 915 | + return value | ||
| 916 | + | ||
| 917 | + | ||
| 918 | +def _mapping_or_empty(value, what): | ||
| 919 | + """Normalize an optional mapping field, rejecting anything else. | ||
| 920 | + | ||
| 921 | + ``None`` means "nothing to record" and becomes ``{}``. Deliberately not | ||
| 922 | + ``value or {}``: that also swallows ``[]``, ``''`` and ``False`` into an empty | ||
| 923 | + config, and a non-mapping would otherwise be written into a manifest the | ||
| 924 | + reader is obliged to reject. | ||
| 925 | + """ | ||
| 926 | + if value is None: | ||
| 927 | + return {} | ||
| 928 | + if not isinstance(value, dict): | ||
| 929 | + raise TypeError( | ||
| 930 | + '{} must be a mapping or None, got {}.'.format(what, type(value).__name__) | ||
| 931 | + ) | ||
| 932 | + return value | ||
| 933 | + | ||
| 934 | + | ||
| 935 | +def _validate_artifact_id(artifact_id): | ||
| 936 | + """``artifact_id`` is the UUID the unit headers are cross-checked against.""" | ||
| 937 | + _non_empty_str(artifact_id, 'quant manifest artifact_id') | ||
| 938 | + try: | ||
| 939 | + uuid.UUID(artifact_id) | ||
| 940 | + except (ValueError, AttributeError, TypeError): | ||
| 941 | + raise ValueError( | ||
| 942 | + 'quant manifest artifact_id {!r} is not a UUID.'.format(artifact_id) | ||
| 943 | + ) from None | ||
| 944 | + | ||
| 945 | + | ||
| 946 | +def _validate_digest(value, what): | ||
| 947 | + """Digests are always ``sha256:`` followed by 64 lowercase hex chars.""" | ||
| 948 | + _non_empty_str(value, what) | ||
| 949 | + if not _DIGEST_PATTERN.fullmatch(value): | ||
| 950 | + raise ValueError( | ||
| 951 | + '{} must look like "{}<64 hex chars>", got {!r}.'.format( | ||
| 952 | + what, DIGEST_PREFIX, value | ||
| 953 | + ) | ||
| 954 | + ) | ||
| 955 | + | ||
| 956 | + | ||
| 957 | +def _validate_source_model_fields(source_model): | ||
| 958 | + """``source_model`` is the fingerprint reload matches the input model on.""" | ||
| 959 | + _check_required(source_model, _REQUIRED_SOURCE_MODEL_FIELDS, 'source_model') | ||
| 960 | + _non_empty_str(source_model['model_type'], 'source_model.model_type') | ||
| 961 | + _non_empty_str(source_model['torch_dtype'], 'source_model.torch_dtype') | ||
| 962 | + | ||
| 963 | + | ||
| 964 | +def _reject_duplicate_keys(pairs): | ||
| 965 | + """Reject duplicate JSON object keys instead of silently keeping the last.""" | ||
| 966 | + result = {} | ||
| 967 | + for key, value in pairs: | ||
| 968 | + if key in result: | ||
| 969 | + raise ValueError( | ||
| 970 | + "quant manifest contains duplicate JSON key {!r}.".format(key) | ||
| 971 | + ) | ||
| 972 | + result[key] = value | ||
| 973 | + return result | ||
| 974 | + | ||
| 975 | + | ||
| 976 | +def _check_required(mapping, fields, what): | ||
| 977 | + """Require the exact frozen field set for a protocol object.""" | ||
| 978 | + if not isinstance(mapping, dict): | ||
| 979 | + raise ValueError("{} must be a JSON object.".format(what)) | ||
| 980 | + missing = [field for field in fields if field not in mapping] | ||
| 981 | + if missing: | ||
| 982 | + raise ValueError("{} is missing required field(s) {}.".format(what, missing)) | ||
| 983 | + unexpected = sorted(set(mapping) - set(fields)) | ||
| 984 | + if unexpected: | ||
| 985 | + raise ValueError("{} contains unknown field(s) {}.".format(what, unexpected)) | ||
| 986 | + | ||
| 987 | + | ||
| 988 | +def _validate_unit_fields(entry): | ||
| 989 | + """Type-check one unit entry against the frozen artifact schema. | ||
| 990 | + | ||
| 991 | + Every field reload depends on is checked here rather than where it is used, so | ||
| 992 | + a malformed manifest cannot get as far as touching the model. | ||
| 993 | + """ | ||
| 994 | + layer_name = entry['layer_name'] | ||
| 995 | + for field in ('layer_name', 'file', 'quant_module_type', 'ori_module_type'): | ||
| 996 | + _non_empty_str(entry[field], "quant manifest unit {}".format(field)) | ||
| 997 | + for field in ('quant_config', 'non_tensor_params'): | ||
| 998 | + if not isinstance(entry[field], dict): | ||
| 999 | + raise ValueError( | ||
| 1000 | + "quant manifest unit '{}' field {} must be a JSON object, got " | ||
| 1001 | + '{}.'.format(layer_name, field, type(entry[field]).__name__) | ||
| 1002 | + ) | ||
| 1003 | + # bool is an int subclass, so it has to be excluded explicitly. | ||
| 1004 | + file_size = entry['file_size'] | ||
| 1005 | + if isinstance(file_size, bool) or not isinstance(file_size, int) or file_size < 0: | ||
| 1006 | + raise ValueError( | ||
| 1007 | + "quant manifest unit '{}' file_size must be a non-negative int, got " | ||
| 1008 | + '{!r}.'.format(layer_name, file_size) | ||
| 1009 | + ) | ||
| 1010 | + _validate_digest( | ||
| 1011 | + entry['sha256'], "quant manifest unit '{}' sha256".format(layer_name) | ||
| 1012 | + ) | ||
| 1013 | + _validate_digest( | ||
| 1014 | + entry['entry_digest'], | ||
| 1015 | + "quant manifest unit '{}' entry_digest".format(layer_name), | ||
| 1016 | + ) | ||
| 1017 | + _validate_tensor_schema(entry['tensors'], layer_name) | ||
| 1018 | + | ||
| 1019 | + | ||
| 1020 | +def _validate_tensor_schema(tensors, layer_name): | ||
| 1021 | + """The tensor schema is a list of ``{key, dtype, shape}``, keys unique.""" | ||
| 1022 | + if not isinstance(tensors, list): | ||
| 1023 | + raise ValueError( | ||
| 1024 | + "quant manifest unit '{}' tensors must be a list, got {}.".format( | ||
| 1025 | + layer_name, type(tensors).__name__ | ||
| 1026 | + ) | ||
| 1027 | + ) | ||
| 1028 | + seen = set() | ||
| 1029 | + for item in tensors: | ||
| 1030 | + _check_required( | ||
| 1031 | + item, | ||
| 1032 | + _REQUIRED_TENSOR_FIELDS, | ||
| 1033 | + "quant manifest unit '{}' tensor".format(layer_name), | ||
| 1034 | + ) | ||
| 1035 | + key = _non_empty_str( | ||
| 1036 | + item['key'], "quant manifest unit '{}' tensor key".format(layer_name) | ||
| 1037 | + ) | ||
| 1038 | + _non_empty_str( | ||
| 1039 | + item['dtype'], | ||
| 1040 | + "quant manifest unit '{}' tensor '{}' dtype".format(layer_name, key), | ||
| 1041 | + ) | ||
| 1042 | + shape = item['shape'] | ||
| 1043 | + if not isinstance(shape, list) or any( | ||
| 1044 | + isinstance(dim, bool) or not isinstance(dim, int) or dim < 0 | ||
| 1045 | + for dim in shape | ||
| 1046 | + ): | ||
| 1047 | + raise ValueError( | ||
| 1048 | + "quant manifest unit '{}' tensor '{}' shape must be a list of " | ||
| 1049 | + 'non-negative ints, got {!r}.'.format(layer_name, key, shape) | ||
| 1050 | + ) | ||
| 1051 | + if key in seen: | ||
| 1052 | + raise ValueError( | ||
| 1053 | + "quant manifest unit '{}' lists tensor '{}' more than once.".format( | ||
| 1054 | + layer_name, key | ||
| 1055 | + ) | ||
| 1056 | + ) | ||
| 1057 | + seen.add(key) | ||
| 1058 | + | ||
| 1059 | + | ||
| 1060 | +def _validate_units(units, result_dir): | ||
| 1061 | + """Validate the unit list: shape, required fields, uniqueness, safe paths.""" | ||
| 1062 | + if not isinstance(units, list) or not units: | ||
| 1063 | + raise ValueError("quant manifest units must be a non-empty list.") | ||
| 1064 | + seen_layers = set() | ||
| 1065 | + seen_files = set() | ||
| 1066 | + for entry in units: | ||
| 1067 | + _check_required(entry, _REQUIRED_UNIT_FIELDS, 'quant manifest unit') | ||
| 1068 | + _validate_unit_fields(entry) | ||
| 1069 | + layer_name = entry['layer_name'] | ||
| 1070 | + if layer_name in seen_layers: | ||
| 1071 | + raise ValueError( | ||
| 1072 | + "quant manifest lists layer '{}' more than once.".format(layer_name) | ||
| 1073 | + ) | ||
| 1074 | + seen_layers.add(layer_name) | ||
| 1075 | + # Raises on absolute paths, directory components and traversal. | ||
| 1076 | + resolve_unit_path(result_dir, entry['file']) | ||
| 1077 | + if entry['file'] in seen_files: | ||
| 1078 | + raise ValueError( | ||
| 1079 | + "quant manifest lists file '{}' more than once.".format(entry['file']) | ||
| 1080 | + ) | ||
| 1081 | + seen_files.add(entry['file']) | ||
| 1082 | + | ||
| 1083 | + | ||
| 1084 | +def validate_result_directory(quant_result_dir, manifest): | ||
| 1085 | + """Check that the manifest and every unit it manages are present.""" | ||
| 1086 | + result_dir = os.path.abspath(os.path.expanduser(quant_result_dir)) | ||
| 1087 | + expected = {entry['file'] for entry in manifest['units']} | {MANIFEST_NAME} | ||
| 1088 | + actual = set(os.listdir(result_dir)) | ||
| 1089 | + missing = sorted(expected - actual) | ||
| 1090 | + if missing: | ||
| 1091 | + raise ValueError( | ||
| 1092 | + "quant result '{}' is missing file(s) listed in its manifest: {}.".format( | ||
| 1093 | + result_dir, missing | ||
| 1094 | + ) | ||
| 1095 | + ) | ||
| 1096 | + | ||
| 1097 | + | ||
| 1098 | +def read_quant_unit(quant_result_dir, layer_name, manifest=None): | ||
| 1099 | + """Read one verified unit back as ``(tensors, non_tensor_params, entry)``. | ||
| 1100 | + | ||
| 1101 | + Checks file size, SHA-256, header, entry digest and tensor schema before | ||
| 1102 | + handing anything back, so a tampered or truncated artifact fails loudly. | ||
| 1103 | + """ | ||
| 1104 | + result_dir = os.path.abspath(os.path.expanduser(quant_result_dir)) | ||
| 1105 | + manifest = manifest or load_quant_manifest(result_dir) | ||
| 1106 | + entries = {entry['layer_name']: entry for entry in manifest['units']} | ||
| 1107 | + entry = entries.get(layer_name) | ||
| 1108 | + if entry is None: | ||
| 1109 | + raise KeyError( | ||
| 1110 | + "quant result '{}' has no unit for layer '{}'.".format( | ||
| 1111 | + result_dir, layer_name | ||
| 1112 | + ) | ||
| 1113 | + ) | ||
| 1114 | + | ||
| 1115 | + file_path = resolve_unit_path(result_dir, entry['file']) | ||
| 1116 | + if not os.path.isfile(file_path): | ||
| 1117 | + raise ValueError("quant result unit '{}' is missing.".format(file_path)) | ||
| 1118 | + if os.path.getsize(file_path) != entry['file_size']: | ||
| 1119 | + raise ValueError("quant result unit '{}' has wrong size.".format(file_path)) | ||
| 1120 | + if sha256_of_file(file_path) != entry['sha256']: | ||
| 1121 | + raise ValueError( | ||
| 1122 | + "quant result unit '{}' failed its SHA-256 check.".format(file_path) | ||
| 1123 | + ) | ||
| 1124 | + digest_input = { | ||
| 1125 | + key: value for key, value in entry.items() if key not in _DERIVED_ENTRY_FIELDS | ||
| 1126 | + } | ||
| 1127 | + if digest_of(digest_input) != entry['entry_digest']: | ||
| 1128 | + raise ValueError( | ||
| 1129 | + "quant result entry of layer '{}' failed its entry_digest check.".format( | ||
| 1130 | + layer_name | ||
| 1131 | + ) | ||
| 1132 | + ) | ||
| 1133 | + | ||
| 1134 | + with safe_open(file_path, framework='pt') as file_handle: | ||
| 1135 | + header = file_handle.metadata() or {} | ||
| 1136 | + tensors = {key: file_handle.get_tensor(key) for key in file_handle.keys()} | ||
| 1137 | + expected_header = unit_header( | ||
| 1138 | + manifest['artifact_id'], layer_name, entry['entry_digest'] | ||
| 1139 | + ) | ||
| 1140 | + if header != expected_header: | ||
| 1141 | + raise ValueError( | ||
| 1142 | + "quant result unit '{}' header does not cross-check with the " | ||
| 1143 | + "manifest.".format(file_path) | ||
| 1144 | + ) | ||
| 1145 | + actual_schema = tensor_schema_list(tensors) | ||
| 1146 | + if actual_schema != entry['tensors']: | ||
| 1147 | + raise ValueError( | ||
| 1148 | + "quant result unit '{}' tensor schema does not match the manifest.".format( | ||
| 1149 | + file_path | ||
| 1150 | + ) | ||
| 1151 | + ) | ||
| 1152 | + return tensors, dict(entry['non_tensor_params']), entry | ||
| 1153 | + | ||
| 1154 | + | ||
| 1155 | +def validate_source_model(model, source_model): | ||
| 1156 | + """The float model handed to reload must be the one that was calibrated. | ||
| 1157 | + | ||
| 1158 | + Compared on model type, structure/config digest and full state_dict digest, | ||
| 1159 | + so a different architecture, a changed dtype or shape, or any altered weight | ||
| 1160 | + is refused before a single module is touched. | ||
| 1161 | + """ | ||
| 1162 | + # Shape first: a non-mapping must fail as a validation error, not as an | ||
| 1163 | + # AttributeError from inside the strict conversion. | ||
| 1164 | + _check_required(source_model, _REQUIRED_SOURCE_MODEL_FIELDS, 'source_model') | ||
| 1165 | + _strict_mapping(source_model, 'source_model') | ||
| 1166 | + actual = compute_model_fingerprint(model) | ||
| 1167 | + mismatched = [ | ||
| 1168 | + (field, source_model[field], actual[field]) | ||
| 1169 | + for field in _REQUIRED_SOURCE_MODEL_FIELDS | ||
| 1170 | + if source_model[field] != actual.get(field) | ||
| 1171 | + ] | ||
| 1172 | + if mismatched: | ||
| 1173 | + details = '; '.join( | ||
| 1174 | + '{}: artifact has {!r} but this model has {!r}'.format(*item) | ||
| 1175 | + for item in mismatched | ||
| 1176 | + ) | ||
| 1177 | + raise ValueError( | ||
| 1178 | + 'the model passed in is not the model this quant result was produced ' | ||
| 1179 | + 'from ({}). Reload expects the same original float model that was ' | ||
| 1180 | + 'calibrated.'.format(details) | ||
| 1181 | + ) | ||
| 1182 | + | ||
| 1183 | + | ||
| 1184 | +def resolve_layer_parent(model, layer_name): | ||
| 1185 | + """Resolve ``layer_name`` to ``(parent_module, attr_name, child_module)``. | ||
| 1186 | + | ||
| 1187 | + Every failure is a hard error: an empty name, an unreachable path, or a path | ||
| 1188 | + that runs through something that is not a module all mean the artifact does | ||
| 1189 | + not describe this model. | ||
| 1190 | + """ | ||
| 1191 | + if not isinstance(layer_name, str) or not layer_name: | ||
| 1192 | + raise ValueError('quant result unit has an empty layer_name.') | ||
| 1193 | + parts = layer_name.split('.') | ||
| 1194 | + if any(not part for part in parts): | ||
| 1195 | + raise ValueError( | ||
| 1196 | + "quant result unit layer_name '{}' has an empty path component.".format( | ||
| 1197 | + layer_name | ||
| 1198 | + ) | ||
| 1199 | + ) | ||
| 1200 | + parent = model | ||
| 1201 | + for depth, part in enumerate(parts[:-1]): | ||
| 1202 | + child = getattr(parent, part, None) | ||
| 1203 | + if not isinstance(child, nn.Module): | ||
| 1204 | + raise ValueError( | ||
| 1205 | + "cannot reach layer '{}': '{}' is not a submodule of the model " | ||
| 1206 | + 'passed in.'.format(layer_name, '.'.join(parts[: depth + 1])) | ||
| 1207 | + ) | ||
| 1208 | + parent = child | ||
| 1209 | + attr_name = parts[-1] | ||
| 1210 | + child = getattr(parent, attr_name, None) | ||
| 1211 | + if not isinstance(child, nn.Module): | ||
| 1212 | + raise ValueError( | ||
| 1213 | + "the model passed in has no submodule '{}', which this quant result " | ||
| 1214 | + 'expects.'.format(layer_name) | ||
| 1215 | + ) | ||
| 1216 | + return parent, attr_name, child | ||
| 1217 | + | ||
| 1218 | + | ||
| 1219 | +def _check_no_nested_units(layer_names): | ||
| 1220 | + """Refuse units whose paths nest, e.g. ``blk`` and ``blk.fc``. | ||
| 1221 | + | ||
| 1222 | + Replacing the outer module would throw away the inner replacement (or make | ||
| 1223 | + the result depend on ordering), so this is rejected up front. | ||
| 1224 | + """ | ||
| 1225 | + ordered = sorted(layer_names) | ||
| 1226 | + for outer in ordered: | ||
| 1227 | + prefix = outer + '.' | ||
| 1228 | + nested = [name for name in ordered if name.startswith(prefix)] | ||
| 1229 | + if nested: | ||
| 1230 | + raise ValueError( | ||
| 1231 | + "quant result cannot be reloaded: layer '{}' contains layer(s) {}, " | ||
| 1232 | + 'so replacing it would discard them.'.format(outer, nested) | ||
| 1233 | + ) | ||
| 1234 | + | ||
| 1235 | + | ||
| 1236 | +def _resolve_quant_module_class(entry): | ||
| 1237 | + """Resolve a unit through the registered algorithm and source module type.""" | ||
| 1238 | + # Imported lazily to keep this module free of an import-time dependency on | ||
| 1239 | + # the algorithm registry. | ||
| 1240 | + from amct_pytorch.algorithms import AlgorithmRegistry | ||
| 1241 | + from amct_pytorch.common.config.utils import get_alg_name_from_config | ||
| 1242 | + from amct_pytorch.quantize_op.base_quant_module import BaseQuantizeModule | ||
| 1243 | + | ||
| 1244 | + layer_name = entry['layer_name'] | ||
| 1245 | + quant_config = entry['quant_config'] | ||
| 1246 | + _strict_mapping(quant_config, "quant_config of layer '{}'".format(layer_name)) | ||
| 1247 | + try: | ||
| 1248 | + algorithms, _ = get_alg_name_from_config(quant_config.get('algorithm')) | ||
| 1249 | + except ValueError as error: | ||
| 1250 | + raise ValueError( | ||
| 1251 | + "quant config of layer '{}' must name exactly one algorithm.".format( | ||
| 1252 | + layer_name | ||
| 1253 | + ) | ||
| 1254 | + ) from error | ||
| 1255 | + if len(algorithms) != 1: | ||
| 1256 | + raise ValueError( | ||
| 1257 | + "quant config of layer '{}' must name exactly one algorithm, got {}.".format( | ||
| 1258 | + layer_name, algorithms | ||
| 1259 | + ) | ||
| 1260 | + ) | ||
| 1261 | + | ||
| 1262 | + algorithm = algorithms[0] | ||
| 1263 | + registered_sources = AlgorithmRegistry.algo.get(algorithm) | ||
| 1264 | + if registered_sources is None: | ||
| 1265 | + raise ValueError( | ||
| 1266 | + "quant config of layer '{}' names algorithm '{}', which is not " | ||
| 1267 | + 'registered in this installation.'.format(layer_name, algorithm) | ||
| 1268 | + ) | ||
| 1269 | + ori_module_type = entry['ori_module_type'] | ||
| 1270 | + quant_cls = registered_sources.get(ori_module_type) | ||
| 1271 | + if quant_cls is None: | ||
| 1272 | + raise ValueError( | ||
| 1273 | + "algorithm '{}' has no registered quantize module for source type " | ||
| 1274 | + "'{}', required by layer '{}'.".format( | ||
| 1275 | + algorithm, ori_module_type, layer_name | ||
| 1276 | + ) | ||
| 1277 | + ) | ||
| 1278 | + if not is_quant_result_supported(quant_cls): | ||
| 1279 | + raise ValueError( | ||
| 1280 | + "algorithm '{}' of layer '{}' does not support quant result reload " | ||
| 1281 | + 'because its registered quantize class does not declare the artifact ' | ||
| 1282 | + 'contract.'.format(algorithm, layer_name) | ||
| 1283 | + ) | ||
| 1284 | + if not isinstance(quant_cls, type): | ||
| 1285 | + raise ValueError( | ||
| 1286 | + "registered quantize module for algorithm '{}' and source type '{}' " | ||
| 1287 | + 'is not a class.'.format(algorithm, ori_module_type) | ||
| 1288 | + ) | ||
| 1289 | + if not issubclass(quant_cls, BaseQuantizeModule): | ||
| 1290 | + raise ValueError( | ||
| 1291 | + "quantize module '{}' of layer '{}' is not a BaseQuantizeModule " | ||
| 1292 | + 'subclass.'.format(quant_cls.__name__, layer_name) | ||
| 1293 | + ) | ||
| 1294 | + if entry['quant_module_type'] != quant_cls.__name__: | ||
| 1295 | + raise ValueError( | ||
| 1296 | + "quant_module_type '{}' of layer '{}' does not match the registered " | ||
| 1297 | + "class '{}' selected by algorithm '{}' and source type '{}'.".format( | ||
| 1298 | + entry['quant_module_type'], | ||
| 1299 | + layer_name, | ||
| 1300 | + quant_cls.__name__, | ||
| 1301 | + algorithm, | ||
| 1302 | + ori_module_type, | ||
| 1303 | + ) | ||
| 1304 | + ) | ||
| 1305 | + return quant_cls | ||
| 1306 | + | ||
| 1307 | + | ||
| 1308 | +def _check_ori_module_type(entry, ori_module): | ||
| 1309 | + """The layer in this model must still be the type that was calibrated.""" | ||
| 1310 | + actual = type(ori_module).__name__ | ||
| 1311 | + if actual != entry['ori_module_type']: | ||
| 1312 | + raise ValueError( | ||
| 1313 | + "the model passed in is not the model expected by quant result layer " | ||
| 1314 | + "'{}': it is a {}, but the artifact was produced from a {}.".format( | ||
| 1315 | + entry['layer_name'], actual, entry['ori_module_type'] | ||
| 1316 | + ) | ||
| 1317 | + ) | ||
| 1318 | + | ||
| 1319 | + | ||
| 1320 | +def _module_device(module): | ||
| 1321 | + """Device of a module's own tensors, defaulting to CPU when it has none.""" | ||
| 1322 | + for tensor in list(module.parameters(recurse=True)) + list(module.buffers()): | ||
| 1323 | + return tensor.device | ||
| 1324 | + return torch.device('cpu') | ||
| 1325 | + | ||
| 1326 | + | ||
| 1327 | +def _build_staged_module(entry, ori_module, tensors, non_tensor_params): | ||
| 1328 | + """Build one replacement module without touching the live model. | ||
| 1329 | + | ||
| 1330 | + The quantize module is constructed on an isolated ``deepcopy`` of the layer on | ||
| 1331 | + purpose: some modules keep a reference to the original ``Parameter`` (GPTQ | ||
| 1332 | + does), and loading a persisted weight into that would mutate the model before | ||
| 1333 | + the transaction commits. | ||
| 1334 | + """ | ||
| 1335 | + staged_source = copy.deepcopy(ori_module) | ||
| 1336 | + quant_cls = _resolve_quant_module_class(entry) | ||
| 1337 | + replacement = quant_cls(staged_source, entry['layer_name'], entry['quant_config']) | ||
| 1338 | + replacement.load_quant_params(tensors, non_tensor_params) | ||
| 1339 | + replacement.to(_module_device(ori_module)) | ||
| 1340 | + return replacement | ||
| 1341 | + | ||
| 1342 | + | ||
| 1343 | +def reload_quant_params(model, quant_result_dir): | ||
| 1344 | + """Rebuild a calibrated model from a published quant result, atomically. | ||
| 1345 | + | ||
| 1346 | + Takes the same original float model the artifact was produced from and puts it | ||
| 1347 | + back into its post-calibration state. Everything is validated and every | ||
| 1348 | + replacement module is built and loaded before the first ``setattr``, so a bad | ||
| 1349 | + artifact leaves the model exactly as it was -- no partially quantized model. | ||
| 1350 | + | ||
| 1351 | + Args: | ||
| 1352 | + model: the original float ``nn.Module`` that was calibrated. | ||
| 1353 | + quant_result_dir: directory holding the published quant result. | ||
| 1354 | + Returns: | ||
| 1355 | + The same model instance, modified in place. | ||
| 1356 | + """ | ||
| 1357 | + result_dir = os.path.abspath(os.path.expanduser(quant_result_dir)) | ||
| 1358 | + # Phase 1: validate the artifact and stage every replacement. | ||
| 1359 | + manifest = load_quant_manifest(result_dir) | ||
| 1360 | + validate_result_directory(result_dir, manifest) | ||
| 1361 | + validate_source_model(model, manifest['source_model']) | ||
| 1362 | + | ||
| 1363 | + entries = list(manifest['units']) | ||
| 1364 | + _check_no_nested_units([entry['layer_name'] for entry in entries]) | ||
| 1365 | + | ||
| 1366 | + replacement_plan = [] | ||
| 1367 | + for entry in entries: | ||
| 1368 | + layer_name = entry['layer_name'] | ||
| 1369 | + parent, attr_name, ori_module = resolve_layer_parent(model, layer_name) | ||
| 1370 | + _check_ori_module_type(entry, ori_module) | ||
| 1371 | + tensors, non_tensor_params, _ = read_quant_unit( | ||
| 1372 | + result_dir, layer_name, manifest | ||
| 1373 | + ) | ||
| 1374 | + replacement = _build_staged_module( | ||
| 1375 | + entry, ori_module, tensors, non_tensor_params | ||
| 1376 | + ) | ||
| 1377 | + replacement_plan.append((parent, attr_name, replacement, ori_module)) | ||
| 1378 | + | ||
| 1379 | + # Phase 2: commit. Nothing here can fail on artifact content any more; an | ||
| 1380 | + # unexpected setattr failure is rolled back so the model stays untouched. | ||
| 1381 | + _commit_replacements(replacement_plan, result_dir, manifest) | ||
| 1382 | + return model | ||
| 1383 | + | ||
| 1384 | + | ||
| 1385 | +def _commit_replacements(replacement_plan, result_dir, manifest): | ||
| 1386 | + """Install the staged modules, rolling back if any ``setattr`` fails.""" | ||
| 1387 | + committed = [] | ||
| 1388 | + try: | ||
| 1389 | + for parent, attr_name, replacement, ori_module in replacement_plan: | ||
| 1390 | + setattr(parent, attr_name, replacement) | ||
| 1391 | + committed.append((parent, attr_name, ori_module)) | ||
| 1392 | + except Exception: | ||
| 1393 | + for parent, attr_name, ori_module in reversed(committed): | ||
| 1394 | + setattr(parent, attr_name, ori_module) | ||
| 1395 | + raise | ||
| 1396 | + LOGGER.logd( | ||
| 1397 | + "Reloaded {} quant unit(s) of artifact '{}' from '{}'.".format( | ||
| 1398 | + len(replacement_plan), manifest['artifact_id'], result_dir | ||
| 1399 | + ), | ||
| 1400 | + RELOAD_LOG_TAG, | ||
| 1401 | + ) | ||
| @@ -25,35 +25,161 @@ from amct_pytorch.classic.optimizer import ( | |||
| 25 | ReplaceNpuQuantModulePass, | 25 | ReplaceNpuQuantModulePass, |
| 26 | ) | 26 | ) |
| 27 | from amct_pytorch.quantize_op.base_quant_module import BaseQuantizeModule | 27 | from amct_pytorch.quantize_op.base_quant_module import BaseQuantizeModule |
| 28 | +from amct_pytorch.classic.quant_result import ( | ||
| 29 | + QuantResultWriter, | ||
| 30 | + compute_model_fingerprint, | ||
| 31 | + is_quant_result_supported, | ||
| 32 | + reload_quant_params, | ||
| 33 | + validate_quant_result_dir, | ||
| 34 | +) | ||
| 28 | from amct_pytorch.common.config import parse_config, set_default_config | 35 | from amct_pytorch.common.config import parse_config, set_default_config |
| 36 | +from amct_pytorch.common.config.utils import get_alg_name_from_config | ||
| 29 | from amct_pytorch.common.utils.check_params import check_params | 37 | from amct_pytorch.common.utils.check_params import check_params |
| 38 | +from amct_pytorch.common.utils.log import LOGGER | ||
| 30 | 39 | ||
| 31 | 40 | ||
| 32 | -@check_params(model=nn.Module, config=(dict, type(None))) | 41 | +@check_params( |
| 33 | -def quantize(model, config=None): | 42 | + model=nn.Module, config=(dict, type(None)), quant_result_dir=(str, type(None)) |
| 43 | +) | ||
| 44 | +def quantize(model, config=None, quant_result_dir=None): | ||
| 34 | """ | 45 | """ |
| 35 | Function: Modify user's model for quantization | 46 | Function: Modify user's model for quantization |
| 36 | Parameter: model: user mode instance of Torch.nn.Module | 47 | Parameter: model: user mode instance of Torch.nn.Module |
| 37 | config: simply config dict from user to set | 48 | config: simply config dict from user to set |
| 49 | + quant_result_dir: optional path of an existing directory to persist | ||
| 50 | + the quant result artifact into. Unrelated files may coexist, | ||
| 51 | + but AMCT never overwrites its own result files. None (default) | ||
| 52 | + keeps the model in-memory only and writes nothing. | ||
| 38 | Return: None | 53 | Return: None |
| 39 | """ | 54 | """ |
| 40 | if config is None: | 55 | if config is None: |
| 41 | config = set_default_config() | 56 | config = set_default_config() |
| 42 | 57 | ||
| 43 | layer_config = parse_config(model, config, AlgorithmRegistry) | 58 | layer_config = parse_config(model, config, AlgorithmRegistry) |
| 59 | + if '' in layer_config: | ||
| 60 | + # Module replacement is keyed by child name and cannot replace the root. | ||
| 61 | + raise ValueError( | ||
| 62 | + "the model passed to quantize() is itself a quantizable layer ({}); " | ||
| 63 | + "each quantized layer must be a named submodule. Wrap it in a parent " | ||
| 64 | + 'nn.Module.'.format(type(model).__name__) | ||
| 65 | + ) | ||
| 66 | + | ||
| 67 | + writer = None | ||
| 68 | + if quant_result_dir is not None: | ||
| 69 | + writer = _create_quant_result_writer(model, layer_config, quant_result_dir) | ||
| 44 | 70 | ||
| 45 | optimizer = ModelOptimizer() | 71 | optimizer = ModelOptimizer() |
| 46 | - optimizer.add_pass(InsertQuantizeModulePass(layer_config)) | 72 | + optimizer.add_pass( |
| 73 | + InsertQuantizeModulePass(layer_config, quant_result_writer=writer) | ||
| 74 | + ) | ||
| 47 | optimizer.do_optimizer(model) | 75 | optimizer.do_optimizer(model) |
| 48 | 76 | ||
| 77 | + if writer is not None: | ||
| 78 | + _finalize_quant_result(model, writer) | ||
| 49 | 79 | ||
| 50 | -@check_params(model=nn.Module) | 80 | + |
| 51 | -def convert(model): | 81 | +def _check_quant_result_algorithms(model, layer_config): |
| 82 | + """Screen selected registered quant classes before writing any files. | ||
| 83 | + | ||
| 84 | + Artifact support is declared by the class in ``AlgorithmRegistry.algo``; | ||
| 85 | + there is no second algorithm-name allowlist to drift out of sync. The | ||
| 86 | + writer still validates concrete values when each layer reaches finalization. | ||
| 87 | + """ | ||
| 88 | + unsupported = {} | ||
| 89 | + modules = dict(model.named_modules()) | ||
| 90 | + for layer_name, cfg in layer_config.items(): | ||
| 91 | + algorithm = cfg.get('algorithm') if isinstance(cfg, dict) else None | ||
| 92 | + alg_names, _ = get_alg_name_from_config(algorithm) | ||
| 93 | + for alg_name in alg_names: | ||
| 94 | + module = modules.get(layer_name) | ||
| 95 | + source_type = type(module).__name__ if module is not None else None | ||
| 96 | + quant_cls = AlgorithmRegistry.algo.get(alg_name, {}).get(source_type) | ||
| 97 | + if not is_quant_result_supported(quant_cls): | ||
| 98 | + unsupported.setdefault(alg_name, layer_name) | ||
| 99 | + | ||
| 100 | + if unsupported: | ||
| 101 | + details = ', '.join( | ||
| 102 | + "'{}' (e.g. layer '{}')".format(name, layer) | ||
| 103 | + for name, layer in sorted(unsupported.items()) | ||
| 104 | + ) | ||
| 105 | + raise ValueError( | ||
| 106 | + "quant_result_dir is not supported yet for algorithm(s) {}. " | ||
| 107 | + "The selected registered quantize class has no artifact contract; " | ||
| 108 | + "run without quant_result_dir, or choose a class that declares " | ||
| 109 | + "artifact support.".format(details) | ||
| 110 | + ) | ||
| 111 | + | ||
| 112 | + | ||
| 113 | +def _create_quant_result_writer(model, layer_config, quant_result_dir): | ||
| 114 | + """Validate the target dir, fingerprint the source model, build the writer. | ||
| 115 | + | ||
| 116 | + Ordering matters: the path and the algorithm declarations are checked before | ||
| 117 | + fingerprinting, and the fingerprint is taken before any quant module is | ||
| 118 | + inserted, so it always describes the original model. | ||
| 119 | + """ | ||
| 120 | + validate_quant_result_dir(quant_result_dir) | ||
| 121 | + if len(layer_config) == 0: | ||
| 122 | + raise ValueError( | ||
| 123 | + "no quantizable layer was found for this model and config, so " | ||
| 124 | + "there is no quant result to write to '{}'.".format(quant_result_dir) | ||
| 125 | + ) | ||
| 126 | + _check_quant_result_algorithms(model, layer_config) | ||
| 127 | + model_fingerprint = compute_model_fingerprint(model) | ||
| 128 | + return QuantResultWriter.create( | ||
| 129 | + quant_result_dir, tuple(layer_config.keys()), model_fingerprint | ||
| 130 | + ) | ||
| 131 | + | ||
| 132 | + | ||
| 133 | +def _finalize_quant_result(model, writer): | ||
| 134 | + """Publish now if the artifact is already complete, else arm a hook. | ||
| 135 | + | ||
| 136 | + Weight-only algorithms finalize during module insertion, so the manifest can | ||
| 137 | + be published as soon as the optimizer succeeds. Activation calibration only | ||
| 138 | + finalizes during the user's forward passes, so publication is deferred to a | ||
| 139 | + root forward hook that fires after a successful top-level forward; if that | ||
| 140 | + forward raises, the hook never runs and no manifest is written. | ||
| 141 | + """ | ||
| 142 | + if writer.publish_if_complete(): | ||
| 143 | + return | ||
| 144 | + | ||
| 145 | + handle_holder = {} | ||
| 146 | + | ||
| 147 | + def _publish_hook(module, args, output): | ||
| 148 | + if writer.publish_if_complete(): | ||
| 149 | + handle = handle_holder.pop('handle', None) | ||
| 150 | + if handle is not None: | ||
| 151 | + # Artifact is immutable once published; stop checking. | ||
| 152 | + handle.remove() | ||
| 153 | + return None | ||
| 154 | + | ||
| 155 | + handle_holder['handle'] = model.register_forward_hook(_publish_hook) | ||
| 156 | + LOGGER.logd( | ||
| 157 | + "Quant result '{}' has {} pending unit(s); the manifest will be " | ||
| 158 | + "published after calibration forward completes.".format( | ||
| 159 | + writer.result_dir, len(writer.pending_units) | ||
| 160 | + ), | ||
| 161 | + 'quantize', | ||
| 162 | + ) | ||
| 163 | + | ||
| 164 | + | ||
| 165 | + | ||
| 166 | +def convert(model, quant_result_dir=None): | ||
| 52 | """ | 167 | """ |
| 53 | Function: Convert quantized calibration model to quantized deployment model | 168 | Function: Convert quantized calibration model to quantized deployment model |
| 54 | - Parameter: model: quantized calibration model instance of Torch.nn.Module | 169 | + Parameter: model: quantized calibration model instance of Torch.nn.Module. |
| 170 | + When quant_result_dir is given, pass the original float | ||
| 171 | + model instead: it is restored to its calibrated state from | ||
| 172 | + the quant result first, then converted. | ||
| 173 | + quant_result_dir: str or None, directory of a published quant | ||
| 174 | + result produced by quantize(..., quant_result_dir=...). | ||
| 175 | + None keeps the original in-memory behaviour. | ||
| 55 | Return: None | 176 | Return: None |
| 56 | """ | 177 | """ |
| 178 | + if quant_result_dir is not None: | ||
| 179 | + # Restores calibration state only; which deploy module each layer becomes | ||
| 180 | + # is still decided by ReplaceNpuQuantModulePass below. | ||
| 181 | + reload_quant_params(model, quant_result_dir) | ||
| 182 | + | ||
| 57 | optimizer = ModelOptimizer() | 183 | optimizer = ModelOptimizer() |
| 58 | optimizer.add_pass(ReplaceNpuQuantModulePass()) | 184 | optimizer.add_pass(ReplaceNpuQuantModulePass()) |
| 59 | optimizer.do_optimizer(model) | 185 | optimizer.do_optimizer(model) |
| @@ -38,6 +38,13 @@ class GPTQuant(BaseQuantizeModule): | |||
| 38 | APIs: forward. | 38 | APIs: forward. |
| 39 | """ | 39 | """ |
| 40 | 40 | ||
| 41 | + # ``weight`` must be persisted, unlike AWQ: GPTQ rewrites it with the | ||
| 42 | + # Hessian-corrected result, which cannot be recomputed from the original | ||
| 43 | + # weight and the scales. Dropping it would silently deploy the un-optimized | ||
| 44 | + # weight after a reload. | ||
| 45 | + _quant_param_keys = ('scale_w', 'offset_w', 'weight') | ||
| 46 | + supports_quant_result = True | ||
| 47 | + | ||
| 41 | def __init__(self, ori_module, layer_name, quant_config): | 48 | def __init__(self, ori_module, layer_name, quant_config): |
| 42 | """ | 49 | """ |
| 43 | Function: init objective. | 50 | Function: init objective. |
| @@ -97,6 +104,9 @@ class GPTQuant(BaseQuantizeModule): | |||
| 97 | ), | 104 | ), |
| 98 | 'GPTQuant', | 105 | 'GPTQuant', |
| 99 | ) | 106 | ) |
| 107 | + # Dumped only after the optimized weight is in place, so the unit | ||
| 108 | + # holds the corrected weight and not the incoming one. | ||
| 109 | + self.maybe_dump_quant_params() | ||
| 100 | if self.cur_batch > self.quant_config.get('batch_num'): | 110 | if self.cur_batch > self.quant_config.get('batch_num'): |
| 101 | return self.fake_quant_forward(inputs) | 111 | return self.fake_quant_forward(inputs) |
| 102 | return F.linear(inputs, self.weight, self.bias) | 112 | return F.linear(inputs, self.weight, self.bias) |
| @@ -114,6 +124,29 @@ class GPTQuant(BaseQuantizeModule): | |||
| 114 | self.fake_quant_cache_ready = True | 124 | self.fake_quant_cache_ready = True |
| 115 | return F.linear(inputs, self.cached_dq_w, self.bias) | 125 | return F.linear(inputs, self.cached_dq_w, self.bias) |
| 116 | 126 | ||
| 127 | + def _on_quant_params_loaded(self): | ||
| 128 | + """Reloaded params are final: skip Hessian accumulation and GPTQ. | ||
| 129 | + | ||
| 130 | + ``cur_batch`` is pushed past ``batch_num`` so ``forward`` goes straight to | ||
| 131 | + fake-quant, and the Hessian state is released because it is only an | ||
| 132 | + accumulator for the optimization that already happened. | ||
| 133 | + """ | ||
| 134 | + # The constructor already put the *original* weight here, so absence in | ||
| 135 | + # the unit must be caught explicitly instead of falling back to it. | ||
| 136 | + self.require_loaded_quant_params('weight') | ||
| 137 | + # MXFP4 quantizes per-group at deploy time and legitimately leaves | ||
| 138 | + # scale_w as None (see cal_scale_offset_static), so it is only required | ||
| 139 | + # for the dtypes that actually produce it. | ||
| 140 | + if self.wts_type not in (MXFP4_E2M1,): | ||
| 141 | + self.require_loaded_quant_params('scale_w') | ||
| 142 | + if self.quant_config.get('weights_cfg').get('symmetric') is False: | ||
| 143 | + self.require_loaded_quant_params('offset_w') | ||
| 144 | + self.cur_batch = self.quant_config.get('batch_num') + 1 | ||
| 145 | + self.hessian = None | ||
| 146 | + self.nsamples = 0 | ||
| 147 | + self.fake_quant_cache_ready = False | ||
| 148 | + self.cached_dq_w = None | ||
| 149 | + | ||
| 117 | def get_opt_weight_and_quant_factor(self): | 150 | def get_opt_weight_and_quant_factor(self): |
| 118 | """ | 151 | """ |
| 119 | Get optimized weights and quantization factors. | 152 | Get optimized weights and quantization factors. |
| @@ -36,6 +36,13 @@ class LinearAWQuant(BaseQuantizeModule): | |||
| 36 | APIs: forward. | 36 | APIs: forward. |
| 37 | """ | 37 | """ |
| 38 | 38 | ||
| 39 | + # ``weight`` is deliberately absent: AWQ never modifies its own weight, it | ||
| 40 | + # scales the original module's weight by 1/scale. Both calibration fake-quant | ||
| 41 | + # and deploy recompute that transform from the original weight and ``scale``, | ||
| 42 | + # so persisting it would only duplicate the source model. | ||
| 43 | + _quant_param_keys = ('scale_w', 'offset_w', 'scale') | ||
| 44 | + supports_quant_result = True | ||
| 45 | + | ||
| 39 | def __init__(self, ori_module, layer_name, quant_config): | 46 | def __init__(self, ori_module, layer_name, quant_config): |
| 40 | """ | 47 | """ |
| 41 | Function: init objective. | 48 | Function: init objective. |
| @@ -56,6 +63,9 @@ class LinearAWQuant(BaseQuantizeModule): | |||
| 56 | if self.quant_config.get('weights_cfg').get("group_size") is not None: | 63 | if self.quant_config.get('weights_cfg').get("group_size") is not None: |
| 57 | self.group_size = self.quant_config.get('weights_cfg').get("group_size") | 64 | self.group_size = self.quant_config.get('weights_cfg').get("group_size") |
| 58 | self.calc_done = False | 65 | self.calc_done = False |
| 66 | + # True only on the reload path, where ori_module.weight was never scaled | ||
| 67 | + # in place and the transform has to be reapplied from weight and scale. | ||
| 68 | + self._scale_reapplied = False | ||
| 59 | enable_quant = quant_config.get("inputs_cfg").get("enable_quant") | 69 | enable_quant = quant_config.get("inputs_cfg").get("enable_quant") |
| 60 | if enable_quant is None or enable_quant: | 70 | if enable_quant is None or enable_quant: |
| 61 | self.act_granularity = quant_config.get('inputs_cfg').get('strategy') | 71 | self.act_granularity = quant_config.get('inputs_cfg').get('strategy') |
| @@ -63,7 +73,7 @@ class LinearAWQuant(BaseQuantizeModule): | |||
| 63 | 73 | ||
| 64 | def forward(self, inputs): | 74 | def forward(self, inputs): |
| 65 | """ | 75 | """ |
| 66 | - Function: LinearAWQuant foward funtion. | 76 | + Function: LinearAWQuant forward function. |
| 67 | Args: | 77 | Args: |
| 68 | inputs: data used for calibration in torch.tensor. | 78 | inputs: data used for calibration in torch.tensor. |
| 69 | """ | 79 | """ |
| @@ -94,13 +104,36 @@ class LinearAWQuant(BaseQuantizeModule): | |||
| 94 | "Calculate awq quant params of layer '{}' success!".format(self.layer_name), | 104 | "Calculate awq quant params of layer '{}' success!".format(self.layer_name), |
| 95 | "LinearAWQuant", | 105 | "LinearAWQuant", |
| 96 | ) | 106 | ) |
| 107 | + self.maybe_dump_quant_params() | ||
| 97 | return output | 108 | return output |
| 98 | 109 | ||
| 110 | + | ||
| 111 | + def scaled_weight(self): | ||
| 112 | + """The search-scaled weight AWQ fake-quant works on. | ||
| 113 | + | ||
| 114 | + During calibration ``apply_scale`` multiplies ``ori_module.weight`` in | ||
| 115 | + place, so that tensor is authoritative and is used as-is. After a reload | ||
| 116 | + there is no in-place result to read: the source model is untouched, so the | ||
| 117 | + transform is reapplied from ``self.weight`` (the original weight) and | ||
| 118 | + ``self.scale``. | ||
| 119 | + | ||
| 120 | + The two are close but not bit-identical. ``self.scale`` is ``1/scale_awq`` | ||
| 121 | + rounded to the weight precision, and dividing by it does not generally | ||
| 122 | + recover ``weight * scale_awq`` -- on bf16 a few hundred of 2048 elements | ||
| 123 | + differ by one ulp. Calibration therefore never takes the derived path, so | ||
| 124 | + ``quant_result_dir=None`` behaves exactly as before this was added. | ||
| 125 | + """ | ||
| 126 | + original = getattr(self, 'ori_module', None) | ||
| 127 | + if original is not None and not self._scale_reapplied: | ||
| 128 | + return original.weight.data | ||
| 129 | + scale = self.scale.to(device=self.weight.device) | ||
| 130 | + return (self.weight.data / scale).to(self.weight.dtype) | ||
| 131 | + | ||
| 99 | 132 | ||
| 100 | def fake_quant_forward(self, inputs): | 133 | def fake_quant_forward(self, inputs): |
| 101 | if not getattr(self, 'fake_quant_cache_ready', False): | 134 | if not getattr(self, 'fake_quant_cache_ready', False): |
| 102 | self.cached_dq_w = quant_dequant_weight( | 135 | self.cached_dq_w = quant_dequant_weight( |
| 103 | - self.ori_module.weight.data, | 136 | + self.scaled_weight(), |
| 104 | self.wts_type, | 137 | self.wts_type, |
| 105 | self.scale_w, | 138 | self.scale_w, |
| 106 | self.offset_w, | 139 | self.offset_w, |
| @@ -109,3 +142,25 @@ class LinearAWQuant(BaseQuantizeModule): | |||
| 109 | self.fake_quant_cache_ready = True | 142 | self.fake_quant_cache_ready = True |
| 110 | x = inputs * self.scale.to(device=inputs.device, dtype=inputs.dtype) | 143 | x = inputs * self.scale.to(device=inputs.device, dtype=inputs.dtype) |
| 111 | return F.linear(x, self.cached_dq_w, self.bias) | 144 | return F.linear(x, self.cached_dq_w, self.bias) |
| 145 | + | ||
| 146 | + def _on_quant_params_loaded(self): | ||
| 147 | + """Reloaded params are final: never search a scale again. | ||
| 148 | + | ||
| 149 | + ``scale`` is what makes the module usable, so a unit without it is an | ||
| 150 | + incomplete artifact rather than something to re-calibrate around. The | ||
| 151 | + fake-quant cache is dropped so it is rebuilt from the loaded params. | ||
| 152 | + """ | ||
| 153 | + self.require_loaded_quant_params('scale') | ||
| 154 | + # Same per-dtype rule as GPTQ: forward() only computes scale_w/offset_w for | ||
| 155 | + # these types, so only they may be demanded back. Failing here beats | ||
| 156 | + # entering calc_done with scale_w=None and breaking at fake-quant time. | ||
| 157 | + if self.wts_type in (INT4, INT8, FLOAT4_E2M1): | ||
| 158 | + self.require_loaded_quant_params('scale_w') | ||
| 159 | + if self.wts_symmetric is False: | ||
| 160 | + self.require_loaded_quant_params('offset_w') | ||
| 161 | + self.calc_done = True | ||
| 162 | + # ori_module here is an untouched float layer, so the in-place scaled | ||
| 163 | + # weight does not exist and must be derived instead. | ||
| 164 | + self._scale_reapplied = True | ||
| 165 | + self.fake_quant_cache_ready = False | ||
| 166 | + self.cached_dq_w = None | ||
| @@ -45,6 +45,19 @@ class MinMaxQuant(BaseQuantizeModule): | |||
| 45 | APIs: forward. | 45 | APIs: forward. |
| 46 | """ | 46 | """ |
| 47 | 47 | ||
| 48 | + # Quant params this algorithm produces. scale_w/offset_w are final right | ||
| 49 | + # after __init__; scale_d/offset_d only exist once activation calibration | ||
| 50 | + # completes; scale_w1/scale_w2 replace scale_w on the fp8/fp4 path. | ||
| 51 | + _quant_param_keys = ( | ||
| 52 | + 'scale_w', | ||
| 53 | + 'offset_w', | ||
| 54 | + 'scale_d', | ||
| 55 | + 'offset_d', | ||
| 56 | + 'scale_w1', | ||
| 57 | + 'scale_w2', | ||
| 58 | + ) | ||
| 59 | + supports_quant_result = True | ||
| 60 | + | ||
| 48 | def __init__(self, ori_module, layer_name, quant_config): | 61 | def __init__(self, ori_module, layer_name, quant_config): |
| 49 | """ | 62 | """ |
| 50 | Function: init objective. | 63 | Function: init objective. |
| @@ -147,8 +160,42 @@ class MinMaxQuant(BaseQuantizeModule): | |||
| 147 | ), | 160 | ), |
| 148 | 'MinMaxQuant', | 161 | 'MinMaxQuant', |
| 149 | ) | 162 | ) |
| 163 | + # Activation params are final only here, so this is this layer's | ||
| 164 | + # single dump point on the activation path. | ||
| 165 | + self.maybe_dump_quant_params() | ||
| 150 | return fp_out | 166 | return fp_out |
| 151 | 167 | ||
| 168 | + def _on_quant_result_writer_attached(self): | ||
| 169 | + """Weight-only MinMax is already final at construction: dump now. | ||
| 170 | + | ||
| 171 | + With activation quantization enabled, scale_d/offset_d are still | ||
| 172 | + missing, so the dump waits for the batch_num-th calibration forward. | ||
| 173 | + """ | ||
| 174 | + if self.weight_compress_only: | ||
| 175 | + self.maybe_dump_quant_params() | ||
| 176 | + | ||
| 177 | + def _on_quant_params_loaded(self): | ||
| 178 | + """Drop fake-quant caches and mark calibration as already complete.""" | ||
| 179 | + if self.scale_w1 is not None or self.scale_w2 is not None: | ||
| 180 | + self.require_loaded_quant_params('scale_w1', 'scale_w2') | ||
| 181 | + else: | ||
| 182 | + self.require_loaded_quant_params('scale_w') | ||
| 183 | + if self.wts_symmetric is False: | ||
| 184 | + self.require_loaded_quant_params('offset_w') | ||
| 185 | + | ||
| 186 | + if not self.weight_compress_only: | ||
| 187 | + self.require_loaded_quant_params('scale_d') | ||
| 188 | + if self.act_symmetric is False: | ||
| 189 | + self.require_loaded_quant_params('offset_d') | ||
| 190 | + | ||
| 191 | + self.fake_quant_cache_ready = False | ||
| 192 | + self.cached_dq_w = None | ||
| 193 | + self.cached_bias = None | ||
| 194 | + if self.scale_d is not None and self.batch_num is not None: | ||
| 195 | + # Calibration is done; the next forward must fake-quant instead of | ||
| 196 | + # re-collecting min/max. | ||
| 197 | + self.cur_batch = self.batch_num | ||
| 198 | + | ||
| 152 | 199 | ||
| 153 | def fake_quant_forward(self, inputs): | 200 | def fake_quant_forward(self, inputs): |
| 154 | if not getattr(self, 'fake_quant_cache_ready', False): | 201 | if not getattr(self, 'fake_quant_cache_ready', False): |
| @@ -15,12 +15,37 @@ | |||
| 15 | # limitations under the License. | 15 | # limitations under the License. |
| 16 | # ---------------------------------------------------------------------------- | 16 | # ---------------------------------------------------------------------------- |
| 17 | 17 | ||
| 18 | +import torch | ||
| 18 | import torch.nn as nn | 19 | import torch.nn as nn |
| 19 | 20 | ||
| 21 | +from amct_pytorch.common.utils.log import LOGGER | ||
| 22 | + | ||
| 23 | +LOG_TAG = 'BaseQuantizeModule' | ||
| 24 | + | ||
| 20 | 25 | ||
| 21 | class BaseQuantizeModule(nn.Module): | 26 | class BaseQuantizeModule(nn.Module): |
| 27 | + """Base class of calibration modules with quant param export and load. | ||
| 28 | + | ||
| 29 | + Each subclass declares its own quant param attribute names so the base class | ||
| 30 | + does not hard-code the algorithm-specific names. The generic export, load | ||
| 31 | + and dump plumbing can then be shared by MinMax, AWQ and GPTQ. | ||
| 32 | + """ | ||
| 33 | + | ||
| 34 | + # Attribute names holding this algorithm's quant params. Values may be | ||
| 35 | + # tensors, JSON-serializable scalars, or None (skipped when exporting). | ||
| 36 | + _quant_param_keys = () | ||
| 37 | + | ||
| 38 | + # A declaration consumed by the Classic artifact path. It only says that | ||
| 39 | + # the class participates in the artifact contract; the writer still checks | ||
| 40 | + # the concrete values produced at each algorithm's finalization point. | ||
| 41 | + supports_quant_result = False | ||
| 42 | + | ||
| 22 | def __init__(self, ori_module, layer_name, quant_config): | 43 | def __init__(self, ori_module, layer_name, quant_config): |
| 23 | super().__init__() | 44 | super().__init__() |
| 45 | + if not isinstance(layer_name, str) or not layer_name: | ||
| 46 | + raise ValueError( | ||
| 47 | + 'layer_name must be a non-empty string, got {!r}.'.format(layer_name) | ||
| 48 | + ) | ||
| 24 | self.ori_module_type = None | 49 | self.ori_module_type = None |
| 25 | self.act_type = None | 50 | self.act_type = None |
| 26 | self.wts_type = None | 51 | self.wts_type = None |
| @@ -33,5 +58,185 @@ class BaseQuantizeModule(nn.Module): | |||
| 33 | self.scale = None | 58 | self.scale = None |
| 34 | self.clip_max = None | 59 | self.clip_max = None |
| 35 | 60 | ||
| 61 | + self.layer_name = layer_name | ||
| 62 | + self.quant_config = quant_config | ||
| 63 | + self._quant_result_writer = None | ||
| 64 | + self._quant_params_dumped = False | ||
| 65 | + | ||
| 36 | def forward(self, inputs): | 66 | def forward(self, inputs): |
| 37 | pass | 67 | pass |
| 68 | + | ||
| 69 | + | ||
| 70 | + def quant_param_keys(cls): | ||
| 71 | + """Validated tuple of declared quant param attribute names.""" | ||
| 72 | + keys = tuple(cls._quant_param_keys) | ||
| 73 | + for key in keys: | ||
| 74 | + if not isinstance(key, str) or not key: | ||
| 75 | + raise TypeError( | ||
| 76 | + "quant param names declared by {} must be non-empty strings, " | ||
| 77 | + "got {!r}.".format(cls.__name__, key) | ||
| 78 | + ) | ||
| 79 | + if len(set(keys)) != len(keys): | ||
| 80 | + raise ValueError( | ||
| 81 | + '{} declares duplicate quant param names.'.format(cls.__name__) | ||
| 82 | + ) | ||
| 83 | + return keys | ||
| 84 | + | ||
| 85 | + def export_quant_params(self): | ||
| 86 | + """Split declared quant params into ``(tensors, metadata)``. | ||
| 87 | + | ||
| 88 | + Tensors come back detached, on CPU and contiguous so they are safe to | ||
| 89 | + serialize; non-tensor values go to metadata; ``None`` values are absent | ||
| 90 | + from both, which is how "this algorithm/config never produced it" is | ||
| 91 | + represented. | ||
| 92 | + """ | ||
| 93 | + tensors = {} | ||
| 94 | + metadata = {} | ||
| 95 | + for key in self.quant_param_keys(): | ||
| 96 | + value = getattr(self, key, None) | ||
| 97 | + if value is None: | ||
| 98 | + continue | ||
| 99 | + if isinstance(value, torch.Tensor): | ||
| 100 | + tensors[key] = value.detach().cpu().contiguous() | ||
| 101 | + else: | ||
| 102 | + metadata[key] = value | ||
| 103 | + return tensors, metadata | ||
| 104 | + | ||
| 105 | + def load_quant_params(self, tensors, metadata=None): | ||
| 106 | + """Restore quant params previously exported for this module. | ||
| 107 | + | ||
| 108 | + Rejects undeclared keys and keys present in both payloads. An existing | ||
| 109 | + tensor attribute is copied into in place (preserving its device, dtype | ||
| 110 | + and shape); an attribute currently ``None`` adopts the persisted dtype | ||
| 111 | + and moves to this module's device. | ||
| 112 | + """ | ||
| 113 | + tensors = dict(tensors or {}) | ||
| 114 | + metadata = dict(metadata or {}) | ||
| 115 | + declared = set(self.quant_param_keys()) | ||
| 116 | + self._check_loadable_keys(tensors, metadata, declared) | ||
| 117 | + | ||
| 118 | + device = self._quant_param_device() | ||
| 119 | + for key, value in tensors.items(): | ||
| 120 | + self._load_tensor_param(key, value, device) | ||
| 121 | + for key, value in metadata.items(): | ||
| 122 | + setattr(self, key, value) | ||
| 123 | + | ||
| 124 | + # Recorded so a subclass hook can tell "the artifact carried this" from | ||
| 125 | + # "the constructor happened to leave an attribute here", which matters for | ||
| 126 | + # params that always pre-exist, such as GPTQ's weight. | ||
| 127 | + self._loaded_quant_param_keys = frozenset(tensors) | frozenset(metadata) | ||
| 128 | + self._on_quant_params_loaded() | ||
| 129 | + LOGGER.logd( | ||
| 130 | + "Loaded quant params of layer '{}': {} tensor(s), {} scalar(s)".format( | ||
| 131 | + self.layer_name, len(tensors), len(metadata) | ||
| 132 | + ), | ||
| 133 | + LOG_TAG, | ||
| 134 | + ) | ||
| 135 | + | ||
| 136 | + def _check_loadable_keys(self, tensors, metadata, declared): | ||
| 137 | + """Validate the incoming key sets against the declared params.""" | ||
| 138 | + unknown = sorted((set(tensors) | set(metadata)) - declared) | ||
| 139 | + if unknown: | ||
| 140 | + raise ValueError( | ||
| 141 | + "quant params {} are not declared by {}; refusing to load them " | ||
| 142 | + "onto layer '{}'.".format(unknown, type(self).__name__, self.layer_name) | ||
| 143 | + ) | ||
| 144 | + overlap = sorted(set(tensors) & set(metadata)) | ||
| 145 | + if overlap: | ||
| 146 | + raise ValueError( | ||
| 147 | + "quant params {} appear as both tensor and non-tensor values " | ||
| 148 | + "for layer '{}'.".format(overlap, self.layer_name) | ||
| 149 | + ) | ||
| 150 | + | ||
| 151 | + | ||
| 152 | + def _load_tensor_param(self, key, value, device): | ||
| 153 | + """Restore one tensor quant param, in place when one already exists.""" | ||
| 154 | + if not isinstance(value, torch.Tensor): | ||
| 155 | + raise TypeError( | ||
| 156 | + "quant param '{}' of layer '{}' must be a torch.Tensor, got {}.".format( | ||
| 157 | + key, self.layer_name, type(value).__name__ | ||
| 158 | + ) | ||
| 159 | + ) | ||
| 160 | + current = getattr(self, key, None) | ||
| 161 | + if not isinstance(current, torch.Tensor): | ||
| 162 | + setattr(self, key, value.detach().to(device)) | ||
| 163 | + return | ||
| 164 | + if tuple(current.shape) != tuple(value.shape): | ||
| 165 | + raise ValueError( | ||
| 166 | + "quant param '{}' of layer '{}' has shape {} but the persisted " | ||
| 167 | + "value has shape {}.".format( | ||
| 168 | + key, self.layer_name, tuple(current.shape), tuple(value.shape) | ||
| 169 | + ) | ||
| 170 | + ) | ||
| 171 | + if current.dtype != value.dtype: | ||
| 172 | + raise ValueError( | ||
| 173 | + "quant param '{}' of layer '{}' has dtype {} but the persisted " | ||
| 174 | + "value has dtype {}.".format( | ||
| 175 | + key, self.layer_name, current.dtype, value.dtype | ||
| 176 | + ) | ||
| 177 | + ) | ||
| 178 | + current.copy_(value.to(current.device)) | ||
| 179 | + | ||
| 180 | + def _quant_param_device(self): | ||
| 181 | + """Device a restored quant param should live on: follow the weight.""" | ||
| 182 | + weight = getattr(self, 'weight', None) | ||
| 183 | + if isinstance(weight, torch.Tensor): | ||
| 184 | + return weight.device | ||
| 185 | + for parameter in self.parameters(recurse=True): | ||
| 186 | + return parameter.device | ||
| 187 | + return torch.device('cpu') | ||
| 188 | + | ||
| 189 | + def attach_quant_result_writer(self, writer): | ||
| 190 | + """Attach the run's result writer, then run the subclass hook. | ||
| 191 | + | ||
| 192 | + Injected right after this module replaces the original one, which is | ||
| 193 | + also the point where an algorithm whose params are already final (pure | ||
| 194 | + weight-only) can dump immediately. | ||
| 195 | + """ | ||
| 196 | + self._quant_result_writer = writer | ||
| 197 | + self._on_quant_result_writer_attached() | ||
| 198 | + | ||
| 199 | + def maybe_dump_quant_params(self): | ||
| 200 | + """Dump this layer's quant params once, if a writer is attached. | ||
| 201 | + | ||
| 202 | + Returns True when this call performed the dump. A failed write leaves | ||
| 203 | + the module undumped so the unit stays pending and no manifest appears. | ||
| 204 | + """ | ||
| 205 | + if self._quant_result_writer is None or self._quant_params_dumped: | ||
| 206 | + return False | ||
| 207 | + # Imported lazily: quant_result lives under classic/, which this module | ||
| 208 | + # must not depend on at import time. | ||
| 209 | + from amct_pytorch.classic.quant_result import to_jsonable | ||
| 210 | + | ||
| 211 | + tensors, non_tensor_params = self.export_quant_params() | ||
| 212 | + self._quant_result_writer.write_unit( | ||
| 213 | + layer_name=self.layer_name, | ||
| 214 | + tensors=tensors, | ||
| 215 | + non_tensor_params=non_tensor_params, | ||
| 216 | + quant_config=to_jsonable( | ||
| 217 | + self.quant_config, | ||
| 218 | + "quant_config of layer '{}'".format(self.layer_name), | ||
| 219 | + ), | ||
| 220 | + quant_module_type=type(self).__name__, | ||
| 221 | + ori_module_type=self.ori_module_type, | ||
| 222 | + ) | ||
| 223 | + self._quant_params_dumped = True | ||
| 224 | + return True | ||
| 225 | + | ||
| 226 | + def require_loaded_quant_params(self, *keys): | ||
| 227 | + """Fail unless the loaded unit carried every one of ``keys``.""" | ||
| 228 | + loaded_keys = getattr(self, '_loaded_quant_param_keys', frozenset()) | ||
| 229 | + missing = sorted(key for key in keys if key not in loaded_keys) | ||
| 230 | + if missing: | ||
| 231 | + raise ValueError( | ||
| 232 | + "quant result of layer '{}' is missing required quant param(s) {} " | ||
| 233 | + "for {}; refusing to reload an incomplete unit.".format( | ||
| 234 | + self.layer_name, missing, type(self).__name__ | ||
| 235 | + ) | ||
| 236 | + ) | ||
| 237 | + | ||
| 238 | + def _on_quant_params_loaded(self): | ||
| 239 | + """Hook: subclasses drop caches / restore progress after a load.""" | ||
| 240 | + | ||
| 241 | + def _on_quant_result_writer_attached(self): | ||
| 242 | + """Hook: subclasses may dump right away if their params are final.""" | ||
| @@ -13,6 +13,10 @@ loguru | |||
| 13 | datasets | 13 | datasets |
| 14 | accelerate | 14 | accelerate |
| 15 | compressed_tensors==0.15.0.1 | 15 | compressed_tensors==0.15.0.1 |
| 16 | +# Quant result artifacts (amct_pytorch/classic/quant_result.py) and the deploy | ||
| 17 | +# export path read/write safetensors directly, so declare it explicitly rather | ||
| 18 | +# than relying on it arriving as a transformers transitive dependency. | ||
| 19 | +safetensors>=0.4.3 | ||
| 16 | scipy | 20 | scipy |
| 17 | torchao==0.12.0 | 21 | torchao==0.12.0 |
| 18 | einops | 22 | einops |
| @@ -710,7 +710,7 @@ class TestConfigParse(unittest.TestCase): | |||
| 710 | try: | 710 | try: |
| 711 | parse_config(model, cfg, AlgorithmRegistry) | 711 | parse_config(model, cfg, AlgorithmRegistry) |
| 712 | except Exception as e: | 712 | except Exception as e: |
| 713 | - self.assertIn('Not support algorithm AA, pls regiter it first', str(e)) | 713 | + self.assertIn('Not support algorithm AA', str(e)) |
| 714 | 714 | ||
| 715 | def test_customize_algo_cfg(self): | 715 | def test_customize_algo_cfg(self): |
| 716 | cfg = { | 716 | cfg = { |
| @@ -17,19 +17,37 @@ | |||
| 17 | # ---------------------------------------------------------------------------- | 17 | # ---------------------------------------------------------------------------- |
| 18 | 18 | ||
| 19 | import sys | 19 | import sys |
| 20 | -from unittest.mock import patch | 20 | +from unittest.mock import MagicMock, patch |
| 21 | 21 | ||
| 22 | import pytest | 22 | import pytest |
| 23 | +import torch | ||
| 23 | from torch import nn | 24 | from torch import nn |
| 24 | 25 | ||
| 26 | +from amct_pytorch.classic.optimizer import InsertQuantizeModulePass | ||
| 27 | +from amct_pytorch.classic.quant_result import ( | ||
| 28 | + MANIFEST_FORMAT, | ||
| 29 | + MANIFEST_NAME, | ||
| 30 | + SCHEMA_VERSION, | ||
| 31 | + compute_model_fingerprint, | ||
| 32 | + is_quant_result_supported, | ||
| 33 | + load_quant_manifest, | ||
| 34 | + read_quant_unit, | ||
| 35 | + validate_result_directory, | ||
| 36 | +) | ||
| 25 | from amct_pytorch.classic.quantize import ( | 37 | from amct_pytorch.classic.quantize import ( |
| 26 | algorithm_register, | 38 | algorithm_register, |
| 27 | convert, | 39 | convert, |
| 28 | quantize, | 40 | quantize, |
| 29 | ) | 41 | ) |
| 42 | +from amct_pytorch.classic.quantize_op.minmax_module import MinMaxQuant | ||
| 43 | +from amct_pytorch.classic.quantize_op.linear_awq_module import LinearAWQuant | ||
| 44 | +from amct_pytorch.classic.quantize_op.gptq_module import GPTQuant | ||
| 45 | +from amct_pytorch.classic.quantize_op.smooth_quant_module import SmoothQuant | ||
| 46 | +from amct_pytorch.common.utils.model_util import ModuleHelper | ||
| 30 | from amct_pytorch.quantize_op import BaseQuantizeModule | 47 | from amct_pytorch.quantize_op import BaseQuantizeModule |
| 31 | 48 | ||
| 32 | QUANTIZE_MODULE = sys.modules["amct_pytorch.classic.quantize"] | 49 | QUANTIZE_MODULE = sys.modules["amct_pytorch.classic.quantize"] |
| 50 | +QUANT_RESULT_MODULE = sys.modules["amct_pytorch.classic.quant_result"] | ||
| 33 | 51 | ||
| 34 | 52 | ||
| 35 | def test_base_quantize_module_keeps_public_import_paths(): | 53 | def test_base_quantize_module_keeps_public_import_paths(): |
| @@ -46,6 +64,19 @@ def test_classic_quantize_imports_without_ptq_package_layer(): | |||
| 46 | assert classic_algorithm_register is algorithm_register | 64 | assert classic_algorithm_register is algorithm_register |
| 47 | 65 | ||
| 48 | 66 | ||
| 67 | +def test_artifact_support_is_declared_by_the_registered_quant_class(): | ||
| 68 | + assert is_quant_result_supported(MinMaxQuant) | ||
| 69 | + assert is_quant_result_supported(LinearAWQuant) | ||
| 70 | + assert is_quant_result_supported(GPTQuant) | ||
| 71 | + assert not is_quant_result_supported(SmoothQuant) | ||
| 72 | + | ||
| 73 | + class _DeclaredWithoutParams(BaseQuantizeModule): | ||
| 74 | + supports_quant_result = True | ||
| 75 | + | ||
| 76 | + # A declaration alone cannot make an unloadable class eligible. | ||
| 77 | + assert not is_quant_result_supported(_DeclaredWithoutParams) | ||
| 78 | + | ||
| 79 | + | ||
| 49 | class _DummyQuantOp(BaseQuantizeModule): | 80 | class _DummyQuantOp(BaseQuantizeModule): |
| 50 | pass | 81 | pass |
| 51 | 82 | ||
| @@ -117,7 +148,7 @@ def test_quantize_uses_default_config_when_none(): | |||
| 117 | args, _ = mock_parse.call_args | 148 | args, _ = mock_parse.call_args |
| 118 | assert args[0] is model | 149 | assert args[0] is model |
| 119 | assert args[1] is sentinel_default | 150 | assert args[1] is sentinel_default |
| 120 | - mock_pass_cls.assert_called_once_with(sentinel_layer) | 151 | + mock_pass_cls.assert_called_once_with(sentinel_layer, quant_result_writer=None) |
| 121 | opt.add_pass.assert_called_once_with(mock_pass_cls.return_value) | 152 | opt.add_pass.assert_called_once_with(mock_pass_cls.return_value) |
| 122 | opt.do_optimizer.assert_called_once_with(model) | 153 | opt.do_optimizer.assert_called_once_with(model) |
| 123 | 154 | ||
| @@ -149,3 +180,457 @@ def test_quantize_rejects_non_dict_config(): | |||
| 149 | model = nn.Linear(4, 4) | 180 | model = nn.Linear(4, 4) |
| 150 | with pytest.raises(TypeError): | 181 | with pytest.raises(TypeError): |
| 151 | quantize(model, "bad_config") | 182 | quantize(model, "bad_config") |
| 183 | + | ||
| 184 | + | ||
| 185 | +# -------------------------------------------------------------------------- | ||
| 186 | +# quant_result_dir: quant result artifact persistence. | ||
| 187 | +# | ||
| 188 | +# The default (None) path keeps the old user-visible behaviour: no directory and | ||
| 189 | +# no artifact writes. With a directory, the manifest appears only once every | ||
| 190 | +# expected unit has been dumped. | ||
| 191 | +# -------------------------------------------------------------------------- | ||
| 192 | + | ||
| 193 | +WEIGHT_ONLY_CFG = { | ||
| 194 | + 'batch_num': 1, | ||
| 195 | + 'quant_cfg': { | ||
| 196 | + 'weights': {'type': 'int8', 'symmetric': True, 'strategy': 'channel'} | ||
| 197 | + }, | ||
| 198 | + 'algorithm': {'minmax'}, | ||
| 199 | +} | ||
| 200 | + | ||
| 201 | +ACTIVATION_CFG = { | ||
| 202 | + 'batch_num': 2, | ||
| 203 | + 'quant_cfg': { | ||
| 204 | + 'weights': {'type': 'int8', 'symmetric': True, 'strategy': 'channel'}, | ||
| 205 | + 'inputs': {'type': 'int8', 'symmetric': True, 'strategy': 'tensor'}, | ||
| 206 | + }, | ||
| 207 | + 'algorithm': {'minmax'}, | ||
| 208 | +} | ||
| 209 | + | ||
| 210 | + | ||
| 211 | +class _TinyModel(nn.Module): | ||
| 212 | + def __init__(self): | ||
| 213 | + super().__init__() | ||
| 214 | + self.fc1 = nn.Linear(64, 32) | ||
| 215 | + self.fc2 = nn.Linear(32, 16) | ||
| 216 | + | ||
| 217 | + def forward(self, inputs): | ||
| 218 | + return self.fc2(self.fc1(inputs)) | ||
| 219 | + | ||
| 220 | + | ||
| 221 | +def _tiny_model(): | ||
| 222 | + # int8 minmax only accepts bf16/fp16 source weights. | ||
| 223 | + return _TinyModel().to(torch.bfloat16) | ||
| 224 | + | ||
| 225 | + | ||
| 226 | +def _tiny_inputs(batch=8): | ||
| 227 | + return torch.randn(batch, 64).to(torch.bfloat16) | ||
| 228 | + | ||
| 229 | + | ||
| 230 | +def test_quantize_without_result_dir_creates_nothing(tmp_path): | ||
| 231 | + model = _tiny_model() | ||
| 232 | + before = sorted(p.name for p in tmp_path.iterdir()) | ||
| 233 | + quantize(model, WEIGHT_ONLY_CFG) | ||
| 234 | + assert sorted(p.name for p in tmp_path.iterdir()) == before | ||
| 235 | + assert not hasattr(model.fc1, 'quant_result_writer') | ||
| 236 | + assert not hasattr(model.fc1, 'quant_params_dumped') | ||
| 237 | + | ||
| 238 | + | ||
| 239 | +def test_quantize_without_result_dir_keeps_pass_signature(): | ||
| 240 | + model = nn.Linear(4, 4) | ||
| 241 | + with ( | ||
| 242 | + patch.object(QUANTIZE_MODULE, "parse_config", return_value={"layer": "cfg"}), | ||
| 243 | + patch.object(QUANTIZE_MODULE, "ModelOptimizer"), | ||
| 244 | + patch.object(QUANTIZE_MODULE, "InsertQuantizeModulePass") as mock_pass_cls, | ||
| 245 | + patch.object(QUANTIZE_MODULE, "QuantResultWriter") as mock_writer_cls, | ||
| 246 | + ): | ||
| 247 | + quantize(model, {"granularity": "tensor"}) | ||
| 248 | + | ||
| 249 | + mock_pass_cls.assert_called_once_with({"layer": "cfg"}, quant_result_writer=None) | ||
| 250 | + mock_writer_cls.create.assert_not_called() | ||
| 251 | + | ||
| 252 | + | ||
| 253 | + | ||
| 254 | +def test_quantize_rejects_non_str_result_dir(bad_dir): | ||
| 255 | + with pytest.raises(TypeError): | ||
| 256 | + quantize(_tiny_model(), WEIGHT_ONLY_CFG, quant_result_dir=bad_dir) | ||
| 257 | + | ||
| 258 | + | ||
| 259 | +def test_quantize_rejects_empty_result_dir(): | ||
| 260 | + with pytest.raises(ValueError, match="non-empty path"): | ||
| 261 | + quantize(_tiny_model(), WEIGHT_ONLY_CFG, quant_result_dir="") | ||
| 262 | + | ||
| 263 | + | ||
| 264 | +def test_quantize_rejects_missing_result_dir(tmp_path): | ||
| 265 | + missing = tmp_path / "missing" | ||
| 266 | + with pytest.raises(ValueError, match="must be an existing directory"): | ||
| 267 | + quantize(_tiny_model(), WEIGHT_ONLY_CFG, quant_result_dir=str(missing)) | ||
| 268 | + | ||
| 269 | + | ||
| 270 | +def test_quantize_preserves_unmanaged_files_in_result_dir(tmp_path): | ||
| 271 | + existing = tmp_path / "with_other_files" | ||
| 272 | + existing.mkdir() | ||
| 273 | + old_file = existing / "old.safetensors" | ||
| 274 | + old_file.write_bytes(b"old") | ||
| 275 | + quantize(_tiny_model(), WEIGHT_ONLY_CFG, quant_result_dir=str(existing)) | ||
| 276 | + assert (existing / MANIFEST_NAME).is_file() | ||
| 277 | + assert old_file.read_bytes() == b"old" | ||
| 278 | + | ||
| 279 | + | ||
| 280 | +def test_quantize_rejects_result_dir_that_is_a_file(tmp_path): | ||
| 281 | + path = tmp_path / "a_file" | ||
| 282 | + path.write_text("x", encoding="utf-8") | ||
| 283 | + with pytest.raises(ValueError, match="not a file"): | ||
| 284 | + quantize(_tiny_model(), WEIGHT_ONLY_CFG, quant_result_dir=str(path)) | ||
| 285 | + | ||
| 286 | + | ||
| 287 | +def test_quantize_rejects_result_dir_when_no_layer_is_quantizable(tmp_path): | ||
| 288 | + out = tmp_path / "out" | ||
| 289 | + out.mkdir() | ||
| 290 | + # float32 weights are not accepted by int8 minmax, so nothing matches. | ||
| 291 | + with pytest.raises(ValueError, match="no quantizable layer"): | ||
| 292 | + quantize(_TinyModel(), WEIGHT_ONLY_CFG, quant_result_dir=str(out)) | ||
| 293 | + assert list(out.iterdir()) == [] | ||
| 294 | + | ||
| 295 | + | ||
| 296 | +def test_quantize_rejects_result_dir_when_the_root_is_the_quant_layer(tmp_path): | ||
| 297 | + out = tmp_path / "out" | ||
| 298 | + out.mkdir() | ||
| 299 | + # parse_config names the root layer '', and module replacement is keyed by | ||
| 300 | + # name, so the root is never actually replaced: its unit could never be | ||
| 301 | + # dumped and the manifest would stay unpublished forever. | ||
| 302 | + root_linear = nn.Linear(64, 32).to(torch.bfloat16) | ||
| 303 | + with pytest.raises(ValueError, match="itself a quantizable layer"): | ||
| 304 | + quantize(root_linear, WEIGHT_ONLY_CFG, quant_result_dir=str(out)) | ||
| 305 | + assert list(out.iterdir()) == [] | ||
| 306 | + | ||
| 307 | + | ||
| 308 | +def test_quantize_rejects_root_quant_layer_without_result_dir(): | ||
| 309 | + root_linear = nn.Linear(64, 32).to(torch.bfloat16) | ||
| 310 | + with pytest.raises(ValueError, match="itself a quantizable layer"): | ||
| 311 | + quantize(root_linear, WEIGHT_ONLY_CFG) | ||
| 312 | + | ||
| 313 | + | ||
| 314 | +def test_weight_only_publishes_without_any_forward(tmp_path): | ||
| 315 | + out = tmp_path / "wo" | ||
| 316 | + out.mkdir() | ||
| 317 | + model = _tiny_model() | ||
| 318 | + quantize(model, WEIGHT_ONLY_CFG, quant_result_dir=str(out)) | ||
| 319 | + | ||
| 320 | + assert sorted(p.name for p in out.iterdir()) == sorted( | ||
| 321 | + [MANIFEST_NAME, "fc1.safetensors", "fc2.safetensors"] | ||
| 322 | + ) | ||
| 323 | + manifest = load_quant_manifest(str(out)) | ||
| 324 | + validate_result_directory(str(out), manifest) | ||
| 325 | + assert manifest["format"] == MANIFEST_FORMAT | ||
| 326 | + assert manifest["schema_version"] == SCHEMA_VERSION | ||
| 327 | + assert [unit["layer_name"] for unit in manifest["units"]] == ["fc1", "fc2"] | ||
| 328 | + assert all(u["quant_module_type"] == "MinMaxQuant" for u in manifest["units"]) | ||
| 329 | + assert all(u["ori_module_type"] == "Linear" for u in manifest["units"]) | ||
| 330 | + # The algorithm is recorded inside the full layer-level quant_config. | ||
| 331 | + assert all( | ||
| 332 | + u["quant_config"]["algorithm"] == {"minmax": {}} for u in manifest["units"] | ||
| 333 | + ) | ||
| 334 | + assert all("batch_num" in u["quant_config"] for u in manifest["units"]) | ||
| 335 | + | ||
| 336 | + tensors, _, _ = read_quant_unit(str(out), "fc1") | ||
| 337 | + assert torch.equal(tensors["scale_w"], model.fc1.scale_w.cpu()) | ||
| 338 | + | ||
| 339 | + | ||
| 340 | +def test_activation_publishes_only_after_the_last_calibration_batch(tmp_path): | ||
| 341 | + out = tmp_path / "act" | ||
| 342 | + out.mkdir() | ||
| 343 | + model = _tiny_model() | ||
| 344 | + quantize(model, ACTIVATION_CFG, quant_result_dir=str(out)) | ||
| 345 | + manifest_path = out / MANIFEST_NAME | ||
| 346 | + | ||
| 347 | + assert list(out.iterdir()) == [], "nothing is final until calibration runs" | ||
| 348 | + | ||
| 349 | + inputs = _tiny_inputs() | ||
| 350 | + model(inputs) | ||
| 351 | + assert not manifest_path.exists(), "batch_num-1 forwards must not publish" | ||
| 352 | + | ||
| 353 | + model(inputs) | ||
| 354 | + assert manifest_path.exists() | ||
| 355 | + manifest = load_quant_manifest(str(out)) | ||
| 356 | + validate_result_directory(str(out), manifest) | ||
| 357 | + assert len(manifest["units"]) == 2 | ||
| 358 | + | ||
| 359 | + tensors, _, _ = read_quant_unit(str(out), "fc2") | ||
| 360 | + assert sorted(tensors) == ["scale_d", "scale_w"] | ||
| 361 | + assert torch.equal(tensors["scale_d"], model.fc2.scale_d.cpu()) | ||
| 362 | + | ||
| 363 | + | ||
| 364 | +def test_publish_hook_is_removed_after_publishing(tmp_path): | ||
| 365 | + out = tmp_path / "hook" | ||
| 366 | + out.mkdir() | ||
| 367 | + model = _tiny_model() | ||
| 368 | + quantize(model, ACTIVATION_CFG, quant_result_dir=str(out)) | ||
| 369 | + assert model._forward_hooks, "a publish hook must be armed while incomplete" | ||
| 370 | + | ||
| 371 | + inputs = _tiny_inputs() | ||
| 372 | + model(inputs) | ||
| 373 | + model(inputs) | ||
| 374 | + assert not model._forward_hooks, "hook must detach once published" | ||
| 375 | + # Further inference must neither fail nor touch the published artifact. | ||
| 376 | + digest_before = (out / MANIFEST_NAME).read_bytes() | ||
| 377 | + model(inputs) | ||
| 378 | + assert (out / MANIFEST_NAME).read_bytes() == digest_before | ||
| 379 | + | ||
| 380 | + | ||
| 381 | +def test_failed_calibration_forward_leaves_no_manifest(tmp_path): | ||
| 382 | + out = tmp_path / "failed" | ||
| 383 | + out.mkdir() | ||
| 384 | + model = _tiny_model() | ||
| 385 | + quantize(model, ACTIVATION_CFG, quant_result_dir=str(out)) | ||
| 386 | + with pytest.raises(RuntimeError): | ||
| 387 | + model(torch.randn(8, 999).to(torch.bfloat16)) # wrong in_features | ||
| 388 | + assert not (out / MANIFEST_NAME).exists() | ||
| 389 | + | ||
| 390 | + | ||
| 391 | +def test_one_failed_unit_write_leaves_no_manifest(tmp_path): | ||
| 392 | + out = tmp_path / "partial" | ||
| 393 | + out.mkdir() | ||
| 394 | + model = _tiny_model() | ||
| 395 | + real_save = QUANT_RESULT_MODULE.save_file | ||
| 396 | + calls = {"n": 0} | ||
| 397 | + | ||
| 398 | + def fail_on_second(*args, **kwargs): | ||
| 399 | + calls["n"] += 1 | ||
| 400 | + if calls["n"] == 2: | ||
| 401 | + raise OSError("disk full") | ||
| 402 | + return real_save(*args, **kwargs) | ||
| 403 | + | ||
| 404 | + with patch.object(QUANT_RESULT_MODULE, "save_file", fail_on_second): | ||
| 405 | + with pytest.raises(OSError): | ||
| 406 | + quantize(model, WEIGHT_ONLY_CFG, quant_result_dir=str(out)) | ||
| 407 | + | ||
| 408 | + names = sorted(p.name for p in out.iterdir()) | ||
| 409 | + assert MANIFEST_NAME not in names, "an incomplete run must not publish" | ||
| 410 | + assert names == ["fc1.safetensors"], "no temp file may survive the failure" | ||
| 411 | + with pytest.raises(FileNotFoundError, match="not a valid quant result"): | ||
| 412 | + load_quant_manifest(str(out)) | ||
| 413 | + | ||
| 414 | + | ||
| 415 | +def test_quantize_records_source_model_fingerprint(tmp_path): | ||
| 416 | + out = tmp_path / "fp" | ||
| 417 | + out.mkdir() | ||
| 418 | + model = _tiny_model() | ||
| 419 | + expected = compute_model_fingerprint(_tiny_model_like(model)) | ||
| 420 | + quantize(model, WEIGHT_ONLY_CFG, quant_result_dir=str(out)) | ||
| 421 | + manifest = load_quant_manifest(str(out)) | ||
| 422 | + # Fingerprint is taken before MinMaxQuant replaces the Linear modules, so | ||
| 423 | + # it describes the source model, not the rewritten one. | ||
| 424 | + assert manifest["source_model"] == expected | ||
| 425 | + assert manifest["source_model"]["model_type"] == "_TinyModel" | ||
| 426 | + | ||
| 427 | + | ||
| 428 | +def _tiny_model_like(model): | ||
| 429 | + """A pristine copy of the pre-quantize model, for fingerprint comparison.""" | ||
| 430 | + clone = _tiny_model() | ||
| 431 | + clone.load_state_dict( | ||
| 432 | + { | ||
| 433 | + "fc1.weight": model.fc1.weight.detach().clone(), | ||
| 434 | + "fc1.bias": model.fc1.bias.detach().clone(), | ||
| 435 | + "fc2.weight": model.fc2.weight.detach().clone(), | ||
| 436 | + "fc2.bias": model.fc2.bias.detach().clone(), | ||
| 437 | + } | ||
| 438 | + ) | ||
| 439 | + return clone | ||
| 440 | + | ||
| 441 | + | ||
| 442 | +AWQ_CFG = { | ||
| 443 | + 'batch_num': 1, | ||
| 444 | + 'quant_cfg': { | ||
| 445 | + 'weights': { | ||
| 446 | + 'type': 'int4', | ||
| 447 | + 'symmetric': True, | ||
| 448 | + 'strategy': 'channel', | ||
| 449 | + } | ||
| 450 | + }, | ||
| 451 | + 'algorithm': {'awq': {'grids_num': 20}}, | ||
| 452 | +} | ||
| 453 | + | ||
| 454 | + | ||
| 455 | +SMOOTHQUANT_CFG = { | ||
| 456 | + 'batch_num': 1, | ||
| 457 | + 'quant_cfg': { | ||
| 458 | + 'weights': {'type': 'int8', 'symmetric': True, 'strategy': 'channel'}, | ||
| 459 | + 'inputs': {'type': 'int8', 'symmetric': True, 'strategy': 'tensor'}, | ||
| 460 | + }, | ||
| 461 | + 'algorithm': {'smoothquant': {'smooth_strength': 0.5}}, | ||
| 462 | +} | ||
| 463 | + | ||
| 464 | + | ||
| 465 | +class _LinearShapedModel(nn.Module): | ||
| 466 | + def __init__(self): | ||
| 467 | + super().__init__() | ||
| 468 | + self.fc1 = nn.Linear(64, 64, bias=False) | ||
| 469 | + | ||
| 470 | + def forward(self, inputs): | ||
| 471 | + return self.fc1(inputs) | ||
| 472 | + | ||
| 473 | + | ||
| 474 | +def test_quantize_rejects_an_algorithm_without_export_support(tmp_path): | ||
| 475 | + """An algorithm that cannot dump must fail before writing any result. | ||
| 476 | + | ||
| 477 | + Regression guard: every quantize module inherits BaseQuantizeModule, so a | ||
| 478 | + mere isinstance or hasattr check lets an unwired algorithm through. | ||
| 479 | + quantize() then returned normally, the result directory stayed empty, and the | ||
| 480 | + manifest was pending forever -- an API that reports success while producing | ||
| 481 | + nothing loadable. SmoothQuant stands in for "not wired yet" now that AWQ and | ||
| 482 | + GPTQ are. | ||
| 483 | + """ | ||
| 484 | + out = tmp_path / "smoothquant" | ||
| 485 | + out.mkdir() | ||
| 486 | + model = _LinearShapedModel().to(torch.bfloat16) | ||
| 487 | + with pytest.raises(ValueError, match="not supported yet for algorithm") as ctx: | ||
| 488 | + quantize(model, SMOOTHQUANT_CFG, quant_result_dir=str(out)) | ||
| 489 | + | ||
| 490 | + assert "_supports_quant_result_export" not in str(ctx.value) | ||
| 491 | + assert list(out.iterdir()) == [], "a rejected algorithm must write nothing" | ||
| 492 | + assert type(model.fc1).__name__ == "Linear", "the model must be left untouched" | ||
| 493 | + | ||
| 494 | + | ||
| 495 | +def test_the_same_unsupported_algorithm_still_works_without_a_result_dir(tmp_path): | ||
| 496 | + model = _LinearShapedModel().to(torch.bfloat16) | ||
| 497 | + quantize(model, SMOOTHQUANT_CFG) | ||
| 498 | + # The legacy in-memory path is unaffected by the export screening. | ||
| 499 | + assert type(model.fc1).__name__ == "SmoothQuant" | ||
| 500 | + | ||
| 501 | + | ||
| 502 | +def test_awq_is_now_accepted_for_export(tmp_path): | ||
| 503 | + """AWQ is wired in Week 4: it must publish a real artifact.""" | ||
| 504 | + out = tmp_path / "awq_ok" | ||
| 505 | + out.mkdir() | ||
| 506 | + model = _LinearShapedModel().to(torch.bfloat16) | ||
| 507 | + quantize(model, AWQ_CFG, quant_result_dir=str(out)) | ||
| 508 | + model(torch.randn(8, 64).to(torch.bfloat16)) | ||
| 509 | + assert (out / "quant_manifest.json").is_file() | ||
| 510 | + | ||
| 511 | + | ||
| 512 | +def test_artifact_is_consumable_by_the_reload_contract(tmp_path): | ||
| 513 | + """Walk the spec's reload validation chain and rebuild from the manifest. | ||
| 514 | + | ||
| 515 | + Week 4 owns ``reload_quant_params()``; this asserts the Week 3 artifact | ||
| 516 | + already carries everything that chain needs, so the schema cannot drift out | ||
| 517 | + from under it. The rebuild uses only manifest fields -- no knowledge of how | ||
| 518 | + the artifact was produced. | ||
| 519 | + """ | ||
| 520 | + out = tmp_path / "consume" | ||
| 521 | + out.mkdir() | ||
| 522 | + model = _tiny_model() | ||
| 523 | + reference = _tiny_model_like(model) | ||
| 524 | + quantize(model, ACTIVATION_CFG, quant_result_dir=str(out)) | ||
| 525 | + inputs = _tiny_inputs() | ||
| 526 | + model(inputs) | ||
| 527 | + model(inputs) | ||
| 528 | + expected = model(inputs) | ||
| 529 | + | ||
| 530 | + # 1. Manifest format, version, required fields, uniqueness, safe paths. | ||
| 531 | + manifest = load_quant_manifest(str(out)) | ||
| 532 | + # 2. The directory holds the manifest and every unit it manages. | ||
| 533 | + validate_result_directory(str(out), manifest) | ||
| 534 | + # 3. The source model identity matches the untouched float model. | ||
| 535 | + assert manifest["source_model"] == compute_model_fingerprint(reference) | ||
| 536 | + | ||
| 537 | + helper = ModuleHelper(reference) | ||
| 538 | + for entry in manifest["units"]: | ||
| 539 | + layer_name = entry["layer_name"] | ||
| 540 | + assert layer_name in helper.named_module_dict | ||
| 541 | + # 4-6. Size, digest, header and tensor schema, all checked by the reader. | ||
| 542 | + tensors, non_tensor_params, _ = read_quant_unit(str(out), layer_name, manifest) | ||
| 543 | + | ||
| 544 | + # A Week 4 loader resolves the class from quant_module_type; MinMaxQuant | ||
| 545 | + # is the only registered name here, so assert the recorded value and use | ||
| 546 | + # it rather than hardcoding the class at the call site. | ||
| 547 | + assert entry["quant_module_type"] == "MinMaxQuant" | ||
| 548 | + assert entry["ori_module_type"] == "Linear" | ||
| 549 | + ori_module = helper.named_module_dict[layer_name] | ||
| 550 | + rebuilt = MinMaxQuant(ori_module, layer_name, entry["quant_config"]) | ||
| 551 | + rebuilt.load_quant_params(tensors, non_tensor_params) | ||
| 552 | + ModuleHelper.replace_module_by_name(reference, layer_name, rebuilt) | ||
| 553 | + | ||
| 554 | + # The rebuilt model fake-quants identically, without recalibrating. | ||
| 555 | + assert torch.allclose(reference(inputs), expected) | ||
| 556 | + | ||
| 557 | + | ||
| 558 | +def _npu_available(): | ||
| 559 | + """Return whether torch exposes a usable, non-mocked NPU backend.""" | ||
| 560 | + if not hasattr(torch, "npu"): | ||
| 561 | + return False | ||
| 562 | + try: | ||
| 563 | + if torch.npu.is_available() is not True: | ||
| 564 | + return False | ||
| 565 | + torch.zeros(1).npu() | ||
| 566 | + except Exception: | ||
| 567 | + return False | ||
| 568 | + return True | ||
| 569 | + | ||
| 570 | + | ||
| 571 | +def test_npu_available_rejects_mocked_backend(): | ||
| 572 | + with patch.object(torch, "npu", MagicMock(), create=True): | ||
| 573 | + assert not _npu_available() | ||
| 574 | + | ||
| 575 | + | ||
| 576 | + | ||
| 577 | +def test_npu_round_trip_restores_params_to_the_device(tmp_path): | ||
| 578 | + """On real NPU hardware, units persist as CPU tensors and reload to npu. | ||
| 579 | + | ||
| 580 | + Units are always stored CPU-side (device is not part of the artifact), so | ||
| 581 | + this is the check that ``_quant_param_device()`` puts them back on the | ||
| 582 | + module's device rather than leaving a CPU tensor in an NPU module. | ||
| 583 | + """ | ||
| 584 | + if not _npu_available(): | ||
| 585 | + pytest.skip("no Ascend NPU available") | ||
| 586 | + | ||
| 587 | + out = tmp_path / "npu" | ||
| 588 | + out.mkdir() | ||
| 589 | + model = _tiny_model().npu() | ||
| 590 | + reference = _tiny_model_like(model).npu() | ||
| 591 | + quantize(model, ACTIVATION_CFG, quant_result_dir=str(out)) | ||
| 592 | + | ||
| 593 | + inputs = _tiny_inputs().npu() | ||
| 594 | + model(inputs) | ||
| 595 | + model(inputs) | ||
| 596 | + expected = model(inputs) | ||
| 597 | + | ||
| 598 | + manifest = load_quant_manifest(str(out)) | ||
| 599 | + validate_result_directory(str(out), manifest) | ||
| 600 | + helper = ModuleHelper(reference) | ||
| 601 | + for entry in manifest["units"]: | ||
| 602 | + layer_name = entry["layer_name"] | ||
| 603 | + tensors, non_tensor_params, _ = read_quant_unit(str(out), layer_name, manifest) | ||
| 604 | + assert all(t.device.type == "cpu" for t in tensors.values()) | ||
| 605 | + rebuilt = MinMaxQuant( | ||
| 606 | + helper.named_module_dict[layer_name], layer_name, entry["quant_config"] | ||
| 607 | + ) | ||
| 608 | + rebuilt.load_quant_params(tensors, non_tensor_params) | ||
| 609 | + assert rebuilt.scale_w.device == rebuilt.weight.device | ||
| 610 | + assert rebuilt.scale_d.device == rebuilt.weight.device | ||
| 611 | + ModuleHelper.replace_module_by_name(reference, layer_name, rebuilt) | ||
| 612 | + | ||
| 613 | + assert torch.equal(reference(inputs), expected) | ||
| 614 | + | ||
| 615 | + | ||
| 616 | +def test_pass_attaches_the_writer_without_exposing_capability_flags(): | ||
| 617 | + class _RecordingQuantOp(nn.Module): | ||
| 618 | + def __init__(self, ori_module, layer_name, quant_config): | ||
| 619 | + super().__init__() | ||
| 620 | + self.ori_module = ori_module | ||
| 621 | + self.attached_writer = None | ||
| 622 | + | ||
| 623 | + def attach_quant_result_writer(self, writer): | ||
| 624 | + self.attached_writer = writer | ||
| 625 | + | ||
| 626 | + def forward(self, inputs): | ||
| 627 | + return self.ori_module(inputs) | ||
| 628 | + | ||
| 629 | + model = _tiny_model() | ||
| 630 | + writer = MagicMock() | ||
| 631 | + quant_pass = InsertQuantizeModulePass( | ||
| 632 | + {"fc1": {"algorithm": {"minmax": {}}}}, quant_result_writer=writer | ||
| 633 | + ) | ||
| 634 | + quant_pass.quantize_ops["fc1"] = _RecordingQuantOp | ||
| 635 | + quant_pass.do_pass(model, model.fc1, "fc1") | ||
| 636 | + assert model.fc1.attached_writer is writer | ||
| @@ -137,12 +137,37 @@ class _KvDeployTestMixin(unittest.TestCase): | |||
| 137 | self.assertIs(cache.update, original_update) | 137 | self.assertIs(cache.update, original_update) |
| 138 | 138 | ||
| 139 | 139 | ||
| 140 | -class TestDeepseekV3AttentionDeploy(_KvDeployTestMixin): | 140 | +class _FakeNpuImportMixin(_KvDeployTestMixin): |
| 141 | + """Installs the stub only for the import, then puts sys.modules back. | ||
| 142 | + | ||
| 143 | + The stub has to exist while the deploy module is imported, but it must not | ||
| 144 | + outlive this class: it lacks most of the real torch_npu surface, so leaking it | ||
| 145 | + breaks any later test whose code path reaches torch_npu (GPTQ's | ||
| 146 | + ``torch_npu.npu.is_available()``, for one). | ||
| 147 | + """ | ||
| 148 | + | ||
| 149 | + __test__ = False | ||
| 150 | + | ||
| 151 | + | ||
| 152 | + def _install_stub_for_import(cls): | ||
| 153 | + cls._saved_module = sys.modules.get('torch_npu') | ||
| 154 | + _install_fake_torch_npu() | ||
| 155 | + | ||
| 156 | + | ||
| 157 | + def tearDownClass(cls): | ||
| 158 | + saved = getattr(cls, '_saved_module', None) | ||
| 159 | + if saved is not None: | ||
| 160 | + sys.modules['torch_npu'] = saved | ||
| 161 | + else: | ||
| 162 | + sys.modules.pop('torch_npu', None) | ||
| 163 | + | ||
| 164 | + | ||
| 165 | +class TestDeepseekV3AttentionDeploy(_FakeNpuImportMixin): | ||
| 141 | __test__ = True | 166 | __test__ = True |
| 142 | 167 | ||
| 143 | 168 | ||
| 144 | def setUpClass(cls): | 169 | def setUpClass(cls): |
| 145 | - _install_fake_torch_npu() | 170 | + cls._install_stub_for_import() |
| 146 | from amct_pytorch.classic.deploy_op.npu_quantization_deepseekv3_attention import ( | 171 | from amct_pytorch.classic.deploy_op.npu_quantization_deepseekv3_attention import ( |
| 147 | NpuDeepseekV3AttentionQuant, | 172 | NpuDeepseekV3AttentionQuant, |
| 148 | ) | 173 | ) |
| @@ -150,12 +175,12 @@ class TestDeepseekV3AttentionDeploy(_KvDeployTestMixin): | |||
| 150 | cls.module_cls = NpuDeepseekV3AttentionQuant | 175 | cls.module_cls = NpuDeepseekV3AttentionQuant |
| 151 | 176 | ||
| 152 | 177 | ||
| 153 | -class TestLongcatFlashMLADeploy(_KvDeployTestMixin): | 178 | +class TestLongcatFlashMLADeploy(_FakeNpuImportMixin): |
| 154 | __test__ = True | 179 | __test__ = True |
| 155 | 180 | ||
| 156 | 181 | ||
| 157 | def setUpClass(cls): | 182 | def setUpClass(cls): |
| 158 | - _install_fake_torch_npu() | 183 | + cls._install_stub_for_import() |
| 159 | from amct_pytorch.classic.deploy_op.npu_quantization_longcat_flashmla import ( | 184 | from amct_pytorch.classic.deploy_op.npu_quantization_longcat_flashmla import ( |
| 160 | NpuLongcatFlashMLA, | 185 | NpuLongcatFlashMLA, |
| 161 | ) | 186 | ) |
校验太严格，建议只校验model type / torch_dtype等影响量化的参数