Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / quantized_hf_storage.py: 92%
194 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Read FineGrained FP8/FP4 HF storage through shared DCP region reads.
17The planner sees logical, unpacked shapes. Each physical read decodes only the
18requested region, expanded to quantization block boundaries. BF16 rounding is
19performed before any model/replacement conversion. Decoded values, including
20the sign of zero, follow the installed Transformers implementation.
21"""
23import json
24from dataclasses import dataclass, replace
25from pathlib import Path
26from typing import Any, Optional, Union
28import torch
29from safetensors import safe_open
30from torch import nn
32from hyper_parallel.components.checkpoint.hf_storage import HuggingFaceStorageReader
33from hyper_parallel.core.distributed_checkpoint.metadata import (
34 ChunkStorageMetadata, Metadata, MetadataIndex, TensorProperties, TensorStorageMetadata,
35)
37_SCRATCH_BYTES = 32 * 1024 * 1024
40def supports_hf_dequantization(config: Any) -> bool:
41 """Return whether shared loading supports the declared method.
43 Args:
44 config: Quantization dictionary or Transformers quantization configuration.
46 Returns:
47 Whether the shared reader supports the quantization method.
48 """
49 method = config.get("quant_method") if isinstance(config, dict) else getattr(config, "quant_method", None)
50 return method == "fp8"
53def validate_quantization_config(config: dict[str, Any]) -> None:
54 """Validate supported storage formats without depending on a model family.
56 Args:
57 config: Quantization configuration from the checkpoint.
59 Raises:
60 ValueError: The method, block dimensions or scale format is unsupported.
61 """
62 if not supports_hf_dequantization(config):
63 raise ValueError("Unsupported HF quantization method for online DCP decoding; expected fp8")
64 block = config.get("weight_block_size") or [128, 128]
65 if (not isinstance(block, (list, tuple)) or len(block) != 2
66 or any(not isinstance(size, int) or size <= 0 for size in block)):
67 raise ValueError("weight_block_size must contain two positive integers")
68 if config.get("scale_fmt", "float") not in ("float", "ue8m0"):
69 raise ValueError("Unsupported FP8 scale format")
72def _validate_scales(scale: torch.Tensor) -> None:
73 if scale.dtype in (torch.float8_e8m0fnu, torch.uint8):
74 codes = scale.contiguous().view(torch.uint8)
75 if torch.any(codes == 255):
76 raise ValueError("UE8M0 scale contains reserved NaN code 255")
77 elif scale.dtype != torch.float32:
78 raise ValueError(f"Unsupported scale dtype: {scale.dtype}")
79 result = scale.view(torch.float8_e8m0fnu).float() if scale.dtype == torch.uint8 else scale.float()
80 if not torch.all(torch.isfinite(result) & (result > 0)):
81 raise ValueError("Quantization scales must be positive and finite")
84@dataclass(frozen=True)
85class _EncodedWeight:
86 scale: str
87 shape: tuple[int, int]
88 block_shape: tuple[int, int]
89 packed_elements: int
92class QuantizedHuggingFaceStorageReader(HuggingFaceStorageReader):
93 """Decode two-dimensional FineGrained FP8/FP4 weights using Transformers.
95 Args:
96 checkpoint_dir: Local HF directory containing a supported quantization_config, weights and scales.
97 """
99 def __init__(self, checkpoint_dir: Union[Path, str]) -> None:
100 """Initialize metadata caches without reading weight payloads."""
101 super().__init__(checkpoint_dir)
102 config = json.loads((self.checkpoint_dir / "config.json").read_text(encoding="utf-8"))
103 self.quantization_config = config.get("quantization_config", {})
104 validate_quantization_config(self.quantization_config)
105 # Keep optional/version-dependent HF quantization imports off ordinary HF loading.
106 try:
107 from transformers.integrations.finegrained_fp8 import Fp8Dequantize # pylint: disable=C0415
108 from transformers.quantizers.quantizer_finegrained_fp8 import ( # pylint: disable=C0415
109 FineGrainedFP8HfQuantizer,
110 )
111 from transformers.utils.quantization_config import FineGrainedFP8Config # pylint: disable=C0415
112 except ImportError as exc:
113 raise ValueError("Online HF decoding requires Transformers FineGrainedFP8/Fp8Dequantize support") from exc
114 quantizer = FineGrainedFP8HfQuantizer(FineGrainedFP8Config(
115 dequantize=True, weight_block_size=tuple(self.quantization_config.get("weight_block_size") or [128, 128]),
116 scale_fmt=self.quantization_config.get("scale_fmt", "float"),
117 ))
118 self._dequantizer = Fp8Dequantize(quantizer)
119 self._dtype_target = nn.Module()
120 self._dtype_target.register_parameter(
121 "weight", nn.Parameter(torch.empty(0, device="meta", dtype=torch.bfloat16)))
122 self._physical: Optional[Metadata] = None
123 self._logical: Optional[Metadata] = None
124 self._encoded: dict[str, _EncodedWeight] = {}
126 def load_metadata(self, **kwargs: Any) -> Metadata:
127 """Read headers and expose unpacked shapes, hiding consumed scales."""
128 if self._logical is not None:
129 return self._logical
130 physical = super().load_metadata(**kwargs)
131 entries = physical.state_dict_metadata
132 index_path = self.checkpoint_dir / "model.safetensors.index.json"
133 if index_path.is_file():
134 selected = json.loads(index_path.read_text(encoding="utf-8"))["weight_map"]
135 missing = set(selected) - entries.keys()
136 if missing:
137 raise ValueError(f"HF index references missing tensors: {sorted(missing)[:10]}")
138 entries = {name: entries[name] for name in selected}
139 logical = {}
140 used_scales = set()
141 for name, entry in entries.items():
142 encoded = self._describe_weight(name, entry, entries)
143 if encoded is None:
144 logical[name] = entry
145 continue
146 self._encoded[name] = encoded
147 used_scales.add(encoded.scale)
148 logical[name] = TensorStorageMetadata(
149 properties=TensorProperties(dtype="torch.bfloat16"), size=encoded.shape,
150 chunks=[ChunkStorageMetadata(offsets=(0, 0), sizes=encoded.shape)],
151 )
152 for name in used_scales:
153 logical.pop(name, None)
154 orphan_scales = [name for name, entry in logical.items()
155 if str(entry.properties.dtype) == "torch.float8_e8m0fnu"]
156 if orphan_scales:
157 raise ValueError(f"Unconsumed quantization scales: {orphan_scales[:10]}")
158 self._physical = physical
159 self._logical = replace(physical, state_dict_metadata=logical)
160 return self._logical
162 def _describe_weight(self, name: str, entry: TensorStorageMetadata,
163 entries: dict[str, Any]) -> Optional[_EncodedWeight]:
164 """Resolve a quantized weight's scale, logical shape, and block layout."""
165 dtype = str(entry.properties.dtype)
166 if not name.endswith("weight"):
167 return None
168 candidates = [name[:-6] + suffix for suffix in ("scale", "weight_scale_inv")
169 if name[:-6] + suffix in entries]
170 if len(candidates) > 1:
171 raise ValueError(f"Ambiguous quantization scales: {name}")
172 packed = 2 if dtype in ("torch.int8", "torch.uint8") else 1
173 if dtype not in ("torch.float8_e4m3fn", "torch.int8", "torch.uint8"):
174 if candidates or dtype.startswith("torch.float8"):
175 raise ValueError(f"Unsupported encoded weight: {name} ({dtype})")
176 return None
177 if not candidates or len(entry.size) != 2:
178 raise ValueError(f"Missing scale or unsupported encoded shape: {name}")
179 rows, columns = entry.size
180 columns *= packed
181 scale = entries[candidates[0]]
182 block_shape = self._scale_block_shape(name, (rows, columns), scale, packed)
183 return _EncodedWeight(candidates[0], (rows, columns), block_shape, packed)
185 def _scale_block_shape(self, name: str, shape: tuple[int, int],
186 scale: TensorStorageMetadata, packed: int) -> tuple[int, int]:
187 """Validate a scale grid and derive its logical block dimensions."""
188 if (len(scale.size) != 2 or min(scale.size) <= 0
189 or str(scale.properties.dtype) not in ("torch.float8_e8m0fnu", "torch.float32", "torch.uint8")):
190 raise ValueError(f"Invalid quantization scale shape or dtype: {name}")
191 rows, columns = shape
192 block_rows, block_columns = self.quantization_config.get("weight_block_size") or [128, 128]
193 expected = ((rows + block_rows - 1) // block_rows, (columns + block_columns - 1) // block_columns)
194 if packed == 2:
195 # FP4 packing and scale-grid granularity are properties of storage,
196 # not of expert/module names. The same layout can belong to any model.
197 if rows % scale.size[0] or columns % scale.size[1]:
198 raise ValueError(f"FP4 weight is not divisible by its scale grid: {name}")
199 block_rows, block_columns = rows // scale.size[0], columns // scale.size[1]
200 if block_columns % packed:
201 raise ValueError(f"FP4 scale blocks must align with packed byte pairs: {name}")
202 elif scale.size != expected:
203 # Row-wise MX-style scaling can coexist with block-wise FP8 scaling.
204 if scale.size != (rows, (columns + 31) // 32):
205 raise ValueError(f"Invalid quantization scale shape: {name}")
206 block_rows, block_columns = 1, 32
207 return block_rows, block_columns
209 def _physical_slice(self, name: str, region: tuple[slice, slice], open_files: Any = None) -> torch.Tensor:
210 info = self._physical.storage_data[MetadataIndex(name, (0, 0))]
211 path = str(self.checkpoint_dir / info.relative_path)
212 if open_files is not None:
213 return open_files.reader(path).get_slice(name)[region]
214 with safe_open(path, framework="pt", device="cpu") as source:
215 return source.get_slice(name)[region]
217 def read_region(self, name: str, offsets: tuple[int, int], lengths: tuple[int, int],
218 open_files: Any = None) -> torch.Tensor:
219 """Read a logical rectangle, with bounded decoding scratch on CPU.
221 Args:
222 name: Original HF weight name.
223 offsets: Logical row and column starts, before FP4 packing.
224 lengths: Logical rectangle shape.
225 open_files: Optional file cache supplied by the DCP reader.
227 Returns:
228 BF16 rectangle with the same rounding as offline conversion; ordinary
229 floating-point tensors retain their stored dtype.
230 """
231 metadata = self.load_metadata()
232 shape = metadata.state_dict_metadata[name].size
233 if (len(offsets) != 2 or len(lengths) != 2
234 or len(shape) != 2
235 or any(start < 0 or length < 0 or start + length > size
236 for start, length, size in zip(offsets, lengths, shape))):
237 raise ValueError(f"Invalid logical read region for {name}: {offsets}, {lengths}")
238 encoded = self._encoded.get(name)
239 if encoded is None:
240 region = tuple(slice(start, start + length) for start, length in zip(offsets, lengths))
241 return self._physical_slice(name, region, open_files)
242 first_row, first_column = offsets
243 row_count, column_count = lengths
244 result = torch.empty(lengths, dtype=torch.bfloat16)
245 if not row_count or not column_count:
246 return result
247 row_block, column_block = encoded.block_shape
248 left = first_column // column_block * column_block
249 right = min((first_column + column_count + column_block - 1) // column_block * column_block,
250 encoded.shape[1])
251 rows_per_chunk = max(row_block, _SCRATCH_BYTES // max(1, (right - left) * 32))
252 rows_per_chunk = rows_per_chunk // row_block * row_block
253 cursor = first_row
254 while cursor < first_row + row_count:
255 start = cursor // row_block * row_block
256 stop = min(start + rows_per_chunk, encoded.shape[0], first_row + row_count)
257 # Read complete quantization blocks even for unaligned end slices.
258 read_stop = min((stop + row_block - 1) // row_block * row_block, encoded.shape[0])
259 value = self._decode_chunk(name, encoded, start, read_stop, left, right, open_files)
260 result[cursor - first_row:stop - first_row].copy_(
261 value[cursor - start:stop - start, first_column - left:first_column - left + column_count])
262 cursor = stop
263 return result
265 def _decode_chunk(self, name: str, encoded: _EncodedWeight, start: int, stop: int,
266 left: int, right: int, open_files: Any) -> torch.Tensor:
267 """Read aligned weight and scale blocks and decode the requested BF16 rectangle."""
268 packed = encoded.packed_elements
269 weight = self._physical_slice(name, (slice(start, stop), slice(left // packed, (right + packed - 1) // packed)),
270 open_files)
271 row_block, column_block = encoded.block_shape
272 scale = self._physical_slice(encoded.scale,
273 (slice(start // row_block, (stop + row_block - 1) // row_block),
274 slice(left // column_block, (right + column_block - 1) // column_block)),
275 open_files)
276 _validate_scales(scale)
277 if packed == 2 and weight.dtype == torch.uint8:
278 weight = weight.view(torch.int8)
279 padded_shape = (scale.shape[0] * row_block, scale.shape[1] * column_block // packed)
280 if weight.shape != padded_shape:
281 # HF infers block sizes from the scale grid. Pad only the global edge
282 # so a partial last block is not mistaken for a different block size.
283 padded = torch.zeros(padded_shape, dtype=weight.dtype)
284 padded[:weight.shape[0], :weight.shape[1]].copy_(weight)
285 weight = padded
286 value = self._dequantizer.convert(
287 {"weight$": weight, "weight_scale_inv": scale},
288 full_layer_name="weight", model=self._dtype_target,
289 )["weight"][:stop - start, :right - left]
290 if value.shape != (stop - start, right - left) or value.dtype != torch.bfloat16:
291 raise ValueError("Transformers Fp8Dequantize did not produce the requested BF16 rectangle")
292 if not torch.all(torch.isfinite(value)):
293 raise ValueError(f"Nonfinite decoded value or BF16 overflow: {name}")
294 return value
296 def _fetch_from_storage(self, items: list, open_files: Any) -> list:
297 ordinary, encoded = [], []
298 for item in items:
299 (encoded if item.storage_index.fqn in self._encoded else ordinary).append(item)
300 fetched = super()._fetch_from_storage(ordinary, open_files)
301 for item in encoded:
302 fetched.append(
303 (
304 item,
305 self.read_region(
306 item.storage_index.fqn,
307 item.storage_offsets,
308 item.lengths,
309 open_files
310 )
311 )
312 )
313 return fetched
316def build_quantized_hf_reader(location: Union[Path, str]) -> Optional[QuantizedHuggingFaceStorageReader]:
317 """Select shared decoding only for a directory declaring supported quantization.
319 Args:
320 location: Resolved local checkpoint directory or file.
322 Returns:
323 Quantized reader, or None for an ordinary floating-point checkpoint.
324 """
325 root = Path(location)
326 config_path = root / "config.json"
327 if not config_path.is_file():
328 return None
329 config = json.loads(config_path.read_text(encoding="utf-8"))
330 quantization = config.get("quantization_config")
331 if not quantization:
332 return None
333 return QuantizedHuggingFaceStorageReader(root)