| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Per-parameter weight digests and reset, for verifying weight updates.""" |
| 4 | |
| 5 | import hashlib |
| 6 | from collections.abc import Iterator |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | from torch.utils._python_dispatch import is_traceable_wrapper_subclass |
| 11 | |
| 12 | from vllm.model_executor.parameter import SharedWeightParameter |
| 13 | |
| 14 | |
| 15 | def _tensors(name: str, tensor: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]: |
| 16 | # Both keep their bytes in inner tensors (e.g. TorchAO's qdata and scale). |
| 17 | if isinstance(tensor, SharedWeightParameter): |
| 18 | for index, partition in tensor.partitions.items(): |
| 19 | yield from _tensors(f"{name}.{index}", partition) |
| 20 | elif is_traceable_wrapper_subclass(tensor): |
| 21 | for attr in tensor.__tensor_flatten__()[0]: |
| 22 | yield from _tensors(f"{name}.{attr}", getattr(tensor, attr)) |
| 23 | else: |
| 24 | yield name, tensor |
| 25 | |
| 26 | |
| 27 | def _weights(model: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: |
| 28 | for name, param in model.named_parameters(): |
| 29 | yield from _tensors(name, param) |
| 30 | |
| 31 | |
| 32 | def compute_tensor_digests(model: nn.Module) -> dict[str, str]: |
| 33 | return { |
| 34 | name: hashlib.sha256( |
| 35 | weight.detach().cpu().contiguous().view(-1).view(torch.uint8).numpy() |
| 36 | ).hexdigest() |
| 37 | for name, weight in _weights(model) |
| 38 | } |
| 39 | |
| 40 | |
| 41 | def zero_weights(model: nn.Module) -> None: |
| 42 | for _, weight in _weights(model): |
| 43 | weight.data.zero_() |