sgl-project--sglang/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py

"""
In-place weight updates for diffusion pipeline modules.

This module provides WeightsUpdater, which swaps model weights at runtime
without restarting the server.  It is the diffusion-engine counterpart of the
LLM engine's ModelRunner.update_weights_from_disk.

Detailed usage of higher level API can be found in

/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py

Key design decisions:

- All-or-nothing with rollback: modules are updated sequentially.  If
  any module fails (shape mismatch, corrupted file, etc.), every module
  that was already updated is rolled back by reloading its weights from
  that module's last successfully-loaded weights directory.  On a full
  successful update, pipeline.model_path is updated to the new model_path;
  target_modules updates keep per-module rollback state for hybrid models.

- Rollback failures propagate: if rollback itself fails, the exception is
  not caught so the caller knows the model is in an inconsistent state.
  This matches the LLM engine behaviour.

- Offload-aware: the diffusion LayerwiseOffloadManager replaces GPU
  parameters with torch.empty((1,)) placeholders while real weights live
  in consolidated pinned CPU buffers.  A naive param.data.copy_() would
  fail with a shape mismatch.  Instead, the updater dynamically detects
  active offload managers and writes new weights directly into their CPU
  buffers via update_cpu_weights(), bypassing the placeholders entirely.
  For any layer that happens to be prefetched on GPU at update time, the
  live GPU tensor is also updated so the change takes effect immediately.
  This requires no extra GPU memory and does not disturb the offload state.

- DTensor-aware: parameters that have been distributed via
  torch.distributed.tensor are updated through distribute_tensor
  so that each shard is correctly placed on the right device mesh.
"""

from __future__ import annotations

import gc
from pathlib import Path
from typing import Any

import torch
from torch.distributed.tensor import DTensor, distribute_tensor

from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin
from sglang.multimodal_gen.runtime.loader.utils import (
    _list_safetensors_files,
    get_param_names_mapping,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import (
    safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    is_layerwise_offloaded_module,
)
from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.weight_sync.tensor_bucket import (
    FlattenedTensorBucket,
    FlattenedTensorMetadata,
)

logger = init_logger(__name__)
_DEFAULT_TENSOR_TARGET_MODULE = "transformer"


def get_updatable_modules(pipeline) -> dict[str, torch.nn.Module]:
    """Return updatable nn.Module components for the given pipeline.

    Works with both the native ComposedPipelineBase backend and the
    DiffusersPipeline wrapper.
    """
    if isinstance(pipeline, DiffusersPipeline):
        diffusers_pipe = pipeline.get_module("diffusers_pipeline")
        if diffusers_pipe is not None and diffusers_pipe.components is not None:
            raw = diffusers_pipe.components
        else:
            raw = {}
    else:
        raw = pipeline.modules
    return {n: m for n, m in raw.items() if isinstance(m, torch.nn.Module)}


def _get_weights_iter(weights_dir: str):
    """Return a (name, tensor) iterator over safetensors in weights_dir."""
    safetensors_files = _list_safetensors_files(weights_dir)
    if not safetensors_files:
        raise FileNotFoundError(f"No safetensors files found in {weights_dir}")
    return safetensors_weights_iterator(safetensors_files)


def _validate_weight_files(
    local_model_path: str,
    modules_to_update: list[tuple[str, torch.nn.Module]],
) -> tuple[dict[str, str], list[str]]:
    """Check that every module has a weights directory with safetensors files.

    Returns:
        (weights_map, missing) where weights_map maps module name to its
        weights directory and missing lists modules without weight files.
    """
    weights_map: dict[str, str] = {}
    missing: list[str] = []
    for module_name, _ in modules_to_update:
        weights_dir = Path(local_model_path) / module_name
        if weights_dir.exists() and _list_safetensors_files(str(weights_dir)):
            weights_map[module_name] = str(weights_dir)
        else:
            missing.append(module_name)
    return weights_map, missing


def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None:
    """Load weights into a module, handling offload-managed parameters.

    For offloaded modules, updates CPU buffers directly via
    update_cpu_weights(); non-offloaded parameters use in-place copy.
    """
    model_params = dict(module.named_parameters())
    weights_iter = _iter_module_weight_updates(module, weights_iter, model_params)

    offload_managers: list = []
    if is_layerwise_offloaded_module(module):
        offload_managers = [m for m in module.layerwise_offload_managers if m.enabled]

    if offload_managers:
        weight_dict = dict(weights_iter)
        offloaded_names: set[str] = set()
        for manager in offload_managers:
            offloaded_names.update(manager.update_cpu_weights(weight_dict))
        remaining = ((n, w) for n, w in weight_dict.items() if n not in offloaded_names)
        load_weights_into_model(remaining, model_params)
    else:
        load_weights_into_model(weights_iter, model_params)


def _build_module_weight_name_mapper(module: torch.nn.Module):
    """Build a chained regex mapper from mapping dicts exposed by the module."""
    mapping_fns = []
    for attr in ("lora_param_names_mapping", "param_names_mapping"):
        mapping = getattr(module, attr, None)
        if not mapping:
            continue
        mapping_fns.append(get_param_names_mapping(mapping))

    if not mapping_fns:
        return None

    def map_name(name: str) -> str:
        mapped_name = name
        for mapping_fn in mapping_fns:
            mapped_name = mapping_fn(mapped_name)[0]
        return mapped_name

    return map_name


def _iter_module_weight_updates(
    module: torch.nn.Module,
    weights_iter,
    model_params: dict,
):
    map_name = _build_module_weight_name_mapper(module)
    module_name = type(module).__name__

    for name, loaded_weight in weights_iter:
        if name in model_params:
            yield name, loaded_weight
            continue

        mapped_name = map_name(name) if map_name is not None else name
        if mapped_name in model_params:
            yield mapped_name, loaded_weight
            continue

        logger.warning(
            "Skipping weight update for %s: parameter %r not found after mapping to %r",
            module_name,
            name,
            mapped_name,
        )


def load_weights_into_model(
    weights_iter, model_params: dict, module_name: str | None = None
) -> None:
    """Copy weights from weights_iter into model_params in-place."""
    for name, loaded_weight in weights_iter:
        if name not in model_params:
            logger.warning("Skipping weight update: parameter %r not found", name)
            continue
        param = model_params[name]
        weight_loader = getattr(param, "weight_loader", None)
        if callable(weight_loader):
            weight_loader(param, loaded_weight.to(param.dtype))
        else:
            dtensor_param = param if isinstance(param, DTensor) else None
            if dtensor_param is None and isinstance(
                getattr(param, "data", None), DTensor
            ):
                dtensor_param = param.data

            if dtensor_param is not None:
                distributed_weight = distribute_tensor(
                    loaded_weight.to(param.dtype),
                    dtensor_param.device_mesh,
                    dtensor_param.placements,
                )
                dtensor_param._local_tensor.copy_(distributed_weight._local_tensor)
            else:
                if param.shape != loaded_weight.shape:
                    module_prefix = f"{module_name}." if module_name else ""
                    raise ValueError(
                        f"Shape mismatch for {module_prefix}{name}: "
                        f"model={param.shape}, loaded={loaded_weight.shape}"
                    )
                param.data.copy_(loaded_weight.to(param.dtype))


class WeightsUpdater:
    """In-place weight updates for diffusion pipeline modules.

    Args:
        pipeline: A ComposedPipelineBase (or DiffusersPipeline) instance
            whose modules will be updated.
    """

    def __init__(self, pipeline):
        self.pipeline = pipeline
        try:
            self._module_weight_dirs = pipeline._weights_updater_module_weight_dirs
        except AttributeError:
            self._module_weight_dirs = {}
            pipeline._weights_updater_module_weight_dirs = self._module_weight_dirs

    def update_weights_from_disk(
        self,
        model_path: str,
        flush_cache: bool = True,
        target_modules: list[str] | None = None,
    ) -> tuple[bool, str]:
        """Update model weights from disk without restarting the server."""
        logger.info(f"Updating weights from disk: {model_path}")

        try:
            modules_to_update = self._collect_modules(target_modules)
        except ValueError as e:
            logger.error(str(e))
            return False, str(e)

        if not modules_to_update:
            error_msg = (
                f"No matching modules found for update. "
                f"Requested: {target_modules}. "
                f"Available nn.Module(s): {list(get_updatable_modules(self.pipeline).keys())}"
            )
            logger.error(error_msg)
            return False, error_msg

        try:
            local_model_path = maybe_download_model(model_path)
        except Exception as e:
            return False, f"Failed to download model: {e}"

        weights_map, missing = _validate_weight_files(
            local_model_path, modules_to_update
        )
        if missing:
            error_msg = (
                f"Cannot update weights: missing weight files for modules: {missing}. "
                f"No partial updates allowed."
            )
            logger.error(error_msg)
            return False, error_msg

        logger.info(
            f"Updating {len(weights_map)} modules: "
            + ", ".join(f"{n} <- {p}" for n, p in weights_map.items())
        )

        success, message = self._apply_weights(modules_to_update, weights_map)

        if success:
            for module_name, _ in modules_to_update:
                self._module_weight_dirs[module_name] = weights_map[module_name]
            if target_modules is None:
                self.pipeline.model_path = local_model_path

        gc.collect()
        torch.cuda.empty_cache()

        if success and flush_cache:
            for _, module in modules_to_update:
                if isinstance(module, TeaCacheMixin):
                    module.reset_teacache_state()

        logger.info(message)
        return success, message

    def _collect_modules(
        self, target_modules: list[str] | None
    ) -> list[tuple[str, torch.nn.Module]]:
        """Resolve target_modules to (name, module) pairs.

        Raises:
            ValueError: If target_modules contains names not found in the pipeline.
        """
        components = get_updatable_modules(self.pipeline)

        if target_modules is None:
            names = list(components.keys())
        else:
            unknown = [n for n in target_modules if n not in components]
            if unknown:
                raise ValueError(
                    f"Module(s) requested for update not found in pipeline: {unknown}. "
                    f"Available Module(s): {list(components.keys())}"
                )
            names = target_modules

        return [(name, components[name]) for name in names]

    def _apply_weights(
        self,
        modules_to_update: list[tuple[str, torch.nn.Module]],
        weights_map: dict[str, str],
    ) -> tuple[bool, str]:
        """Load weights into each module; rollback on first failure."""
        updated_modules: list[str] = []

        for module_name, module in modules_to_update:
            try:
                weights_iter = _get_weights_iter(weights_map[module_name])
                _load_weights_into_module(module, weights_iter)
                updated_modules.append(module_name)
            except Exception as e:
                rollback_list = updated_modules + [module_name]
                logger.error(
                    f"Weight update failed for module '{module_name}': {e}. "
                    f"Rolling back {len(rollback_list)} module(s) "
                    f"(including partially-loaded '{module_name}'): "
                    f"{rollback_list}.",
                    exc_info=True,
                )
                self._rollback(rollback_list)
                return False, (
                    f"Failed to update module '{module_name}': {e}. "
                    f"All modules rolled back to original weights."
                )

        names = ", ".join(updated_modules)
        return True, f"Updated {len(updated_modules)} modules ({names})."

    def _rollback(self, updated_modules: list[str]) -> None:
        """Restore updated_modules to original weights.

        If rollback itself fails the exception propagates so the caller
        knows the model is in an inconsistent state.
        """
        if not updated_modules:
            return
        original_path: str | None = None
        for name in updated_modules:
            module = self.pipeline.get_module(name)
            if module is None:
                continue
            weights_dir = self._module_weight_dirs.get(name)
            if weights_dir is None:
                if original_path is None:
                    original_path = maybe_download_model(self.pipeline.model_path)
                weights_dir = str(Path(original_path) / name)
            weights_dir = Path(weights_dir)
            if not weights_dir.exists():
                continue
            weights_iter = _get_weights_iter(str(weights_dir))
            _load_weights_into_module(module, weights_iter)

    def update_weights_from_tensor(
        self,
        named_tensors: Any,
        load_format: str | None = None,
        target_modules: list[str] | None = None,
    ) -> tuple[bool, str]:
        if target_modules is None:
            target_modules = [_DEFAULT_TENSOR_TARGET_MODULE]
        try:
            modules_to_update = self._collect_modules(target_modules)
        except ValueError as e:
            logger.error(str(e))
            return False, str(e)

        if not modules_to_update:
            error_msg = (
                f"No matching modules found for update. "
                f"Requested: {target_modules}. "
                f"Available nn.Module(s): {list(get_updatable_modules(self.pipeline).keys())}"
            )
            logger.error(error_msg)
            return False, error_msg

        try:
            module_payloads = self._resolve_module_payloads(
                named_tensors=named_tensors,
                modules_to_update=modules_to_update,
            )
        except ValueError as e:
            logger.error(str(e))
            return False, str(e)

        updated_modules: list[str] = []
        for module_name, module in modules_to_update:
            try:
                payload = module_payloads[module_name]
                weights_iter = self._materialize_weights_iter(payload, load_format)
                _load_weights_into_module(module, weights_iter)
                updated_modules.append(module_name)
            except Exception as e:
                error_msg = (
                    f"Failed to update module '{module_name}' from tensor: {e}. "
                    f"The pipeline may be partially updated. "
                    f"Please discard the whole weights and reload from a known-good checkpoint."
                )
                logger.error(error_msg, exc_info=True)
                return False, error_msg

        gc.collect()
        torch.cuda.empty_cache()
        names = ", ".join(updated_modules)
        message = f"Updated {len(updated_modules)} modules from tensor ({names})."
        logger.info(message)
        return True, message

    def _resolve_module_payloads(
        self,
        named_tensors: Any,
        modules_to_update: list[tuple[str, torch.nn.Module]],
    ) -> dict[str, Any]:
        module_names = [name for name, _ in modules_to_update]
        if isinstance(named_tensors, dict):
            missing = [name for name in module_names if name not in named_tensors]
            if missing:
                raise ValueError(
                    f"Missing tensor payload for module(s): {missing}. "
                    f"Provided modules: {list(named_tensors.keys())}"
                )
            return {name: named_tensors[name] for name in module_names}

        if len(module_names) == 1:
            return {module_names[0]: named_tensors}

        raise ValueError(
            "Ambiguous tensor payload for multi-module update. "
            "Provide a dict mapping module_name -> module payload, "
            f"requested modules: {module_names}."
        )

    def _materialize_weights_iter(self, module_payload: Any, load_format: str | None):
        if load_format == "flattened_bucket":
            if not isinstance(module_payload, dict):
                raise ValueError(
                    "flattened_bucket payload must be a dict with "
                    "'flattened_tensor' and 'metadata'."
                )
            flattened_tensor = module_payload.get("flattened_tensor")
            metadata = module_payload.get("metadata")
            if flattened_tensor is None or metadata is None:
                raise ValueError(
                    "flattened_bucket payload missing 'flattened_tensor' or 'metadata'."
                )
            return self._reconstruct_from_flattened_bucket(flattened_tensor, metadata)

        if isinstance(module_payload, (list, tuple)):
            return iter(module_payload)

        raise ValueError(
            f"Unsupported module payload type for load_format={load_format}: "
            f"{type(module_payload).__name__}"
        )

    def _reconstruct_from_flattened_bucket(self, flattened_tensor: Any, metadata: Any):
        if not isinstance(flattened_tensor, torch.Tensor):
            raise ValueError(
                "flattened_bucket 'flattened_tensor' must be a torch.Tensor."
            )
        if not isinstance(metadata, list):
            raise ValueError("flattened_bucket 'metadata' must be a list.")

        converted_metadata: list[FlattenedTensorMetadata] = []
        for meta in metadata:
            converted_metadata.append(
                FlattenedTensorMetadata(
                    name=meta.name,
                    shape=torch.Size(meta.shape),
                    dtype=self._normalize_torch_dtype(meta.dtype),
                    start_idx=int(meta.start_idx),
                    end_idx=int(meta.end_idx),
                    numel=int(meta.numel),
                )
            )

        bucket = FlattenedTensorBucket(
            flattened_tensor=flattened_tensor,
            metadata=converted_metadata,
        )
        return bucket.reconstruct_tensors()

    def _normalize_torch_dtype(self, dtype: Any) -> torch.dtype:
        if isinstance(dtype, torch.dtype):
            return dtype
        if isinstance(dtype, str):
            name = dtype.split(".")[-1]
            normalized = getattr(torch, name, None)
            if isinstance(normalized, torch.dtype):
                return normalized
        raise ValueError(f"Unsupported dtype in flattened_bucket metadata: {dtype!r}")