97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Generator
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import gguf
|
|
import numpy as np
|
|
import torch
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import (
|
|
DiffusersPipelineLoader,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class MappedTensor:
|
|
name: str
|
|
tensor: Any
|
|
tensor_type: Any
|
|
row_slice: slice | None = None
|
|
swap_scale_shift: bool = False
|
|
|
|
|
|
class GGUFAdapter(ABC):
|
|
"""Base class for model-specific GGUF adapters."""
|
|
|
|
def __init__(
|
|
self,
|
|
gguf_file: str,
|
|
model: torch.nn.Module,
|
|
source: DiffusersPipelineLoader.ComponentSource,
|
|
od_config: OmniDiffusionConfig,
|
|
) -> None:
|
|
self.gguf_file = gguf_file
|
|
self.model = model
|
|
self.source = source
|
|
self.od_config = od_config
|
|
|
|
@staticmethod
|
|
def is_compatible(
|
|
od_config: OmniDiffusionConfig,
|
|
model: torch.nn.Module,
|
|
source: DiffusersPipelineLoader.ComponentSource,
|
|
) -> bool:
|
|
return False
|
|
|
|
@abstractmethod
|
|
def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
raise NotImplementedError
|
|
|
|
|
|
# FIXME(Isotr0py): Sync implemnentation with upstream vLLM?
|
|
def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]:
|
|
"""
|
|
Iterate over the quant weights in the model gguf files and convert
|
|
them to torch tensors.
|
|
Be careful of the order of yielding weight types and weights data,
|
|
we have to yield all weight types first before yielding any weights.
|
|
Otherwise it would cause issue when loading weights with for packed
|
|
layer with different quant types.
|
|
"""
|
|
|
|
reader = gguf.GGUFReader(gguf_file)
|
|
|
|
for tensor in reader.tensors:
|
|
weight_type = tensor.tensor_type
|
|
name = tensor.name
|
|
|
|
if weight_type.name not in ("F32", "F16"):
|
|
weight_type_name = name.replace("weight", "qweight_type")
|
|
weight_type = torch.tensor(weight_type)
|
|
yield weight_type_name, weight_type
|
|
|
|
for tensor in reader.tensors:
|
|
weight = tensor.data
|
|
weight_type = tensor.tensor_type
|
|
name = tensor.name
|
|
if weight_type.name not in ("F32", "F16"):
|
|
name = name.replace("weight", "qweight")
|
|
if weight_type.name == "BF16" and tensor.data.dtype == np.uint8:
|
|
# BF16 is currently the only "quantization" type that isn't
|
|
# actually quantized but is read as a raw byte tensor.
|
|
# Reinterpret as `torch.bfloat16` tensor.
|
|
weight = weight.view(np.uint16)
|
|
if reader.byte_order == "S":
|
|
# GGUF endianness != system endianness
|
|
weight = weight.byteswap()
|
|
param = torch.tensor(weight).view(torch.bfloat16)
|
|
else:
|
|
param = torch.tensor(weight)
|
|
yield name, param
|