[Quantization] Migrate GGUF diffusion model support to OOT plugin (#4769)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py authored and GitHub committed
1 parent
ff6f0da271
commit
b95905deb1
35 files changed
+272
-1614
No files matched your search
@@ -11,7 +11,6 @@ from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from torch import nn
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.config.load import LoadConfig
|
||||
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
|
||||
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
|
||||
get_checkpoint_adapter,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
@@ -381,10 +379,7 @@ class DiffusersPipelineLoader:
|
||||
# DiffusionPipeline.from_pretrained() internally; it does
|
||||
# not use our native customized pipeline classes.
|
||||
cast(DiffusersAdapterPipeline, model).load_weights()
|
||||
elif self._is_gguf_quantization():
|
||||
self._load_weights_with_gguf(model)
|
||||
else:
|
||||
# Quantization does not happen in `load_weights` but after it
|
||||
self.load_weights(model)
|
||||
|
||||
# Process weights after loading for quantization (e.g., FP8 online quantization)
|
||||
@@ -454,7 +449,7 @@ class DiffusersPipelineLoader:
|
||||
def _is_expected_quantized_weight(name: str) -> bool:
|
||||
"""Return True if *name* is a quantization-specific parameter.
|
||||
|
||||
Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
|
||||
Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra
|
||||
parameters that have no counterpart in an unquantized checkpoint.
|
||||
These are expected to be absent and should not trigger a load error.
|
||||
"""
|
||||
@@ -467,8 +462,6 @@ class DiffusersPipelineLoader:
|
||||
".weight_scale",
|
||||
".weight_scale_inv",
|
||||
".input_scale",
|
||||
# GGUF
|
||||
".qweight_type",
|
||||
# INT8 (weight_scale already covered above)
|
||||
)
|
||||
return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
|
||||
@@ -500,116 +493,6 @@ class DiffusersPipelineLoader:
|
||||
if unexpected_missing:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
|
||||
|
||||
def _get_gguf_model_from_config(self) -> str | None:
|
||||
"""Extract gguf_model from quant_config (handles both dict and object)."""
|
||||
if self.quant_config is None:
|
||||
return None
|
||||
if isinstance(self.quant_config, dict):
|
||||
return self.quant_config.get("gguf_model")
|
||||
return getattr(self.quant_config, "gguf_model", None)
|
||||
|
||||
def _is_gguf_quantization(self) -> bool:
|
||||
"""Check whether or not this pipeline loader is pointing at a GGUF config."""
|
||||
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
|
||||
|
||||
maybe_gguf_model = self._get_gguf_model_from_config()
|
||||
is_gguf = (
|
||||
isinstance(self.quant_config, DiffusionGGUFConfig)
|
||||
or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
|
||||
or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
|
||||
)
|
||||
if is_gguf:
|
||||
if not maybe_gguf_model:
|
||||
raise ValueError("GGUF quantization requires gguf_model")
|
||||
return True
|
||||
|
||||
if isinstance(self.quant_config, dict):
|
||||
return False
|
||||
|
||||
return bool(maybe_gguf_model)
|
||||
|
||||
def _is_transformer_source(self, source: "ComponentSource") -> bool:
|
||||
if source.subfolder == "transformer":
|
||||
return True
|
||||
return source.prefix.startswith("transformer.")
|
||||
|
||||
def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
|
||||
# Avoid model.state_dict() here because GGUF uses UninitializedParameter
|
||||
# which raises during detach(). Collect names directly.
|
||||
return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
|
||||
|
||||
def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
|
||||
if os.path.isfile(gguf_model):
|
||||
return gguf_model
|
||||
# repo_id/filename.gguf
|
||||
if "/" in gguf_model and gguf_model.endswith(".gguf"):
|
||||
repo_id, filename = gguf_model.rsplit("/", 1)
|
||||
return hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename=filename,
|
||||
revision=revision,
|
||||
cache_dir=self.load_config.download_dir,
|
||||
)
|
||||
# repo_id:quant_type
|
||||
if "/" in gguf_model and ":" in gguf_model:
|
||||
repo_id, quant_type = gguf_model.rsplit(":", 1)
|
||||
return download_gguf(
|
||||
repo_id,
|
||||
quant_type,
|
||||
cache_dir=self.load_config.download_dir,
|
||||
revision=revision,
|
||||
ignore_patterns=self.load_config.ignore_patterns,
|
||||
)
|
||||
raise ValueError(
|
||||
f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, "
|
||||
"<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)"
|
||||
)
|
||||
|
||||
def _get_gguf_weights_iterator(
|
||||
self,
|
||||
source: "ComponentSource",
|
||||
model: nn.Module,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
gguf_model = self._get_gguf_model_from_config()
|
||||
if gguf_model is None:
|
||||
raise ValueError("GGUF quantization requires quantization_config.gguf_model")
|
||||
gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
|
||||
adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
|
||||
weights_iter = adapter.weights_iterator()
|
||||
return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
|
||||
|
||||
def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
|
||||
sources = self._get_weight_sources(model)
|
||||
loaded: set[str] = set()
|
||||
loadable_names: set[str] | None = None
|
||||
|
||||
for source in sources:
|
||||
if self._is_transformer_source(source):
|
||||
loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
|
||||
|
||||
# GGUF checkpoints can be transformer-only or partially quantized.
|
||||
# Only fall back to HF if this source still has missing loadable weights.
|
||||
loadable_names = loadable_names or self._get_model_loadable_names(model)
|
||||
has_missing_for_source = any(
|
||||
name.startswith(source.prefix) and name not in loaded for name in loadable_names
|
||||
)
|
||||
if not has_missing_for_source:
|
||||
continue
|
||||
|
||||
hf_iter = self._get_weights_iterator(source)
|
||||
hf_iter = (
|
||||
(name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded
|
||||
)
|
||||
loaded |= model.load_weights(hf_iter)
|
||||
else:
|
||||
loaded |= model.load_weights(self._get_weights_iterator(source))
|
||||
|
||||
weights_to_load = self._get_expected_parameter_names(model)
|
||||
weights_not_loaded = weights_to_load - loaded
|
||||
if weights_not_loaded:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
|
||||
return loaded
|
||||
|
||||
def _init_from_load_format(
|
||||
self,
|
||||
load_format: str,
|
||||
|
||||
Reference in new issue
Block a user