| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable, Hashable |
| 5 | from fractions import Fraction |
| 6 | from typing import Any |
| 7 | from weakref import WeakValueDictionary |
| 8 | |
| 9 | import torch |
| 10 | from torch.nn import Parameter |
| 11 | |
| 12 | from vllm.distributed import ( |
| 13 | get_tensor_model_parallel_rank, |
| 14 | get_tensor_model_parallel_world_size, |
| 15 | ) |
| 16 | from vllm.logger import init_logger |
| 17 | |
| 18 | __all__ = [ |
| 19 | "BasevLLMParameter", |
| 20 | "PackedvLLMParameter", |
| 21 | "PerTensorScaleParameter", |
| 22 | "ModelWeightParameter", |
| 23 | "ChannelQuantScaleParameter", |
| 24 | "GroupQuantScaleParameter", |
| 25 | "PackedColumnParameter", |
| 26 | "RowvLLMParameter", |
| 27 | ] |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | |
| 32 | class BasevLLMParameter(Parameter): |
| 33 | """Base parameter for vLLM linear layers. Extends the torch.nn.parameter |
| 34 | by taking in a linear weight loader. Will copy the loaded weight |
| 35 | into the parameter when the provided weight loader is called. |
| 36 | """ |
| 37 | |
| 38 | def __new__(cls, data: torch.Tensor | None, **kwargs): |
| 39 | return super().__new__(cls, data=data, requires_grad=False) |
| 40 | |
| 41 | def __init__(self, data: torch.Tensor, weight_loader: Callable): |
| 42 | """Initialize the BasevLLMParameter. |
| 43 | |
| 44 | Args: |
| 45 | data: torch tensor with the parameter data |
| 46 | weight_loader: weight loader callable |
| 47 | |
| 48 | """ |
| 49 | # During weight loading, we often do something like: |
| 50 | # narrowed_tensor = param.data.narrow(0, offset, len) |
| 51 | # narrowed_tensor.copy_(real_weight) |
| 52 | # expecting narrowed_tensor and param.data to share the same storage. |
| 53 | # However, on TPUs, narrowed_tensor will lazily propagate to the base |
| 54 | # tensor, which is param.data, leading to the redundant memory usage. |
| 55 | # This sometimes causes OOM errors during model loading. To avoid this, |
| 56 | # we sync the param tensor after its weight loader is called. |
| 57 | from vllm.platforms import current_platform |
| 58 | |
| 59 | if current_platform.use_sync_weight_loader(): |
| 60 | weight_loader = current_platform.make_synced_weight_loader(weight_loader) |
| 61 | |
| 62 | self._weight_loader = weight_loader |
| 63 | self.tp_rank = get_tensor_model_parallel_rank() |
| 64 | self.tp_size = get_tensor_model_parallel_world_size() |
| 65 | |
| 66 | @property |
| 67 | def weight_loader(self) -> Callable: |
| 68 | # NOTE(@ksayers) some models such as mamba_mixer2 override the |
| 69 | # weight loader to support custom loading. In the future, model-specific |
| 70 | # weight loading should be implemented via Model.load_weights. In the |
| 71 | # meantime, support deleting and overriding `weight_loader` attribute |
| 72 | if self._weight_loader is None: |
| 73 | raise AttributeError( |
| 74 | f"{self.__class__.__name__} weight_loader attribute has been deleted" |
| 75 | ) |
| 76 | return self._weight_loader |
| 77 | |
| 78 | @weight_loader.setter |
| 79 | def weight_loader(self, value: Callable): |
| 80 | self._weight_loader = value |
| 81 | |
| 82 | @weight_loader.deleter |
| 83 | def weight_loader(self): |
| 84 | self._weight_loader = None # type: ignore[assignment] |
| 85 | |
| 86 | def _is_1d_and_scalar(self, loaded_weight: torch.Tensor): |
| 87 | cond1 = self.data.ndim == 1 and self.data.numel() == 1 |
| 88 | cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1 |
| 89 | return cond1 and cond2 |
| 90 | |
| 91 | def _assert_and_load(self, loaded_weight: torch.Tensor): |
| 92 | assert self.data.shape == loaded_weight.shape or self._is_1d_and_scalar( |
| 93 | loaded_weight |
| 94 | ) |
| 95 | self.data.copy_(loaded_weight) |
| 96 | |
| 97 | def load_column_parallel_weight(self, loaded_weight: torch.Tensor): |
| 98 | self._assert_and_load(loaded_weight) |
| 99 | |
| 100 | def load_row_parallel_weight(self, loaded_weight: torch.Tensor): |
| 101 | self._assert_and_load(loaded_weight) |
| 102 | |
| 103 | def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 104 | self._assert_and_load(loaded_weight) |
| 105 | |
| 106 | def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 107 | self._assert_and_load(loaded_weight) |
| 108 | |
| 109 | def _shard_id_as_int(self, shard_id: str | int) -> int: |
| 110 | if isinstance(shard_id, int): |
| 111 | return shard_id |
| 112 | |
| 113 | # if not int, assume shard_id for qkv |
| 114 | # map to int and return |
| 115 | qkv_idxs = {"q": 0, "k": 1, "v": 2} |
| 116 | assert isinstance(shard_id, str) |
| 117 | assert shard_id in qkv_idxs |
| 118 | return qkv_idxs[shard_id] |
| 119 | |
| 120 | @classmethod |
| 121 | def __torch_function__(cls, func, types, args=(), kwargs=None): |
| 122 | if kwargs is None: |
| 123 | kwargs = {} |
| 124 | return super().__torch_function__(func, types, args, kwargs) |
| 125 | |
| 126 | |
| 127 | class _ColumnvLLMParameter(BasevLLMParameter): |
| 128 | """Private class defining weight loading functionality |
| 129 | (load_merged_column_weight, load_qkv_weight) |
| 130 | for parameters being loaded into linear layers with column |
| 131 | parallelism. This includes QKV and MLP layers which are |
| 132 | not already fused on disk. Requires an output dimension |
| 133 | to be defined. Called within the weight loader of |
| 134 | each of the column parallel linear layers. |
| 135 | """ |
| 136 | |
| 137 | def __init__(self, output_dim: int, **kwargs): |
| 138 | self._output_dim = output_dim |
| 139 | super().__init__(**kwargs) |
| 140 | |
| 141 | @property |
| 142 | def output_dim(self): |
| 143 | return self._output_dim |
| 144 | |
| 145 | def load_column_parallel_weight(self, loaded_weight: torch.Tensor): |
| 146 | shard_size = self.data.shape[self.output_dim] |
| 147 | loaded_weight = loaded_weight.narrow( |
| 148 | self.output_dim, self.tp_rank * shard_size, shard_size |
| 149 | ) |
| 150 | assert self.data.shape == loaded_weight.shape |
| 151 | self.data.copy_(loaded_weight) |
| 152 | |
| 153 | def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 154 | shard_offset: int = kwargs["shard_offset"] |
| 155 | shard_size: int = kwargs["shard_size"] |
| 156 | |
| 157 | # TODO: move these to PackedColumnParameter and PackedvLLMParameter |
| 158 | if ( |
| 159 | isinstance(self, (PackedColumnParameter, PackedvLLMParameter)) |
| 160 | and self.packed_dim == self.output_dim |
| 161 | ): |
| 162 | shard_size, shard_offset = self.adjust_shard_indexes_for_packing( |
| 163 | shard_offset=shard_offset, shard_size=shard_size |
| 164 | ) |
| 165 | |
| 166 | param_data = self.data |
| 167 | |
| 168 | param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) |
| 169 | loaded_weight = loaded_weight.narrow( |
| 170 | self.output_dim, self.tp_rank * shard_size, shard_size |
| 171 | ) |
| 172 | assert param_data.shape == loaded_weight.shape |
| 173 | param_data.copy_(loaded_weight) |
| 174 | |
| 175 | def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 176 | shard_offset: int = kwargs["shard_offset"] |
| 177 | shard_size: int = kwargs["shard_size"] |
| 178 | shard_id: str = kwargs["shard_id"] |
| 179 | num_heads: int = kwargs["num_heads"] |
| 180 | |
| 181 | # TODO: move these to PackedColumnParameter and PackedvLLMParameter |
| 182 | if ( |
| 183 | isinstance(self, (PackedColumnParameter, PackedvLLMParameter)) |
| 184 | and self.output_dim == self.packed_dim |
| 185 | ): |
| 186 | shard_size, shard_offset = self.adjust_shard_indexes_for_packing( |
| 187 | shard_offset=shard_offset, shard_size=shard_size |
| 188 | ) |
| 189 | |
| 190 | param_data = self.data |
| 191 | shard_id_int = self.tp_rank if shard_id == "q" else self.tp_rank // num_heads |
| 192 | param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) |
| 193 | loaded_weight = loaded_weight.narrow( |
| 194 | self.output_dim, shard_id_int * shard_size, shard_size |
| 195 | ) |
| 196 | |
| 197 | assert param_data.shape == loaded_weight.shape |
| 198 | param_data.copy_(loaded_weight) |
| 199 | |
| 200 | |
| 201 | class RowvLLMParameter(BasevLLMParameter): |
| 202 | """Parameter class defining weight_loading functionality |
| 203 | (load_row_parallel_weight) for parameters being loaded |
| 204 | into linear layers with row parallel functionality. |
| 205 | Requires an input_dim to be defined. |
| 206 | """ |
| 207 | |
| 208 | def __init__(self, input_dim: int, **kwargs): |
| 209 | self._input_dim = input_dim |
| 210 | super().__init__(**kwargs) |
| 211 | |
| 212 | @property |
| 213 | def input_dim(self): |
| 214 | return self._input_dim |
| 215 | |
| 216 | def load_row_parallel_weight(self, loaded_weight: torch.Tensor): |
| 217 | shard_size = self.data.shape[self.input_dim] |
| 218 | loaded_weight = loaded_weight.narrow( |
| 219 | self.input_dim, self.tp_rank * shard_size, shard_size |
| 220 | ) |
| 221 | |
| 222 | if len(loaded_weight.shape) == 0: |
| 223 | loaded_weight = loaded_weight.reshape(1) |
| 224 | |
| 225 | assert self.data.shape == loaded_weight.shape |
| 226 | self.data.copy_(loaded_weight) |
| 227 | |
| 228 | |
| 229 | class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): |
| 230 | """Parameter class for linear layer weights. Uses both column and |
| 231 | row parallelism. |
| 232 | """ |
| 233 | |
| 234 | pass |
| 235 | |
| 236 | |
| 237 | class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): |
| 238 | """Parameter class for weight scales loaded for weights with |
| 239 | grouped quantization. Uses both column and row parallelism. |
| 240 | """ |
| 241 | |
| 242 | pass |
| 243 | |
| 244 | |
| 245 | class ChannelQuantScaleParameter(_ColumnvLLMParameter): |
| 246 | """Parameter class for weight scales loaded for weights with |
| 247 | channel-wise quantization. Equivalent to _ColumnvLLMParameter. |
| 248 | """ |
| 249 | |
| 250 | pass |
| 251 | |
| 252 | |
| 253 | class PerTensorScaleParameter(BasevLLMParameter): |
| 254 | """Parameter class for scales where the number of scales is |
| 255 | equivalent to the number of logical matrices in fused linear |
| 256 | layers (e.g. for QKV, there are 3 scales loaded from disk). |
| 257 | This is relevant to weights with per-tensor quantization. |
| 258 | Adds functionality to map the scalers to a shard during |
| 259 | weight loading. |
| 260 | |
| 261 | Note: additional parameter manipulation may be handled |
| 262 | for each quantization config specifically, within |
| 263 | process_weights_after_loading |
| 264 | """ |
| 265 | |
| 266 | def __init__(self, **kwargs): |
| 267 | super().__init__(**kwargs) |
| 268 | |
| 269 | # For row parallel layers, no sharding needed |
| 270 | # load weight into parameter as is |
| 271 | def load_row_parallel_weight(self, *args, **kwargs): |
| 272 | super().load_row_parallel_weight(*args, **kwargs) |
| 273 | |
| 274 | def load_merged_column_weight(self, *args, **kwargs): |
| 275 | self._load_into_shard_id(*args, **kwargs) |
| 276 | |
| 277 | def load_qkv_weight(self, *args, **kwargs): |
| 278 | self._load_into_shard_id(*args, **kwargs) |
| 279 | |
| 280 | def load_column_parallel_weight(self, *args, **kwargs): |
| 281 | super().load_row_parallel_weight(*args, **kwargs) |
| 282 | |
| 283 | def _load_into_shard_id( |
| 284 | self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs |
| 285 | ): |
| 286 | """Slice the parameter data based on the shard id for |
| 287 | loading. |
| 288 | """ |
| 289 | param_data = self.data |
| 290 | shard_id = self._shard_id_as_int(shard_id) |
| 291 | |
| 292 | # AutoFP8 scales do not have a shape |
| 293 | # compressed-tensors scales do have a shape |
| 294 | if len(loaded_weight.shape) != 0: |
| 295 | assert loaded_weight.shape[0] == 1 |
| 296 | loaded_weight = loaded_weight[0] |
| 297 | |
| 298 | param_data = param_data[shard_id] |
| 299 | assert param_data.shape == loaded_weight.shape |
| 300 | param_data.copy_(loaded_weight) |
| 301 | |
| 302 | |
| 303 | class PackedColumnParameter(_ColumnvLLMParameter): |
| 304 | """Parameter for model parameters which are packed on disk |
| 305 | and support column parallelism only. See PackedvLLMParameter |
| 306 | for more details on the packed properties. |
| 307 | """ |
| 308 | |
| 309 | def __init__( |
| 310 | self, |
| 311 | packed_factor: int | Fraction, |
| 312 | packed_dim: int, |
| 313 | marlin_tile_size: int | None = None, |
| 314 | **kwargs, |
| 315 | ): |
| 316 | self._packed_factor = packed_factor |
| 317 | self._packed_dim = packed_dim |
| 318 | self._marlin_tile_size = marlin_tile_size |
| 319 | super().__init__(**kwargs) |
| 320 | |
| 321 | @property |
| 322 | def packed_dim(self): |
| 323 | return self._packed_dim |
| 324 | |
| 325 | @property |
| 326 | def packed_factor(self): |
| 327 | return self._packed_factor |
| 328 | |
| 329 | @property |
| 330 | def marlin_tile_size(self): |
| 331 | return self._marlin_tile_size |
| 332 | |
| 333 | def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): |
| 334 | return _adjust_shard_indexes_for_packing( |
| 335 | shard_size=shard_size, |
| 336 | shard_offset=shard_offset, |
| 337 | packed_factor=self.packed_factor, |
| 338 | marlin_tile_size=self.marlin_tile_size, |
| 339 | ) |
| 340 | |
| 341 | |
| 342 | class PackedvLLMParameter(ModelWeightParameter): |
| 343 | """Parameter for model weights which are packed on disk. |
| 344 | Example: GPTQ Marlin weights are int4 or int8, packed into int32. |
| 345 | Extends the ModelWeightParameter to take in the |
| 346 | packed factor, the packed dimension, and optionally, marlin |
| 347 | tile size for marlin kernels. Adjusts the shard_size and |
| 348 | shard_offset for fused linear layers model weight loading |
| 349 | by accounting for packing and optionally, marlin tile size. |
| 350 | """ |
| 351 | |
| 352 | def __init__( |
| 353 | self, |
| 354 | packed_factor: int | Fraction, |
| 355 | packed_dim: int, |
| 356 | marlin_tile_size: int | None = None, |
| 357 | **kwargs, |
| 358 | ): |
| 359 | self._packed_factor = packed_factor |
| 360 | self._packed_dim = packed_dim |
| 361 | self._marlin_tile_size = marlin_tile_size |
| 362 | super().__init__(**kwargs) |
| 363 | |
| 364 | @property |
| 365 | def packed_dim(self): |
| 366 | return self._packed_dim |
| 367 | |
| 368 | @property |
| 369 | def packed_factor(self): |
| 370 | return self._packed_factor |
| 371 | |
| 372 | @property |
| 373 | def marlin_tile_size(self): |
| 374 | return self._marlin_tile_size |
| 375 | |
| 376 | def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): |
| 377 | return _adjust_shard_indexes_for_packing( |
| 378 | shard_size=shard_size, |
| 379 | shard_offset=shard_offset, |
| 380 | packed_factor=self.packed_factor, |
| 381 | marlin_tile_size=self.marlin_tile_size, |
| 382 | ) |
| 383 | |
| 384 | |
| 385 | class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): |
| 386 | """Parameter class for weight scales loaded for weights with |
| 387 | block-wise quantization. Uses both column and row parallelism. |
| 388 | """ |
| 389 | |
| 390 | pass |
| 391 | |
| 392 | |
| 393 | class SharedWeightParameter(BasevLLMParameter): |
| 394 | """Parameter for weights with many shared tensors across a model. |
| 395 | |
| 396 | For example, when applying transforms to the "gate" and "up" partitions of |
| 397 | `MergedColumnParallelLinear`, the transform weights must stay separate |
| 398 | tensors in order to allow for tensor memory sharing between layers. |
| 399 | """ |
| 400 | |
| 401 | # global registry for sharing tensors based on passed `data_key` |
| 402 | # this dict holds weaksrefs to avoid memory leak after model cleanup |
| 403 | tensors_registry: WeakValueDictionary = WeakValueDictionary() |
| 404 | |
| 405 | # local container for strong references to shared tensors |
| 406 | # this set compensates for the fact that torch.nn.Parameter |
| 407 | # and Parameter subclasses do not hold reliable references to tensors |
| 408 | local_tensors: set[torch.Tensor] |
| 409 | |
| 410 | # dictionary mapping partition indices to associated parameters |
| 411 | partitions: dict[int, ModelWeightParameter | Parameter] |
| 412 | |
| 413 | def __new__(cls, **kwargs): |
| 414 | return super().__new__(cls, data=None, **kwargs) |
| 415 | |
| 416 | def __init__(self, input_dim: int = 1, output_dim: int = 0, **kwargs): |
| 417 | weight_loader: Callable = kwargs.get("weight_loader") # type: ignore[assignment] |
| 418 | super().__init__(data=None, weight_loader=weight_loader) |
| 419 | |
| 420 | self.local_tensors = set() |
| 421 | self.partitions = {} |
| 422 | self.kwargs = { |
| 423 | "input_dim": input_dim, |
| 424 | "output_dim": output_dim, |
| 425 | "weight_loader": self._fake_weight_loader, |
| 426 | } |
| 427 | |
| 428 | if self.tp_size > 1: |
| 429 | raise NotImplementedError( |
| 430 | f"{self.__class__.__name__} does not " |
| 431 | "currently support tensor parallelism" |
| 432 | ) |
| 433 | |
| 434 | def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: Any): |
| 435 | """Add a partition to the weight parameter. Partitions whose `data_key` |
| 436 | is the same will share tensor data |
| 437 | |
| 438 | Args: |
| 439 | index: index of partition to add |
| 440 | data_key: hashable key used to key shared tensors |
| 441 | *args: arguments for `torch.empty` |
| 442 | **kwargs: keyword arguments for `torch.empty` |
| 443 | |
| 444 | """ |
| 445 | # load (shared) tensor using `data_key` |
| 446 | if data_key not in self.tensors_registry: |
| 447 | data = torch.empty(*args, **kwargs) |
| 448 | self.tensors_registry[data_key] = data |
| 449 | else: |
| 450 | data = self.tensors_registry[data_key] |
| 451 | |
| 452 | # create associated model parameter |
| 453 | self.partitions[index] = ModelWeightParameter(data=data, **self.kwargs) # type: ignore[arg-type] |
| 454 | |
| 455 | # hold local reference, since ModelWeightParameter does not |
| 456 | # see https://github.com/pytorch/pytorch/issues/75932 |
| 457 | self.local_tensors.add(data) |
| 458 | |
| 459 | def load_column_parallel_weight(self, loaded_weight: torch.Tensor): |
| 460 | assert len(self.partitions) == 1 and 0 in self.partitions |
| 461 | partition = self.partitions[0] |
| 462 | |
| 463 | ModelWeightParameter.load_column_parallel_weight(partition, loaded_weight) |
| 464 | |
| 465 | def load_row_parallel_weight(self, loaded_weight: torch.Tensor): |
| 466 | assert len(self.partitions) == 1 and 0 in self.partitions |
| 467 | partition = self.partitions[0] |
| 468 | |
| 469 | ModelWeightParameter.load_row_parallel_weight(partition, loaded_weight) |
| 470 | |
| 471 | def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 472 | partition_id = kwargs.pop("shard_id") |
| 473 | partition_id = self._shard_id_as_int(partition_id) |
| 474 | partition = self.partitions[partition_id] |
| 475 | |
| 476 | input_dim = self.kwargs.get("input_dim") |
| 477 | shard_size = partition.data.size(input_dim) // self.tp_size |
| 478 | shard_offset = self.tp_rank * shard_size |
| 479 | |
| 480 | ModelWeightParameter.load_merged_column_weight( |
| 481 | partition, loaded_weight, shard_offset=shard_offset, shard_size=shard_size |
| 482 | ) |
| 483 | |
| 484 | def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs): |
| 485 | partition_id = self._shard_id_as_int(kwargs.pop("shard_id")) |
| 486 | partition = self.partitions[partition_id] |
| 487 | |
| 488 | input_dim = self.kwargs.get("input_dim") |
| 489 | shard_size = partition.data.size(input_dim) // self.tp_size |
| 490 | shard_offset = self.tp_rank * shard_size |
| 491 | shard_id = "q" # fake first partition |
| 492 | num_heads = kwargs.get("num_heads") |
| 493 | |
| 494 | ModelWeightParameter.load_qkv_weight( |
| 495 | partition, |
| 496 | loaded_weight, |
| 497 | shard_offset=shard_offset, |
| 498 | shard_size=shard_size, |
| 499 | shard_id=shard_id, |
| 500 | num_heads=num_heads, |
| 501 | ) |
| 502 | |
| 503 | def process_weights_after_loading(self): |
| 504 | for key in self.partitions: |
| 505 | self.partitions[key] = torch.nn.Parameter( |
| 506 | data=self.partitions[key].data, requires_grad=False |
| 507 | ) |
| 508 | |
| 509 | @property |
| 510 | def data(self): |
| 511 | raise ValueError( |
| 512 | "Accessing `data` of a `SharedWeightParameter` is not allowed. " |
| 513 | "Instead, use `get_partition` to get the weight of " |
| 514 | "the particular partition you want to access" |
| 515 | ) |
| 516 | |
| 517 | def _fake_weight_loader( |
| 518 | self, |
| 519 | param: BasevLLMParameter, |
| 520 | loaded_weight: torch.Tensor, |
| 521 | loaded_weight_shard_id: str | int | None, |
| 522 | ): |
| 523 | raise ValueError( |
| 524 | "When loading partition weights of " |
| 525 | f"{self.__class__.__name__}, use methods provided by " |
| 526 | f"{self.__class__.__name__}, not partition loader" |
| 527 | ) |
| 528 | |
| 529 | |
| 530 | def permute_param_layout_( |
| 531 | param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs |
| 532 | ) -> BasevLLMParameter: |
| 533 | """Permute a parameter's layout to the specified input and output dimensions, |
| 534 | useful for forcing the parameter into a known layout, for example, if I need |
| 535 | a packed (quantized) weight matrix to be in the layout |
| 536 | {input_dim = 0, output_dim = 1, packed_dim = 0} |
| 537 | then I can call: |
| 538 | permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) |
| 539 | to ensure x is in the correct layout (permuting it to the correct layout if |
| 540 | required, asserting if it cannot get it to the correct layout) |
| 541 | """ |
| 542 | curr_input_dim = getattr(param, "input_dim", None) |
| 543 | curr_output_dim = getattr(param, "output_dim", None) |
| 544 | |
| 545 | if curr_input_dim is None or curr_output_dim is None: |
| 546 | assert param.data.dim() == 2, ( |
| 547 | "permute_param_layout_ only supports 2D parameters when either " |
| 548 | "input_dim or output_dim is not set" |
| 549 | ) |
| 550 | |
| 551 | # if one of the dimensions is not set, set it to the opposite of the other |
| 552 | # we can only do this since we asserted the parameter is 2D above |
| 553 | if curr_input_dim is None: |
| 554 | assert curr_output_dim is not None, "either input or output dim must be set" |
| 555 | curr_input_dim = (curr_output_dim + 1) % 2 |
| 556 | if curr_output_dim is None: |
| 557 | assert curr_input_dim is not None, "either input or output dim must be set" |
| 558 | curr_output_dim = (curr_input_dim + 1) % 2 |
| 559 | |
| 560 | # create permutation from the current layout to the layout with |
| 561 | # self.input_dim at input_dim and self.output_dim at output_dim preserving |
| 562 | # other dimensions |
| 563 | perm = [ |
| 564 | i for i in range(param.data.dim()) if i not in [curr_input_dim, curr_output_dim] |
| 565 | ] |
| 566 | perm.insert(input_dim, curr_input_dim) |
| 567 | perm.insert(output_dim, curr_output_dim) |
| 568 | |
| 569 | if "packed_dim" in kwargs: |
| 570 | assert ( |
| 571 | hasattr(param, "packed_dim") |
| 572 | and param.packed_dim == perm[kwargs["packed_dim"]] |
| 573 | ), "permute_param_layout_ currently doesn't support repacking" |
| 574 | |
| 575 | param.data = param.data.permute(*perm) |
| 576 | if hasattr(param, "_input_dim"): |
| 577 | param._input_dim = input_dim |
| 578 | if hasattr(param, "_output_dim"): |
| 579 | param._output_dim = output_dim |
| 580 | if "packed_dim" in kwargs and hasattr(param, "_packed_dim"): |
| 581 | param._packed_dim = kwargs["packed_dim"] |
| 582 | |
| 583 | return param |
| 584 | |
| 585 | |
| 586 | def _adjust_shard_indexes_for_marlin(shard_size, shard_offset, marlin_tile_size): |
| 587 | return shard_size * marlin_tile_size, shard_offset * marlin_tile_size |
| 588 | |
| 589 | |
| 590 | def _adjust_shard_indexes_for_packing( |
| 591 | shard_size, shard_offset, packed_factor, marlin_tile_size |
| 592 | ): |
| 593 | shard_size = round(shard_size // packed_factor) |
| 594 | shard_offset = round(shard_offset // packed_factor) |
| 595 | if marlin_tile_size is not None: |
| 596 | return _adjust_shard_indexes_for_marlin( |
| 597 | shard_size=shard_size, |
| 598 | shard_offset=shard_offset, |
| 599 | marlin_tile_size=marlin_tile_size, |
| 600 | ) |
| 601 | |
| 602 | return shard_size, shard_offset |