235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from collections.abc import Sequence
|
|
from typing import TYPE_CHECKING, Literal
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
LinearBase,
|
|
ReplicatedLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.models.transformers.utils import init_on_device_without_buffers
|
|
from vllm.model_executor.models.utils import maybe_prefix
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.quantization import build_quant_config
|
|
|
|
if TYPE_CHECKING:
|
|
from transformers import PretrainedConfig, PreTrainedModel
|
|
from transformers.models.auto.auto_factory import _BaseAutoModelClass
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizationConfig,
|
|
)
|
|
|
|
|
|
Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]
|
|
|
|
|
|
def make_attention_mask(
|
|
hidden_states: torch.Tensor,
|
|
seq_lengths: list[int],
|
|
) -> torch.Tensor | None:
|
|
"""Build a padding mask only when the batch contains padding.
|
|
|
|
``seq_lengths`` must be non-empty and contain valid sequence lengths.
|
|
Returns a boolean mask of shape ``(len(seq_lengths), max(seq_lengths))``
|
|
on the device of ``hidden_states``, with True marking valid tokens,
|
|
or None when all sequence lengths are equal.
|
|
|
|
Passing an all-true CUDA mask makes attention backends inspect a device
|
|
scalar to select the dense path, which introduces a host-device sync.
|
|
"""
|
|
max_seq_len = max(seq_lengths)
|
|
if all(seq_len == max_seq_len for seq_len in seq_lengths):
|
|
return None
|
|
|
|
attention_mask = hidden_states.new_zeros(
|
|
len(seq_lengths),
|
|
max_seq_len,
|
|
dtype=torch.bool,
|
|
)
|
|
for i, seq_len in enumerate(seq_lengths):
|
|
attention_mask[i, :seq_len] = True
|
|
return attention_mask
|
|
|
|
|
|
def replace_linear_class(
|
|
linear: nn.Linear,
|
|
style: Style = "replicate",
|
|
quant_config: QuantizationConfig | None = None,
|
|
*,
|
|
prefix: str = "",
|
|
) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
|
|
"""
|
|
Replace nn.Linear with one of vLLM's tensor parallel linear classes.
|
|
|
|
Args:
|
|
linear: `nn.Linear` to be replaced.
|
|
style: Tensor parallel style of the new linear, e.g. "colwise".
|
|
quant_config: Quantization config for the new linear.
|
|
Returns:
|
|
The new linear.
|
|
"""
|
|
|
|
if not isinstance(style, str):
|
|
raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")
|
|
|
|
vllm_linear_maps = {
|
|
"colwise": (ColumnParallelLinear, {}),
|
|
"colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
|
|
"rowwise": (RowParallelLinear, {}),
|
|
"rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
|
|
"replicate": (ReplicatedLinear, {}),
|
|
}
|
|
vllm_linear_cls, vllm_linear_kwargs = vllm_linear_maps[style]
|
|
|
|
return vllm_linear_cls(
|
|
input_size=linear.in_features,
|
|
output_size=linear.out_features,
|
|
bias=linear.bias is not None,
|
|
quant_config=quant_config,
|
|
prefix=prefix,
|
|
return_bias=False,
|
|
**vllm_linear_kwargs,
|
|
)
|
|
|
|
|
|
def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig):
|
|
"""Recursively replace modules in the model as needed.
|
|
Currently, this replaces:
|
|
- `nn.Linear` with vLLM's tensor parallel linear classes
|
|
"""
|
|
# Prefix the patterns because we always start from `self.model`
|
|
quant_config = build_quant_config(od_config.quantization_config)
|
|
|
|
def _recursive_replace(module: nn.Module, prefix: str):
|
|
for child_name, child_module in module.named_children():
|
|
new_module = child_module
|
|
qual_name = maybe_prefix(prefix, child_name)
|
|
# Replace modules as needed
|
|
if isinstance(child_module, nn.Linear):
|
|
style = "replicate"
|
|
new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name)
|
|
else:
|
|
_recursive_replace(child_module, prefix=qual_name)
|
|
if new_module is not child_module:
|
|
setattr(module, child_name, new_module)
|
|
|
|
_recursive_replace(model, prefix="")
|
|
|
|
|
|
def recursive_replace_linear_with_quantization_config(
|
|
model: nn.Module,
|
|
quant_config: QuantizationConfig | None,
|
|
prefix: str = "",
|
|
skip_modules: Sequence[str] = (),
|
|
) -> None:
|
|
for child_name, child_module in model.named_children():
|
|
qual_name = maybe_prefix(prefix, child_name)
|
|
if any(qual_name == path or qual_name.startswith(path + ".") for path in skip_modules):
|
|
continue
|
|
if isinstance(child_module, nn.Linear):
|
|
setattr(
|
|
model,
|
|
child_name,
|
|
replace_linear_class(child_module, quant_config=quant_config, prefix=qual_name),
|
|
)
|
|
else:
|
|
recursive_replace_linear_with_quantization_config(
|
|
child_module, quant_config, prefix=qual_name, skip_modules=skip_modules
|
|
)
|
|
|
|
|
|
def init_parameters(
|
|
module: nn.Module,
|
|
dtype: torch.dtype | None,
|
|
device: torch.device | None = None,
|
|
):
|
|
for name, param in module.named_parameters(recurse=False):
|
|
if param.device == torch.device("meta"):
|
|
new_param = nn.Parameter(
|
|
torch.empty_like(
|
|
param.data,
|
|
dtype=dtype,
|
|
device=device,
|
|
),
|
|
requires_grad=param.requires_grad,
|
|
)
|
|
setattr(module, name, new_param)
|
|
for child in module.children():
|
|
init_parameters(child, dtype, device)
|
|
|
|
|
|
def init_parameters_preserving_vllm_linear(
|
|
module: nn.Module,
|
|
dtype: torch.dtype | None,
|
|
device: torch.device | None = None,
|
|
) -> None:
|
|
if isinstance(module, LinearBase):
|
|
return
|
|
for name, param in module.named_parameters(recurse=False):
|
|
if param.is_meta:
|
|
setattr(
|
|
module,
|
|
name,
|
|
nn.Parameter(torch.empty_like(param, dtype=dtype, device=device), requires_grad=param.requires_grad),
|
|
)
|
|
for child in module.children():
|
|
init_parameters_preserving_vllm_linear(child, dtype, device)
|
|
|
|
|
|
def create_transformers_model(
|
|
auto_cls: _BaseAutoModelClass,
|
|
od_config: OmniDiffusionConfig,
|
|
hf_config: PretrainedConfig,
|
|
dtype: torch.dtype | None = None,
|
|
device: torch.device | None = None,
|
|
) -> PreTrainedModel:
|
|
"""Create a HuggingFace model using the given auto class and model name."""
|
|
dtype = dtype or od_config.dtype
|
|
device = device or torch.get_default_device()
|
|
with init_on_device_without_buffers("meta"):
|
|
model = auto_cls.from_config(hf_config)
|
|
recursive_replace_linear(model, od_config)
|
|
init_parameters(model, dtype=dtype, device=device)
|
|
return model
|
|
|
|
|
|
def create_transformers_model_with_vllm_linears(
|
|
auto_cls: _BaseAutoModelClass,
|
|
hf_config: PretrainedConfig,
|
|
quant_config: QuantizationConfig | None,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
prefix: str = "",
|
|
skip_modules: Sequence[str] = (),
|
|
) -> PreTrainedModel:
|
|
"""Build a Transformers model with vLLM linears."""
|
|
with init_on_device_without_buffers("meta"):
|
|
model = auto_cls.from_config(hf_config, dtype=dtype)
|
|
recursive_replace_linear_with_quantization_config(model, quant_config, prefix=prefix, skip_modules=skip_modules)
|
|
init_parameters_preserving_vllm_linear(model, dtype=dtype, device=device)
|
|
return model
|
|
|
|
|
|
def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
|
|
"""Load a JSON config file from a local path or HuggingFace Hub repo."""
|
|
if local_files_only:
|
|
path = os.path.join(model_path, *filename.split("/"))
|
|
with open(path) as f:
|
|
return json.load(f)
|
|
else:
|
|
from vllm_omni.transformers_utils.repo_utils import hf_api
|
|
|
|
cached = hf_api().hf_hub_download(repo_id=model_path, filename=filename)
|
|
with open(cached) as f:
|
|
return json.load(f)
|