| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | |
| 8 | def update_tensor_inplace(dst: torch.Tensor, src: torch.Tensor): |
| 9 | assert dst.dtype == src.dtype, "Tensors must have the same dtype" |
| 10 | |
| 11 | # update tensor shape and stride |
| 12 | dst.as_strided_(src.shape, src.stride()) |
| 13 | |
| 14 | # If not the same underlying storage move tensor data |
| 15 | if dst.data_ptr() != src.data_ptr(): |
| 16 | dst.copy_(src) |
| 17 | del src |
| 18 | |
| 19 | |
| 20 | # Newly generated tensors need to replace existing tensors that are |
| 21 | # already registered as parameters by vLLM (and won't be freed) |
| 22 | def replace_parameter( |
| 23 | mod: torch.nn.Module, name: str, new: torch.Tensor | torch.nn.Parameter |
| 24 | ) -> None: |
| 25 | old = getattr(mod, name) |
| 26 | if ( |
| 27 | type(old) is type(new) |
| 28 | and old.dtype == new.dtype |
| 29 | and old.untyped_storage().nbytes() == new.untyped_storage().nbytes() |
| 30 | ): |
| 31 | # If we can just update in-place to avoid re-registering |
| 32 | # can be faster if the underlying storage is the same |
| 33 | update_tensor_inplace(old, new) |
| 34 | else: |
| 35 | # Fallback re-register parameter, convert to Parameter if necessary |
| 36 | # this not only ensures we don't register a tensor as a parameter, but |
| 37 | # also ensures that all parameter subclasses get re-registered as |
| 38 | # parameters for `torch.compile` compatibility |
| 39 | if not isinstance(new, torch.nn.Parameter): |
| 40 | new = torch.nn.Parameter(new, requires_grad=False) |
| 41 | mod.register_parameter(name, torch.nn.Parameter(new, requires_grad=False)) |