310 lines
12 KiB
Python
310 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from collections.abc import Generator, Iterable
|
|
from dataclasses import dataclass, field
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.utils import WeightsMapper
|
|
from vllm.model_executor.utils import get_packed_modules_mapping
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
MODEL_OPT_SCALE_SUFFIXES = (
|
|
".input_scale",
|
|
".weight_scale",
|
|
".weight_scale_2",
|
|
".weight_scale_inv",
|
|
)
|
|
DEFAULT_PACKED_MODULES_MAPPING = {
|
|
"to_qkv": ("to_q", "to_k", "to_v"),
|
|
"add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
|
|
"w13": ("w1", "w3"),
|
|
}
|
|
FP8_DTYPES = tuple(
|
|
dtype
|
|
for dtype in (
|
|
getattr(torch, "float8_e4m3fn", None),
|
|
getattr(torch, "float8_e5m2", None),
|
|
getattr(torch, "float8_e4m3fnuz", None),
|
|
getattr(torch, "float8_e5m2fnuz", None),
|
|
)
|
|
if dtype is not None
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _AdaptState:
|
|
scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict)
|
|
pending_weights: dict[str, list[tuple[str, str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
|
|
skipped_scales: int = 0
|
|
dequantized_weights: int = 0
|
|
|
|
|
|
class ModelOptFp8CheckpointAdapter:
|
|
def __init__(self, model: nn.Module, source: object):
|
|
self._loadable_tensors = self._get_model_loadable_tensors(model)
|
|
self._weights_mapper = self._get_weights_mapper(model)
|
|
self._checkpoint_key_mapper = getattr(model, "remap_checkpoint_key", None)
|
|
self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model"
|
|
|
|
@classmethod
|
|
def is_compatible(
|
|
cls,
|
|
source: object,
|
|
quant_config: object | None,
|
|
use_safetensors: bool,
|
|
) -> bool:
|
|
return use_safetensors and cls._is_transformer_source(source) and cls._is_checkpoint_quant_config(quant_config)
|
|
|
|
@staticmethod
|
|
def _is_transformer_source(source: object) -> bool:
|
|
if getattr(source, "subfolder", None) == "transformer":
|
|
return True
|
|
return str(getattr(source, "prefix", "")).startswith("transformer.")
|
|
|
|
@staticmethod
|
|
def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
|
|
return (
|
|
quant_config is not None
|
|
and hasattr(quant_config, "get_name")
|
|
and quant_config.get_name() == "modelopt"
|
|
and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
|
|
)
|
|
|
|
@staticmethod
|
|
def _get_model_loadable_tensors(model: nn.Module) -> dict[str, torch.Tensor]:
|
|
loadable_tensors: dict[str, torch.Tensor] = {name: param for name, param in model.named_parameters()}
|
|
loadable_tensors.update({name: buffer for name, buffer in model.named_buffers()})
|
|
return loadable_tensors
|
|
|
|
@staticmethod
|
|
def _is_scale(name: str) -> bool:
|
|
return name.endswith(MODEL_OPT_SCALE_SUFFIXES)
|
|
|
|
@staticmethod
|
|
def _is_fp8_tensor(tensor: torch.Tensor) -> bool:
|
|
return tensor.dtype in FP8_DTYPES
|
|
|
|
@staticmethod
|
|
def _get_weight_scale_name(weight_name: str) -> str | None:
|
|
if weight_name.endswith(".weight"):
|
|
return weight_name[: -len(".weight")] + ".weight_scale"
|
|
return None
|
|
|
|
@classmethod
|
|
def _get_weights_mapper(cls, model: nn.Module) -> WeightsMapper:
|
|
mapping = {
|
|
packed_name: tuple(shard_names) for packed_name, shard_names in DEFAULT_PACKED_MODULES_MAPPING.items()
|
|
}
|
|
mapping.update(
|
|
{
|
|
str(packed_name): tuple(str(shard_name) for shard_name in shard_names)
|
|
for packed_name, shard_names in get_packed_modules_mapping(model).items()
|
|
}
|
|
)
|
|
|
|
orig_to_new_substr = {".to_out.0.": ".to_out."}
|
|
orig_to_new_prefix: dict[str, str] = {}
|
|
for packed_name, shard_names in mapping.items():
|
|
for shard_name in shard_names:
|
|
orig_to_new_substr[f".{shard_name}."] = f".{packed_name}."
|
|
orig_to_new_prefix[f"{shard_name}."] = f"{packed_name}."
|
|
return WeightsMapper(
|
|
orig_to_new_substr=orig_to_new_substr,
|
|
orig_to_new_prefix=orig_to_new_prefix,
|
|
)
|
|
|
|
def _resolve_target_and_output_names(self, name: str) -> tuple[str | None, str]:
|
|
if name in self._loadable_tensors:
|
|
return name, name
|
|
|
|
if callable(self._checkpoint_key_mapper):
|
|
candidate = self._checkpoint_key_mapper(name)
|
|
if candidate in self._loadable_tensors:
|
|
return candidate, candidate
|
|
|
|
for candidate in self._weights_mapper.apply_list([name]):
|
|
if candidate != name and candidate in self._loadable_tensors:
|
|
# Preserve shard names so downstream packed-weight loaders can route them.
|
|
return candidate, name
|
|
return None, name
|
|
|
|
@staticmethod
|
|
def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor:
|
|
if scale.numel() == 1:
|
|
return scale.reshape(())
|
|
if len(weight_shape) == 2 and scale.ndim == 1 and scale.shape[0] == weight_shape[0]:
|
|
return scale.reshape(-1, 1)
|
|
if tuple(scale.shape) == tuple(weight_shape):
|
|
return scale
|
|
if (
|
|
len(weight_shape) == 2
|
|
and scale.ndim == 4
|
|
and scale.shape[1] == 1
|
|
and scale.shape[3] == 1
|
|
and weight_shape[0] % scale.shape[0] == 0
|
|
and weight_shape[1] % scale.shape[2] == 0
|
|
):
|
|
block_n = weight_shape[0] // scale.shape[0]
|
|
block_k = weight_shape[1] // scale.shape[2]
|
|
return scale.expand(scale.shape[0], block_n, scale.shape[2], block_k).reshape(weight_shape)
|
|
raise ValueError(f"Unsupported ModelOpt FP8 weight_scale shape {tuple(scale.shape)} for weight {weight_shape}")
|
|
|
|
def _dequantize_weight(
|
|
self,
|
|
name: str,
|
|
loaded_weight: torch.Tensor,
|
|
state: _AdaptState,
|
|
target_dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
scale_name = self._get_weight_scale_name(name)
|
|
if scale_name is None or scale_name not in state.scale_tensors:
|
|
raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weight {name!r}")
|
|
|
|
scale = state.scale_tensors[scale_name].to(dtype=target_dtype, device=loaded_weight.device)
|
|
scale = self._reshape_weight_scale(scale, loaded_weight.shape)
|
|
return loaded_weight.to(dtype=target_dtype) * scale
|
|
|
|
def _flush_pending_weights(
|
|
self,
|
|
scale_name: str,
|
|
state: _AdaptState,
|
|
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
for (
|
|
weight_name,
|
|
output_name,
|
|
weight_tensor,
|
|
target_dtype,
|
|
) in state.pending_weights.pop(scale_name, []):
|
|
yield output_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
|
|
state.dequantized_weights += 1
|
|
|
|
def _handle_scale_tensor(
|
|
self,
|
|
name: str,
|
|
output_name: str,
|
|
tensor: torch.Tensor,
|
|
target_name: str | None,
|
|
state: _AdaptState,
|
|
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
state.scale_tensors[name] = tensor
|
|
if target_name is None:
|
|
state.skipped_scales += 1
|
|
else:
|
|
yield output_name, tensor
|
|
yield from self._flush_pending_weights(name, state)
|
|
|
|
def _target_dtype_for_dequantization(
|
|
self,
|
|
tensor: torch.Tensor,
|
|
target_name: str | None,
|
|
) -> torch.dtype | None:
|
|
if target_name is None or not self._is_fp8_tensor(tensor):
|
|
return None
|
|
|
|
target_dtype = self._loadable_tensors[target_name].dtype
|
|
if target_dtype in FP8_DTYPES:
|
|
return None
|
|
return target_dtype
|
|
|
|
def _maybe_dequantize_or_defer_weight(
|
|
self,
|
|
name: str,
|
|
output_name: str,
|
|
tensor: torch.Tensor,
|
|
target_dtype: torch.dtype,
|
|
state: _AdaptState,
|
|
) -> torch.Tensor | None:
|
|
scale_name = self._get_weight_scale_name(name)
|
|
if scale_name is None:
|
|
raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}")
|
|
|
|
if scale_name not in state.scale_tensors:
|
|
state.pending_weights.setdefault(scale_name, []).append((name, output_name, tensor, target_dtype))
|
|
return None
|
|
|
|
state.dequantized_weights += 1
|
|
return self._dequantize_weight(name, tensor, state, target_dtype)
|
|
|
|
@staticmethod
|
|
def _check_pending_weights(state: _AdaptState) -> None:
|
|
if not state.pending_weights:
|
|
return
|
|
|
|
missing_scale_names = ", ".join(repr(name) for name in sorted(state.pending_weights))
|
|
raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weights: {missing_scale_names}")
|
|
|
|
def _log_adaptation_summary(self, state: _AdaptState) -> None:
|
|
if not state.skipped_scales and not state.dequantized_weights:
|
|
return
|
|
|
|
logger.info_once(
|
|
"Adapted ModelOpt FP8 %s weights: dequantized %d full-precision weights, skipped %d scale tensors",
|
|
self._source_label,
|
|
state.dequantized_weights,
|
|
state.skipped_scales,
|
|
)
|
|
|
|
def adapt(
|
|
self,
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
state = _AdaptState()
|
|
|
|
for name, tensor in weights:
|
|
target_name, output_name = self._resolve_target_and_output_names(name)
|
|
if self._is_scale(name):
|
|
yield from self._handle_scale_tensor(name, output_name, tensor, target_name, state)
|
|
continue
|
|
|
|
target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
|
|
if target_dtype is not None:
|
|
tensor = self._maybe_dequantize_or_defer_weight(name, output_name, tensor, target_dtype, state)
|
|
if tensor is None:
|
|
continue
|
|
yield output_name, tensor
|
|
|
|
self._check_pending_weights(state)
|
|
self._log_adaptation_summary(state)
|
|
|
|
|
|
class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
|
|
_PRE_QUANT_SCALE_SUFFIX = ".pre_quant_scale"
|
|
|
|
@staticmethod
|
|
def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
|
|
return (
|
|
quant_config is not None
|
|
and hasattr(quant_config, "get_name")
|
|
and quant_config.get_name() == "modelopt_fp4"
|
|
and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
|
|
)
|
|
|
|
def adapt(
|
|
self,
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
def validated_weights() -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
for name, tensor in weights:
|
|
if name.endswith(self._PRE_QUANT_SCALE_SUFFIX):
|
|
raise ValueError(
|
|
f"ModelOpt NVFP4 checkpoint tensor {name!r} is unsupported: "
|
|
"vLLM 0.25.0 does not consume pre_quant_scale. Export the checkpoint "
|
|
"with pre-quant scales folded into the weights."
|
|
)
|
|
yield name, tensor
|
|
|
|
yield from super().adapt(validated_weights())
|
|
|
|
|
|
class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
|
|
@staticmethod
|
|
def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
|
|
return (
|
|
quant_config is not None
|
|
and hasattr(quant_config, "get_name")
|
|
and quant_config.get_name() == "modelopt_mixed"
|
|
)
|