nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
SolitaryThinker
2025-03-16 20:18:03 +00:00
parent 8140cb269f
commit dd10588fb7
3 changed files with 2 additions and 514 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/v1/models/loader/utils.py
-167
View File
@@ -1,20 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
"""Utilities for selecting and loading models."""
import contextlib
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Type
import torch
import transformers
from torch import nn
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from vllm.config import ModelConfig, ModelImpl
from fastvideo.v1.logger import init_logger
from vllm.model_executor.models import ModelRegistry
from vllm.model_executor.models.adapters import (as_classification_model,
as_embedding_model,
as_reward_model)
logger = init_logger(__name__)
@@ -26,159 +15,3 @@ def set_default_torch_dtype(dtype: torch.dtype):
torch.set_default_dtype(dtype)
yield
torch.set_default_dtype(old_dtype)
def load_hf_config_from_subdir(model_path: str, component_name: str, trust_remote_code: bool = False, revision: Optional[str] = None):
"""
Load a HuggingFace config from a component subdirectory.
Args:
model_path: Path to the model directory
component_name: Name of the component subdirectory (e.g., "text_encoder", "vae")
trust_remote_code: Whether to trust remote code when loading the config
revision: Optional revision to use when loading from HuggingFace Hub
Returns:
The loaded HuggingFace config
"""
import os
from transformers import AutoConfig
component_path = os.path.join(model_path, component_name)
if not os.path.exists(component_path):
raise ValueError(f"Component directory {component_path} does not exist")
config_path = os.path.join(component_path, "config.json")
if not os.path.exists(config_path):
raise ValueError(f"Config file {config_path} does not exist")
logger.info(f"Loading config from {config_path}")
config = AutoConfig.from_pretrained(
component_path,
trust_remote_code=trust_remote_code,
revision=revision
)
return config
def is_transformers_impl_compatible(
arch: str,
module: Optional[transformers.PreTrainedModel] = None) -> bool:
mod = module or getattr(transformers, arch, None)
if mod is None:
return False
if hasattr(mod, "supports_backend"):
return mod.is_backend_compatible()
else:
return mod._supports_flex_attn
def resolve_transformers_fallback(model_config: ModelConfig,
architectures: list[str]):
for i, arch in enumerate(architectures):
if arch == "TransformersModel":
continue
auto_map: dict[str, str] = getattr(model_config.hf_config, "auto_map",
None) or dict()
# Make sure that config class is always initialized before model class,
# otherwise the model class won't be able to access the config class,
# the expected auto_map should have correct order like:
# "auto_map": {
# "AutoConfig": "<your-repo-name>--<config-name>",
# "AutoModel": "<your-repo-name>--<config-name>",
# "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
# },
auto_modules = {
name: get_class_from_dynamic_module(module, model_config.model)
for name, module in sorted(auto_map.items(), key=lambda x: x[0])
}
custom_model_module = auto_modules.get("AutoModel")
# TODO(Isotr0py): Further clean up these raises.
# perhaps handled them in _ModelRegistry._raise_for_unsupported?
if model_config.model_impl == ModelImpl.TRANSFORMERS:
if not is_transformers_impl_compatible(arch, custom_model_module):
raise ValueError(
f"The Transformers implementation of {arch} is not "
"compatible with vLLM.")
architectures[i] = "TransformersModel"
if model_config.model_impl == ModelImpl.AUTO:
if not is_transformers_impl_compatible(arch, custom_model_module):
raise ValueError(
f"{arch} has no vLLM implementation and the Transformers "
"implementation is not compatible with vLLM.")
logger.warning(
"%s has no vLLM implementation, falling back to Transformers "
"implementation. Some features may not be supported and "
"performance may not be optimal.", arch)
architectures[i] = "TransformersModel"
return architectures
def get_model_architecture(
model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:
architectures = getattr(model_config.hf_config, "architectures", [])
# Special handling for quantized Mixtral.
# FIXME(woosuk): This is a temporary hack.
mixtral_supported = [
"fp8", "compressed-tensors", "gptq_marlin", "awq_marlin"
]
if (model_config.quantization is not None
and model_config.quantization not in mixtral_supported
and "MixtralForCausalLM" in architectures):
architectures = ["QuantMixtralForCausalLM"]
vllm_supported_archs = ModelRegistry.get_supported_archs()
is_vllm_supported = any(arch in vllm_supported_archs
for arch in architectures)
if (not is_vllm_supported
or model_config.model_impl == ModelImpl.TRANSFORMERS):
architectures = resolve_transformers_fallback(model_config,
architectures)
model_cls, arch = ModelRegistry.resolve_model_cls(architectures)
if model_config.task == "embed":
model_cls = as_embedding_model(model_cls)
elif model_config.task == "classify":
model_cls = as_classification_model(model_cls)
elif model_config.task == "reward":
model_cls = as_reward_model(model_cls)
return model_cls, arch
def get_architecture_class_name(model_config: ModelConfig) -> str:
return get_model_architecture(model_config)[1]
@dataclass
class ParamMapping:
"""
A class to handle parameter mapping for model weight loading.
It creates a bidirectional mapping between packed parameters and their
constituent parts.
"""
packed_mapping: Dict[str, List[str]]
inverse_packed_mapping: Dict[str, Tuple[str,
int]] = field(default_factory=dict)
def __post_init__(self):
for packed_name, sub_params in self.packed_mapping.items():
# Skip self-contained cases (e.g., {"W_pack": ["W_pack"]})
if len(sub_params) == 1 and sub_params[0] == packed_name:
continue
for index, param_name in enumerate(sub_params):
self.inverse_packed_mapping[param_name] = (
packed_name,
index,
)
def get_sub_modules(self,
module_name: str) -> Optional[Tuple[str, List[str]]]:
for key, value in self.packed_mapping.items():
if module_name.endswith(key):
return key, value
return None
fastvideo/v1/models/loader/weight_utils.py
+2 -347
View File
@@ -6,7 +6,6 @@
"""Utilities for downloading and initializing model weights."""
import fnmatch
import glob
import hashlib
import json
import os
@@ -14,34 +13,16 @@ import tempfile
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
from typing import Generator, List, Optional, Tuple, Union
import filelock
import gguf
import huggingface_hub.constants
import numpy as np
import torch
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
from safetensors.torch import load_file, safe_open, save_file
from safetensors.torch import safe_open
from tqdm.auto import tqdm
# from vllm.config import LoadConfig, ModelConfig
from fastvideo.v1.distributed import get_tensor_model_parallel_rank
from fastvideo.v1.logger import init_logger
# from vllm.model_executor.layers.quantization import (QuantizationConfig,
# get_quantization_config)
from fastvideo.v1.platforms import current_platform
# from vllm.utils import PlaceholderModule
# try:
# from runai_model_streamer import SafetensorsStreamer
# except (ImportError, OSError):
# # see https://github.com/run-ai/runai-model-streamer/issues/26
# # OSError will be raised on arm64 platform
# runai_model_streamer = PlaceholderModule(
# "runai_model_streamer") # type: ignore[assignment]
# SafetensorsStreamer = runai_model_streamer.placeholder_attr(
# "SafetensorsStreamer")
logger = init_logger(__name__)
@@ -99,129 +80,6 @@ def _shared_pointers(tensors):
return failing
def convert_bin_to_safetensor_file(
pt_filename: str,
sf_filename: str,
) -> None:
loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
if "state_dict" in loaded:
loaded = loaded["state_dict"]
shared = _shared_pointers(loaded)
for shared_weights in shared:
for name in shared_weights[1:]:
loaded.pop(name)
# For tensors to be contiguous
loaded = {k: v.contiguous() for k, v in loaded.items()}
dirname = os.path.dirname(sf_filename)
os.makedirs(dirname, exist_ok=True)
save_file(loaded, sf_filename, metadata={"format": "pt"})
# check file size
sf_size = os.stat(sf_filename).st_size
pt_size = os.stat(pt_filename).st_size
if (sf_size - pt_size) / pt_size > 0.01:
raise RuntimeError(f"""The file size different is more than 1%:
- {sf_filename}: {sf_size}
- {pt_filename}: {pt_size}
""")
# check if the tensors are the same
reloaded = load_file(sf_filename)
for k in loaded:
pt_tensor = loaded[k]
sf_tensor = reloaded[k]
if not torch.equal(pt_tensor, sf_tensor):
raise RuntimeError(f"The output tensors do not match for key {k}")
# TODO(woosuk): Move this to other place.
# def get_quant_config(model_config: ModelConfig,
# load_config: LoadConfig) -> QuantizationConfig:
# quant_cls = get_quantization_config(model_config.quantization)
# # GGUF doesn't have config file
# if model_config.quantization == "gguf":
# return quant_cls.from_config({})
# # Read the quantization config from the HF model config, if available.
# hf_quant_config = getattr(model_config.hf_config, "quantization_config",
# None)
# # some vision model may keep quantization_config in their text_config
# hf_text_config = getattr(model_config.hf_config, "text_config", None)
# if hf_quant_config is None and hf_text_config is not None:
# hf_quant_config = getattr(hf_text_config, "quantization_config", None)
# if hf_quant_config is None:
# # compressed-tensors uses a compressions_config
# hf_quant_config = getattr(model_config.hf_config, "compression_config",
# None)
# if hf_quant_config is not None:
# return quant_cls.from_config(hf_quant_config)
# # In case of bitsandbytes/QLoRA, get quant config from the adapter model.
# if model_config.quantization == "bitsandbytes":
# if (not load_config.model_loader_extra_config
# or "qlora_adapter_name_or_path"
# not in load_config.model_loader_extra_config):
# return quant_cls.from_config({"adapter_name_or_path": ""})
# model_name_or_path = load_config.model_loader_extra_config[
# "qlora_adapter_name_or_path"]
# else:
# model_name_or_path = model_config.model
# is_local = os.path.isdir(model_name_or_path)
# if not is_local:
# # Download the config files.
# with get_lock(model_name_or_path, load_config.download_dir):
# hf_folder = snapshot_download(
# model_name_or_path,
# revision=model_config.revision,
# allow_patterns="*.json",
# cache_dir=load_config.download_dir,
# local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
# tqdm_class=DisabledTqdm,
# )
# else:
# hf_folder = model_name_or_path
# possible_config_filenames = quant_cls.get_config_filenames()
# # If the quantization config is not found, use the default config.
# if not possible_config_filenames:
# return quant_cls()
# config_files = glob.glob(os.path.join(hf_folder, "*.json"))
# quant_config_files = [
# f for f in config_files if any(
# f.endswith(x) for x in possible_config_filenames)
# ]
# if len(quant_config_files) == 0:
# raise ValueError(
# f"Cannot find the config file for {model_config.quantization}")
# if len(quant_config_files) > 1:
# raise ValueError(
# f"Found multiple config files for {model_config.quantization}: "
# f"{quant_config_files}")
# quant_config_file = quant_config_files[0]
# with open(quant_config_file) as f:
# config = json.load(f)
# if model_config.quantization == "bitsandbytes":
# config["adapter_name_or_path"] = model_name_or_path
# elif model_config.quantization == "modelopt":
# if config["producer"]["name"] == "modelopt":
# return quant_cls.from_config(config)
# else:
# raise ValueError(
# f"Unsupported quantization config"
# f" found for {model_config.quantization} in {f}.")
# return quant_cls.from_config(config)
def download_weights_from_hf(
model_name_or_path: str,
cache_dir: Optional[str],
@@ -371,53 +229,6 @@ def filter_files_not_needed_for_inference(
_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501
def np_cache_weights_iterator(
model_name_or_path: str, cache_dir: Optional[str], hf_folder: str,
hf_weights_files: List[str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""Iterate over the weights in the model np files.
Will dump the model weights to numpy files if they are not already dumped.
"""
enable_tqdm = not torch.distributed.is_initialized(
) or torch.distributed.get_rank() == 0
# Convert the model weights from torch tensors to numpy arrays for
# faster loading.
np_folder = os.path.join(hf_folder, "np")
os.makedirs(np_folder, exist_ok=True)
weight_names_file = os.path.join(np_folder, "weight_names.json")
# Use file lock to prevent multiple processes from
# dumping the same model weights to numpy at the same time.
with get_lock(model_name_or_path, cache_dir):
if not os.path.exists(weight_names_file):
weight_names: List[str] = []
for bin_file in tqdm(
hf_weights_files,
desc="Loading np_cache checkpoint shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
):
state = torch.load(bin_file,
map_location="cpu",
weights_only=True)
for name, param in state.items():
param_path = os.path.join(np_folder, name)
with open(param_path, "wb") as f:
np.save(f, param.cpu().detach().numpy())
weight_names.append(name)
with open(weight_names_file, "w") as f:
json.dump(weight_names, f)
with open(weight_names_file) as f:
weight_names = json.load(f)
for name in weight_names:
param_path = os.path.join(np_folder, name)
with open(param_path, "rb") as f:
param = np.load(f)
yield name, torch.from_numpy(param)
def safetensors_weights_iterator(
hf_weights_files: List[str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
@@ -436,23 +247,6 @@ def safetensors_weights_iterator(
yield name, param
# def runai_safetensors_weights_iterator(
# hf_weights_files: List[str]
# ) -> Generator[Tuple[str, torch.Tensor], None, None]:
# """Iterate over the weights in the model safetensor files."""
# enable_tqdm = not torch.distributed.is_initialized(
# ) or torch.distributed.get_rank() == 0
# with SafetensorsStreamer() as streamer:
# for st_file in tqdm(
# hf_weights_files,
# desc="Loading safetensors using Runai Model Streamer",
# disable=not enable_tqdm,
# bar_format=_BAR_FORMAT,
# ):
# streamer.stream_file(st_file)
# yield from streamer.get_tensors()
def pt_weights_iterator(
hf_weights_files: List[str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
@@ -470,61 +264,6 @@ def pt_weights_iterator(
del state
def get_gguf_extra_tensor_names(
gguf_file: str, gguf_to_hf_name_map: Dict[str, str]) -> List[str]:
reader = gguf.GGUFReader(gguf_file)
expected_gguf_keys = set(gguf_to_hf_name_map.keys())
exact_gguf_keys = set([tensor.name for tensor in reader.tensors])
extra_keys = expected_gguf_keys - exact_gguf_keys
return [gguf_to_hf_name_map[key] for key in extra_keys]
def gguf_quant_weights_iterator(
gguf_file: str, gguf_to_hf_name_map: Dict[str, str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""
Iterate over the quant weights in the model gguf files and convert
them to torch tensors
"""
reader = gguf.GGUFReader(gguf_file)
for tensor in reader.tensors:
if tensor.name in gguf_to_hf_name_map:
weight_type = tensor.tensor_type
name = gguf_to_hf_name_map[tensor.name]
if weight_type.name != "F32":
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:
if tensor.name in gguf_to_hf_name_map:
weight = tensor.data
weight_type = tensor.tensor_type
name = gguf_to_hf_name_map[tensor.name]
if weight_type.name != "F32":
name = name.replace("weight", "qweight")
param = torch.tensor(weight)
yield name, param
def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
"""convert PySafeSlice object from safetensors to torch.Tensor
PySafeSlice object supports indexing, which is done before loading the
actual tensor and can reduce the amount of memory being read into the
memory. However, it does not support more advanced functionalities
like `.view()` or `.t()`. Therefore, if we need to modify the loaded
tensor with these more complicated operators, we need to convert to
tensor first.
"""
if not isinstance(x, torch.Tensor):
x = x[:]
return x
def default_weight_loader(param: torch.Tensor,
loaded_weight: torch.Tensor) -> None:
"""Default weight loader."""
@@ -546,90 +285,6 @@ def default_weight_loader(param: torch.Tensor,
raise
def row_parallel_weight_loader(param: torch.Tensor,
loaded_weight: torch.Tensor) -> None:
"""Load weights that are row-parallelized."""
tp_rank = get_tensor_model_parallel_rank()
shard_dim = 0 if param.dim() != 1 else None
if shard_dim is not None:
shard_size = param.data.shape[shard_dim]
start_idx = tp_rank * shard_size
loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)
return default_weight_loader(param, loaded_weight)
LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
"""Create a weight loader that shards the weights along the given axis"""
def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
tp_rank = get_tensor_model_parallel_rank()
shard_size = param.data.shape[shard_axis]
start_idx = tp_rank * shard_size
loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
return default_weight_loader(param, loaded_weight)
return loader
def composed_weight_loader(
loader: LoaderFunction, fn: Callable[[torch.Tensor],
torch.Tensor]) -> LoaderFunction:
"""Create a weight loader that post-processes the weights after loading"""
def composed_loader(param: torch.Tensor,
loaded_weight: torch.Tensor) -> None:
loader(param, loaded_weight)
param.data.copy_(fn(param))
return
return composed_loader
def initialize_dummy_weights(
model: torch.nn.Module,
low: float = -1e-3,
high: float = 1e-3,
seed: int = 1234,
) -> None:
"""Initialize model weights with random values.
The model weights must be randomly initialized for accurate performance
measurements. Additionally, the model weights should not cause NaNs in the
forward pass. We empirically found that initializing the weights with
values between -1e-3 and 1e-3 works well for most models.
We use per-parameter random seed, so that dummy weights are consistent,
even if the model is partitioned across multiple devices. When the seed
is fixed, the random values generated by this function only depends on
the parameter's number of elements and its data type.
"""
for param in model.state_dict().values():
if torch.is_floating_point(param):
if current_platform.is_tpu():
# XLA device does not support torch.Generator()
param.uniform_(low, high)
continue
generator = torch.Generator(device=param.data.device)
generator.manual_seed(seed)
if torch.finfo(param.data.dtype).bits < 16:
# uniform_ doesn't support < 16-bit datatypes (FP8)
dtype = param.data.dtype
tmp_param = param.data.to(torch.float16)
tmp_param = tmp_param.uniform_(low, high,
generator=generator).to(dtype)
param.data.copy_(tmp_param)
else:
param.uniform_(low, high, generator=generator)
def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
"""Remap the name of FP8 k/v_scale parameters.
fastvideo/v1/models/schedulers/base.py
View File