55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Compatibility helpers for vLLM model weight loading."""
|
|
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.model_executor.models.utils import (
|
|
AutoWeightsLoader as _VllmAutoWeightsLoader,
|
|
)
|
|
from vllm.model_executor.models.utils import (
|
|
WeightsMapper,
|
|
)
|
|
|
|
|
|
class AutoWeightsLoader(_VllmAutoWeightsLoader):
|
|
"""Preserve vLLM's pre-0.29 weight-filtering arguments.
|
|
|
|
vLLM removed ``skip_prefixes`` and ``skip_substrs`` from
|
|
:class:`AutoWeightsLoader` in favor of mapper-based filtering. Omni models
|
|
still use those arguments to exclude checkpoint components owned by a
|
|
different stage. Convert the filters to a ``WeightsMapper`` so the old
|
|
behavior is retained without bypassing the current loader implementation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
module: nn.Module,
|
|
*,
|
|
skip_prefixes: list[str] | None = None,
|
|
skip_substrs: list[str] | None = None,
|
|
ignore_unexpected_prefixes: list[str] | None = None,
|
|
ignore_unexpected_suffixes: list[str] | None = None,
|
|
) -> None:
|
|
super().__init__(
|
|
module,
|
|
ignore_unexpected_prefixes=ignore_unexpected_prefixes,
|
|
ignore_unexpected_suffixes=ignore_unexpected_suffixes,
|
|
)
|
|
self._skip_mapper = WeightsMapper(
|
|
orig_to_new_prefix={prefix: None for prefix in skip_prefixes or ()},
|
|
orig_to_new_substr={substring: None for substring in skip_substrs or ()},
|
|
)
|
|
|
|
def load_weights(
|
|
self,
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|
*,
|
|
mapper: WeightsMapper | None = None,
|
|
) -> set[str]:
|
|
mapper = (mapper or WeightsMapper()) | self._skip_mapper
|
|
return super().load_weights(weights, mapper=mapper)
|