Generator that yields shards of weights loaded from the model files in huggingface format. Each shard contains a dictionary of tensors, with weights normalized and cast to the specified dtype (except layer norm weights which are converted to float32).
(model_files: List[str], metadata: Metadata, tie_word_embeddings: bool, n_layers: int)
| 294 | return torch.stack(output_weights) |
| 295 | |
| 296 | def load_weights(model_files: List[str], metadata: Metadata, tie_word_embeddings: bool, n_layers: int): |
| 297 | """ |
| 298 | Generator that yields shards of weights loaded from the model files in huggingface format. |
| 299 | Each shard contains a dictionary of tensors, with weights normalized and cast to the specified dtype |
| 300 | (except layer norm weights which are converted to float32). |
| 301 | """ |
| 302 | weights = {} |
| 303 | for model_path in model_files: |
| 304 | ext = os.path.splitext(model_path)[1] |
| 305 | if ext == ".safetensors": |
| 306 | with safetensors.safe_open(model_path, framework="pt") as f: |
| 307 | for k in f.keys(): |
| 308 | assert(k not in weights) |
| 309 | weights[k] = f.get_tensor(k) |
| 310 | dtype = metadata.quant.dtype |
| 311 | |
| 312 | # convert weights |
| 313 | progress = 0 |
| 314 | dequant_block_size = None |
| 315 | if metadata.original_quantization_config is not None: |
| 316 | dequant_block_size = torch.tensor(metadata.original_quantization_config["weight_block_size"]) |
| 317 | tensors = {} |
| 318 | |
| 319 | def load_and_dequantize(weight_name: str, scale_name: str) -> torch.Tensor: |
| 320 | t = weights[weight_name] |
| 321 | if scale_name in weights: |
| 322 | scale = weights[scale_name] |
| 323 | t = blockwise_dequantize(t, scale, dequant_block_size) |
| 324 | return t |
| 325 | |
| 326 | def quantize(t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| 327 | if dtype not in [torch.float32, torch.float16]: |
| 328 | if isinstance(metadata.quant, KQuant): |
| 329 | t = k_quantize(t.to(torch.float32), metadata.quant.name) |
| 330 | elif metadata.quant.block_size is None: |
| 331 | return per_tensor_quantize(t, dtype) |
| 332 | else: |
| 333 | quant_block_size = torch.tensor(metadata.quant.block_size) |
| 334 | return blockwise_quantize(t, quant_block_size, dtype) |
| 335 | return t.to(dtype), None |
| 336 | |
| 337 | def conv(weight_name: str, scale_name: str) -> Tuple[torch.Tensor, torch.Tensor]: |
| 338 | nonlocal progress |
| 339 | progress += 1 |
| 340 | t = load_and_dequantize(weight_name, scale_name) |
| 341 | print(f"\rConverting tensor {progress}: {t.shape}", end="", flush=True) |
| 342 | return quantize(t) |
| 343 | |
| 344 | def conv_experts(weight_and_scale_names: List[Tuple[str, str]]) -> Tuple[torch.Tensor, torch.Tensor]: |
| 345 | nonlocal progress |
| 346 | progress += 1 |
| 347 | expert_weights = [weights[weight_name] for weight_name, _ in weight_and_scale_names] |
| 348 | if weight_and_scale_names[0][1] in weights: |
| 349 | for i in range(len(weight_and_scale_names)): |
| 350 | scale = weights[weight_and_scale_names[i][1]] |
| 351 | expert_weights[i] = blockwise_dequantize(expert_weights[i], scale, dequant_block_size) |
| 352 | t = torch.stack(expert_weights) |
| 353 | print(f"\rConverting tensor {progress}: {t.shape}", end="", flush=True) |
no test coverage detected