| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utils for model executor.""" |
| 4 | |
| 5 | import copy |
| 6 | from contextlib import contextmanager |
| 7 | from contextvars import ContextVar |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 13 | |
| 14 | _weights_pre_processed: ContextVar[bool] = ContextVar( |
| 15 | "weights_pre_processed", default=False |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | @contextmanager |
| 20 | def weights_already_processed(): |
| 21 | """Mark that weights are already in post-processed (runtime) format, so |
| 22 | ``process_weights_after_loading`` skips tensor transforms (used by the |
| 23 | weight cache IPC loader). |
| 24 | """ |
| 25 | token = _weights_pre_processed.set(True) |
| 26 | try: |
| 27 | yield |
| 28 | finally: |
| 29 | _weights_pre_processed.reset(token) |
| 30 | |
| 31 | |
| 32 | def is_weights_pre_processed() -> bool: |
| 33 | return _weights_pre_processed.get() |
| 34 | |
| 35 | |
| 36 | def set_weight_attrs( |
| 37 | weight: torch.Tensor, |
| 38 | weight_attrs: dict[str, Any] | None, |
| 39 | ): |
| 40 | """Set attributes on a weight tensor. |
| 41 | |
| 42 | This method is used to set attributes on a weight tensor. This method |
| 43 | will not overwrite existing attributes. |
| 44 | |
| 45 | Args: |
| 46 | weight: The weight tensor. |
| 47 | weight_attrs: A dictionary of attributes to set on the weight tensor. |
| 48 | |
| 49 | """ |
| 50 | if weight_attrs is None: |
| 51 | return |
| 52 | for key, value in weight_attrs.items(): |
| 53 | assert not hasattr(weight, key), f"Overwriting existing tensor attribute: {key}" |
| 54 | |
| 55 | # NOTE(woosuk): During weight loading, we often do something like: |
| 56 | # narrowed_tensor = param.data.narrow(0, offset, len) |
| 57 | # narrowed_tensor.copy_(real_weight) |
| 58 | # expecting narrowed_tensor and param.data to share the same storage. |
| 59 | # However, on TPUs, narrowed_tensor will lazily propagate to the base |
| 60 | # tensor, which is param.data, leading to the redundant memory usage. |
| 61 | # This sometimes causes OOM errors during model loading. To avoid this, |
| 62 | # we sync the param tensor after its weight loader is called. |
| 63 | # TODO(woosuk): Remove this hack once we have a better solution. |
| 64 | from vllm.platforms import current_platform |
| 65 | |
| 66 | if current_platform.use_sync_weight_loader() and key == "weight_loader": |
| 67 | value = current_platform.make_synced_weight_loader(value) |
| 68 | setattr(weight, key, value) |
| 69 | |
| 70 | |
| 71 | def replace_parameter( |
| 72 | layer: torch.nn.Module, |
| 73 | param_name: str, |
| 74 | new_tensor: torch.Tensor | None, |
| 75 | prefer_copy: bool = False, |
| 76 | ): |
| 77 | """Replace a parameter of a layer while maintaining the ability to reload the |
| 78 | weight. Called within implementations of the `process_weights_after_loading` method. |
| 79 | |
| 80 | Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as |
| 81 | ``is_shuffled``) are carried over to the replacement parameter, except |
| 82 | ``weight_loader``, which is always taken from the existing parameter. |
| 83 | |
| 84 | Attributes of the existing parameter are otherwise dropped when a new |
| 85 | parameter is registered, but kept when ``prefer_copy`` reuses it in place. |
| 86 | |
| 87 | This function should not be called on weights which are tied/shared |
| 88 | |
| 89 | Args: |
| 90 | layer: Layer containing parameter to replace |
| 91 | param_name: Name of parameter to replace |
| 92 | new_tensor: New data of the new parameter, or None to set the parameter |
| 93 | to None |
| 94 | prefer_copy: If True and the existing parameter is compatible with |
| 95 | ``new_tensor`` (same shape, dtype, and device), copy ``new_tensor`` |
| 96 | into the existing parameter in place rather than re-registering |
| 97 | a new parameter. This preserves the parameter's storage address |
| 98 | (``data_ptr``), which is required for captured CUDA graphs to |
| 99 | remain valid across weight updates (e.g. in RL training loops). |
| 100 | |
| 101 | """ |
| 102 | # should not be used on a tied/shared param |
| 103 | |
| 104 | # If new_tensor is None, set the parameter to None |
| 105 | if new_tensor is None: |
| 106 | setattr(layer, param_name, None) |
| 107 | return |
| 108 | |
| 109 | old_param: torch.nn.Parameter | None = getattr(layer, param_name, None) |
| 110 | |
| 111 | new_tensor_attrs = dict(new_tensor.__dict__) |
| 112 | |
| 113 | # `weight_loader` is the only attribute not ported over from new_tensor, |
| 114 | # old_param.weight_loader only is supported. `_weight_loader` is |
| 115 | # `BasevLLMParameter`'s backing field for its `weight_loader` property, so |
| 116 | # a loader riding along under that name has to be dropped as well. |
| 117 | new_tensor_attrs.pop("weight_loader", None) |
| 118 | new_tensor_attrs.pop("_weight_loader", None) |
| 119 | |
| 120 | if isinstance(new_tensor, torch.nn.Parameter): |
| 121 | new_tensor = new_tensor.data |
| 122 | |
| 123 | if ( |
| 124 | prefer_copy |
| 125 | and old_param is not None |
| 126 | and old_param.shape == new_tensor.shape |
| 127 | and old_param.dtype == new_tensor.dtype |
| 128 | and old_param.device == new_tensor.device |
| 129 | ): |
| 130 | old_param.copy_(new_tensor) |
| 131 | for attr_name, attr in new_tensor_attrs.items(): |
| 132 | setattr(old_param, attr_name, attr) |
| 133 | return |
| 134 | |
| 135 | new_param = torch.nn.Parameter(new_tensor, requires_grad=False) |
| 136 | |
| 137 | for attr_name, attr in new_tensor_attrs.items(): |
| 138 | setattr(new_param, attr_name, attr) |
| 139 | |
| 140 | if old_param is not None and hasattr(old_param, "weight_loader"): |
| 141 | weight_loader = old_param.weight_loader |
| 142 | set_weight_attrs(new_param, {"weight_loader": weight_loader}) |
| 143 | |
| 144 | setattr(layer, param_name, new_param) |
| 145 | |
| 146 | |
| 147 | def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]: |
| 148 | parent_map = getattr(model, "packed_modules_mapping", None) |
| 149 | parent_map = copy.deepcopy(parent_map) if parent_map is not None else {} |
| 150 | |
| 151 | # don't infer mapping if the model has defined it explicitly. |
| 152 | if parent_map: |
| 153 | return parent_map |
| 154 | |
| 155 | # We only check main components instead of whole model submodules |
| 156 | for child in model.children(): |
| 157 | child_map = getattr(child, "packed_modules_mapping", None) |
| 158 | child_map = copy.deepcopy(child_map) if child_map is not None else {} |
| 159 | |
| 160 | if any((k in parent_map and parent_map[k] != v) for k, v in child_map.items()): |
| 161 | raise ValueError( |
| 162 | f"Can't update {type(model).__name__}'s packed_modules_mapping " |
| 163 | f"safely because of conflicts from {type(child).__name__}." |
| 164 | ) |
| 165 | else: |
| 166 | parent_map.update(child_map) |
| 167 | return parent_map |
| 168 | |
| 169 | |
| 170 | def get_moe_expert_mapping( |
| 171 | model: torch.nn.Module, |
| 172 | ) -> list[tuple[str, str, int, str]]: |
| 173 | """Get the expert mapping from a model. |
| 174 | |
| 175 | It will be retrieved from the first module that has a `get_expert_mapping` method. |
| 176 | If the model manually implements `get_expert_mapping`, it will be used. |
| 177 | Otherwise, it will use the first RoutedExperts layer.""" |
| 178 | for _, module in model.named_modules(): |
| 179 | get_mapping = getattr(module, "get_expert_mapping", None) |
| 180 | if get_mapping is not None: |
| 181 | return get_mapping() |
| 182 | raise ValueError("No module in the model has a `get_expert_mapping` method.") |
| 183 | |
| 184 | |
| 185 | def maybe_disable_graph_partition(current_backend: str) -> dict[str, bool]: |
| 186 | if current_backend == "inductor" and is_torch_equal_or_newer("2.9.0.dev"): |
| 187 | return {"graph_partition": False} |
| 188 | else: |
| 189 | return {} |