260 lines
8.1 KiB
Python
260 lines
8.1 KiB
Python
import threading
|
|
from contextlib import contextmanager, nullcontext
|
|
from dataclasses import dataclass
|
|
from typing import Iterator, Optional, Union
|
|
|
|
import torch
|
|
from torch.distributed.fsdp import MixedPrecisionPolicy
|
|
|
|
from sglang.multimodal_gen.runtime.platforms import current_platform
|
|
from sglang.multimodal_gen.runtime.utils.precision_types import PRECISION_TO_TYPE
|
|
|
|
|
|
def precision_to_dtype(precision: str, field_name: str = "precision") -> torch.dtype:
|
|
try:
|
|
return PRECISION_TO_TYPE[precision]
|
|
except KeyError as exc:
|
|
raise ValueError(
|
|
f"Unsupported {field_name}={precision!r}; "
|
|
f"expected one of {sorted(PRECISION_TO_TYPE)}"
|
|
) from exc
|
|
|
|
|
|
def resolve_precision(
|
|
server_args,
|
|
component_or_precision_attr: str,
|
|
*,
|
|
precision_attr: Optional[str] = None,
|
|
field_name: Optional[str] = None,
|
|
) -> torch.dtype:
|
|
component_precision = server_args.component_precisions.get(
|
|
component_or_precision_attr
|
|
)
|
|
if component_precision is not None:
|
|
return precision_to_dtype(
|
|
component_precision,
|
|
f"component_precisions.{component_or_precision_attr}",
|
|
)
|
|
precision_attr = precision_attr or component_or_precision_attr
|
|
precision = getattr(server_args.pipeline_config, precision_attr)
|
|
return precision_to_dtype(precision, field_name or precision_attr)
|
|
|
|
|
|
def resolve_decode_precision(
|
|
server_args,
|
|
component_name: str = "vae",
|
|
*,
|
|
quality: str | None = None,
|
|
) -> torch.dtype:
|
|
component_precision = server_args.component_precisions.get(component_name)
|
|
if component_precision is not None:
|
|
return precision_to_dtype(
|
|
component_precision, f"component_precisions.{component_name}"
|
|
)
|
|
|
|
pipeline_config = server_args.pipeline_config
|
|
if component_name in ("audio_vae", "vocoder"):
|
|
return resolve_precision(
|
|
server_args,
|
|
component_name,
|
|
precision_attr="audio_vae_precision",
|
|
)
|
|
|
|
if quality == "high":
|
|
high_precision = getattr(pipeline_config, "vae_decode_precision_high", None)
|
|
if high_precision is not None:
|
|
return precision_to_dtype(high_precision, "vae_decode_precision_high")
|
|
|
|
decode_precision = getattr(pipeline_config, "vae_decode_precision", None)
|
|
if decode_precision is not None:
|
|
return precision_to_dtype(decode_precision, "vae_decode_precision")
|
|
return resolve_precision(
|
|
server_args,
|
|
component_name,
|
|
precision_attr="vae_precision",
|
|
)
|
|
|
|
|
|
def resolve_component_precision_override(
|
|
server_args, module_name: str
|
|
) -> Optional[torch.dtype]:
|
|
exact_precision = server_args.component_precisions.get(module_name)
|
|
if exact_precision is None:
|
|
return None
|
|
return precision_to_dtype(exact_precision, f"component_precisions.{module_name}")
|
|
|
|
|
|
def resolve_component_precision(server_args, module_name: str) -> Optional[torch.dtype]:
|
|
exact_precision = resolve_component_precision_override(server_args, module_name)
|
|
if exact_precision is not None:
|
|
return exact_precision
|
|
|
|
pipeline_config = server_args.pipeline_config
|
|
|
|
if module_name in ("audio_vae", "vocoder"):
|
|
precision_attr = "audio_vae_precision"
|
|
elif module_name in ("vae", "video_vae", "diffusion_decoder"):
|
|
precision_attr = "vae_precision"
|
|
elif module_name in (
|
|
"transformer",
|
|
"transformer_2",
|
|
"audio_dit",
|
|
"video_dit",
|
|
"connectors",
|
|
"dual_tower_bridge",
|
|
):
|
|
precision_attr = "dit_precision"
|
|
elif module_name == "image_encoder":
|
|
precision_attr = "image_encoder_precision"
|
|
elif module_name == "text_encoder" or module_name.startswith("text_encoder_"):
|
|
precisions = getattr(pipeline_config, "text_encoder_precisions", None)
|
|
if not precisions:
|
|
return None
|
|
suffix = module_name.removeprefix("text_encoder")
|
|
index = 0 if suffix == "" else int(suffix.removeprefix("_")) - 1
|
|
if index < 0 or index >= len(precisions):
|
|
raise ValueError(
|
|
f"No configured precision for {module_name!r}; "
|
|
f"text_encoder_precisions has {len(precisions)} entries"
|
|
)
|
|
precision = precisions[index]
|
|
return precision_to_dtype(precision, f"text_encoder_precisions[{index}]")
|
|
else:
|
|
return None
|
|
|
|
if not hasattr(pipeline_config, precision_attr):
|
|
return None
|
|
return resolve_precision(server_args, precision_attr)
|
|
|
|
|
|
def autocast_enabled(dtype: torch.dtype, disable_autocast: bool) -> bool:
|
|
return (
|
|
dtype != torch.float32
|
|
and not disable_autocast
|
|
and current_platform.is_amp_supported()
|
|
)
|
|
|
|
|
|
def autocast_enabled_for_device(
|
|
tensor: torch.Tensor, dtype: torch.dtype, disable_autocast: bool
|
|
) -> bool:
|
|
return tensor.device.type == current_platform.device_type and autocast_enabled(
|
|
dtype, disable_autocast
|
|
)
|
|
|
|
|
|
def autocast_context(
|
|
dtype: torch.dtype,
|
|
disable_autocast: bool,
|
|
*,
|
|
enabled: Optional[bool] = None,
|
|
):
|
|
autocast_is_enabled = (
|
|
autocast_enabled(dtype, disable_autocast) if enabled is None else enabled
|
|
)
|
|
if not autocast_is_enabled and current_platform.is_mps():
|
|
return nullcontext()
|
|
return torch.autocast(
|
|
device_type=current_platform.device_type,
|
|
dtype=dtype,
|
|
enabled=autocast_is_enabled,
|
|
)
|
|
|
|
|
|
def get_module_dtype(module, default: torch.dtype = torch.float32) -> torch.dtype:
|
|
try:
|
|
return next(module.parameters()).dtype
|
|
except (AttributeError, StopIteration):
|
|
dtype = getattr(module, "dtype", None)
|
|
return dtype if isinstance(dtype, torch.dtype) else default
|
|
|
|
|
|
def align_tensor_to_module_dtype(
|
|
tensor: torch.Tensor,
|
|
module,
|
|
*,
|
|
device: Optional[Union[torch.device, str]] = None,
|
|
default_dtype: torch.dtype = torch.float32,
|
|
) -> torch.Tensor:
|
|
dtype = get_module_dtype(module, default=default_dtype)
|
|
if device is None:
|
|
try:
|
|
device = next(module.parameters()).device
|
|
except (AttributeError, StopIteration):
|
|
device = tensor.device
|
|
if not tensor.is_floating_point():
|
|
return tensor.to(device=device)
|
|
return tensor.to(device=device, dtype=dtype)
|
|
|
|
|
|
@contextmanager
|
|
def temporary_module_dtype(
|
|
module,
|
|
dtype: torch.dtype,
|
|
*,
|
|
enabled: bool = True,
|
|
restore_dtype: Optional[torch.dtype] = None,
|
|
) -> Iterator:
|
|
if not enabled:
|
|
yield module
|
|
return
|
|
|
|
original_dtype = restore_dtype or get_module_dtype(module)
|
|
module = module.to(dtype=dtype)
|
|
try:
|
|
yield module
|
|
finally:
|
|
module.to(dtype=original_dtype)
|
|
|
|
|
|
@dataclass
|
|
class MixedPrecisionState:
|
|
param_dtype: torch.dtype | None = None
|
|
reduce_dtype: torch.dtype | None = None
|
|
output_dtype: torch.dtype | None = None
|
|
compute_dtype: torch.dtype | None = None
|
|
mp_policy: MixedPrecisionPolicy | None = None
|
|
|
|
|
|
class _MixedPrecisionContext(threading.local):
|
|
state: MixedPrecisionState | None = None
|
|
|
|
|
|
_mixed_precision_state = _MixedPrecisionContext()
|
|
|
|
|
|
def get_mixed_precision_state() -> MixedPrecisionState:
|
|
"""Get the current mixed precision state."""
|
|
state = _mixed_precision_state.state
|
|
if state is None:
|
|
raise ValueError("Mixed precision state not set")
|
|
return state
|
|
|
|
|
|
def set_mixed_precision_policy(
|
|
param_dtype: torch.dtype,
|
|
reduce_dtype: torch.dtype,
|
|
output_dtype: torch.dtype | None = None,
|
|
mp_policy: MixedPrecisionPolicy | None = None,
|
|
):
|
|
"""Set mixed precision policy for the current thread.
|
|
|
|
Args:
|
|
param_dtype: Parameter dtype used for training
|
|
reduce_dtype: Reduction dtype used for gradients
|
|
output_dtype: Optional output dtype
|
|
"""
|
|
state = MixedPrecisionState(
|
|
param_dtype=param_dtype,
|
|
reduce_dtype=reduce_dtype,
|
|
output_dtype=output_dtype,
|
|
mp_policy=mp_policy,
|
|
)
|
|
_mixed_precision_state.state = state
|
|
|
|
|
|
def get_compute_dtype() -> torch.dtype:
|
|
"""Get the current compute dtype from mixed precision policy."""
|
|
state = _mixed_precision_state.state
|
|
return torch.get_default_dtype() if state is None else state.param_dtype
|