wkutakandRahul Steiger
43a476c1e5
Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster> Signed-off-by: Wojciech Kutak <wkutak@nvidia.com> Co-authored-by: Rahul Steiger <rsteiger@nvidia.com>
107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Configuration parsing and denoising-step policy for mixed precision."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Literal, cast
|
|
|
|
ReasonerPolicy = Literal["native", "a16"]
|
|
|
|
_REASONER_POLICIES = frozenset({"native", "a16"})
|
|
_CONFIG_FIELDS = frozenset({"enabled", "first_steps", "last_steps", "reasoner"})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Cosmos3MixedPrecisionConfig:
|
|
"""Validated precision policy shared by every quantization strategy."""
|
|
|
|
first_steps: int = 3
|
|
last_steps: int = 3
|
|
reasoner: ReasonerPolicy = "a16"
|
|
|
|
@classmethod
|
|
def from_additional_config(
|
|
cls,
|
|
additional_config: Mapping[str, object] | None,
|
|
) -> Cosmos3MixedPrecisionConfig | None:
|
|
"""Parse Cosmos3 fields from vLLM-Omni's additional configuration."""
|
|
_, config = cls.resolve_additional_config(additional_config)
|
|
return config
|
|
|
|
@classmethod
|
|
def resolve_additional_config(
|
|
cls,
|
|
additional_config: Mapping[str, object] | None,
|
|
) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]:
|
|
"""Return whether an override was supplied and its effective policy.
|
|
|
|
The separate presence bit distinguishes an absent runtime override from
|
|
an explicit disable. That distinction matters when a checkpoint owns a
|
|
default policy.
|
|
"""
|
|
if additional_config is not None and not isinstance(additional_config, Mapping):
|
|
raise TypeError("additional_config must be a mapping")
|
|
values = additional_config or {}
|
|
if "cosmos3_mixed_precision" not in values:
|
|
return False, None
|
|
raw_config = values["cosmos3_mixed_precision"]
|
|
if not isinstance(raw_config, Mapping):
|
|
raise TypeError("cosmos3_mixed_precision must be a mapping")
|
|
unknown = set(raw_config) - _CONFIG_FIELDS
|
|
if unknown:
|
|
raise ValueError(f"Unknown cosmos3_mixed_precision fields: {sorted(unknown)}")
|
|
|
|
enabled = raw_config.get("enabled", True)
|
|
if not isinstance(enabled, bool):
|
|
raise TypeError("cosmos3_mixed_precision.enabled must be a boolean")
|
|
if not enabled:
|
|
conflicting = set(raw_config) - {"enabled"}
|
|
if conflicting:
|
|
raise ValueError(
|
|
"cosmos3_mixed_precision.enabled=false cannot be combined with "
|
|
f"schedule fields: {sorted(conflicting)}"
|
|
)
|
|
return True, None
|
|
|
|
first_steps = _non_negative_int(
|
|
raw_config.get("first_steps", 3),
|
|
"cosmos3_mixed_precision.first_steps",
|
|
)
|
|
last_steps = _non_negative_int(
|
|
raw_config.get("last_steps", 3),
|
|
"cosmos3_mixed_precision.last_steps",
|
|
)
|
|
|
|
reasoner = str(raw_config.get("reasoner", "a16")).lower()
|
|
if reasoner not in _REASONER_POLICIES:
|
|
raise ValueError(
|
|
f"cosmos3_mixed_precision.reasoner must be one of {sorted(_REASONER_POLICIES)}, got {reasoner!r}"
|
|
)
|
|
config = cls(
|
|
first_steps=first_steps,
|
|
last_steps=last_steps,
|
|
reasoner=cast(ReasonerPolicy, reasoner),
|
|
)
|
|
if first_steps == 0 and last_steps == 0 and reasoner == "native":
|
|
return True, None
|
|
return True, config
|
|
|
|
def use_high_precision(self, step_index: int, num_steps: int) -> bool:
|
|
"""Select the high path for an actual scheduler-step index."""
|
|
if num_steps <= 0:
|
|
raise ValueError(f"num_steps must be positive, got {num_steps}")
|
|
if step_index < 0 or step_index >= num_steps:
|
|
raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}")
|
|
return step_index < self.first_steps or step_index >= num_steps - self.last_steps
|
|
|
|
|
|
def _non_negative_int(value: object, name: str) -> int:
|
|
"""Validate an integer configuration field without accepting booleans."""
|
|
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
|
|
raise TypeError(f"{name} must be a non-negative integer")
|
|
return value
|