| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import MethodType |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | __all__ = ["sanitize_layer_refs", "restore_layer_refs"] |
| 8 | |
| 9 | |
| 10 | layer_ref_sentinel = object() |
| 11 | |
| 12 | |
| 13 | def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor: |
| 14 | """Removes references to layer held by tensor attributes. Specifically, removes the |
| 15 | `__self__` attribute of weight loader methods attached to the tensor. |
| 16 | |
| 17 | Used by `capture_layer_to_meta` to avoid circular references to layers in |
| 18 | `LAYERWISE_INFO`, leading to modules never being cleaned up. Without sanitation, |
| 19 | tensors will reference layers, and the WeakKeyDictionary will never evict entries, |
| 20 | even when the model is deleted. |
| 21 | |
| 22 | Args: |
| 23 | tensor: tensor to be sanitized |
| 24 | layer: layer whose references should be removed |
| 25 | |
| 26 | Returns: |
| 27 | sanitized tensor |
| 28 | |
| 29 | """ |
| 30 | for key, value in tensor.__dict__.items(): |
| 31 | if isinstance(value, MethodType) and value.__self__ is layer: |
| 32 | tensor.__dict__[key] = value.__func__.__get__(layer_ref_sentinel) |
| 33 | |
| 34 | return tensor |
| 35 | |
| 36 | |
| 37 | def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor: |
| 38 | """Restores references to layer held by tensor attributes. |
| 39 | |
| 40 | Used by `restore_layer_on_meta` to add back layer references, allowing for proper |
| 41 | weight loading. |
| 42 | |
| 43 | Args: |
| 44 | tensor: tensor to be sanitized |
| 45 | layer: layer whose references should be removed |
| 46 | |
| 47 | Returns: |
| 48 | sanitized tensor |
| 49 | |
| 50 | """ |
| 51 | for key, value in tensor.__dict__.items(): |
| 52 | if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel: |
| 53 | tensor.__dict__[key] = value.__func__.__get__(layer) |
| 54 | |
| 55 | return tensor |