diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py index 4e3f540ad..1fe0f8b0a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py @@ -17,6 +17,23 @@ class UpdateWeightFromDiskReqInput: target_modules: list[str] | None = None +@dataclass +class UpdateWeightFromTensorReqInput: + """Request to update model weights from tensor payloads for diffusion models.""" + + serialized_named_tensors: list[str | bytes] + load_format: str | None = None + target_modules: list[str] | None = None + + +@dataclass +class UpdateWeightFromTensorCheckerReqInput: + """Request to verify live module weights against expected SHA-256 values.""" + + target_module: str + expected_named_tensors_sha256: dict[str, str] + + @dataclass class GetWeightsChecksumReqInput: """Compute SHA-256 checksum of loaded module weights for verification.""" diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py index 7bc0054f7..cd17a4943 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py @@ -5,6 +5,8 @@ from fastapi import APIRouter, Request from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( GetWeightsChecksumReqInput, UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.srt.utils.json_response import orjson_response @@ -29,6 +31,82 @@ async def update_weights_from_disk(request: Request): target_modules=body.get("target_modules"), ) + try: + response = await async_scheduler_client.forward(req) + except Exception as e: + return orjson_response( + {"success": False, "message": str(e)}, + status_code=500, + ) + + result = response.output + return orjson_response( + result, + status_code=200 if result["success"] else 400, + ) + + +@router.post("/update_weights_from_tensor") +async def update_weights_from_tensor(request: Request): + """Update model weights from serialized tensor payloads.""" + body = await request.json() + serialized_named_tensors = body.get("serialized_named_tensors") + if not serialized_named_tensors: + return orjson_response( + {"success": False, "message": "serialized_named_tensors is required"}, + status_code=400, + ) + + req = UpdateWeightFromTensorReqInput( + serialized_named_tensors=serialized_named_tensors, + load_format=body.get("load_format"), + target_modules=body.get("target_modules"), + ) + + try: + response = await async_scheduler_client.forward(req) + except Exception as e: + return orjson_response( + {"success": False, "message": str(e)}, + status_code=500, + ) + + result = response.output + return orjson_response( + result, + status_code=200 if result["success"] else 400, + ) + + +@router.post("/update_weights_from_tensor_checker") +async def update_weights_from_tensor_checker(request: Request): + """Verify live module weights against expected SHA-256 values.""" + body = await request.json() + target_module = body.get("target_module") + if not target_module: + return orjson_response( + {"success": False, "message": "target_module is required"}, + status_code=400, + ) + + expected_named_tensors_sha256 = body.get("expected_named_tensors_sha256") + if ( + not isinstance(expected_named_tensors_sha256, dict) + or not expected_named_tensors_sha256 + ): + return orjson_response( + { + "success": False, + "message": "expected_named_tensors_sha256 is required", + }, + status_code=400, + ) + + req = UpdateWeightFromTensorCheckerReqInput( + target_module=target_module, + expected_named_tensors_sha256=expected_named_tensors_sha256, + ) + try: response = await async_scheduler_client.forward(req) except Exception as e: diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 145611797..5498bdd35 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -37,14 +37,8 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import ( post_process_sample, save_outputs, ) -from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum -from sglang.multimodal_gen.runtime.loader.weights_updater import ( - WeightsUpdater, - get_updatable_modules, -) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( configure_layerwise_offload_modules, - iter_materialized_weights, ) from sglang.multimodal_gen.runtime.pipelines_core import ( ComposedPipelineBase, @@ -54,6 +48,9 @@ from sglang.multimodal_gen.runtime.pipelines_core import ( ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.post_training.gpu_worker_post_training_mixin import ( + GPUWorkerPostTrainingMixin, +) from sglang.multimodal_gen.runtime.realtime.session import ( RealtimeSessionCache, ) @@ -102,7 +99,7 @@ class _ExpandedOutputParts: trajectory_decoded_parts: list[list[torch.Tensor]] | None = None -class GPUWorker: +class GPUWorker(GPUWorkerPostTrainingMixin): """ A worker that executes the model on a single GPU. """ @@ -894,48 +891,6 @@ class GPUWorker: status = self.pipeline.get_lora_status() return OutputBatch(output=status) - 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 inplace without restarting the server.""" - if not self.pipeline: - return False, "Pipeline is not initialized" - - updater = WeightsUpdater(self.pipeline) - success, message = updater.update_weights_from_disk( - model_path, - flush_cache=flush_cache, - target_modules=target_modules, - ) - if success: - self.server_args.model_path = model_path - self.pipeline.model_path = model_path - return success, message - - def get_weights_checksum( - self, module_names: list[str] | None = None - ) -> dict[str, str]: - """Compute SHA-256 checksum of each module's weights.""" - if not self.pipeline: - return {"error": "Pipeline is not initialized"} - - all_modules = get_updatable_modules(self.pipeline) - names = module_names if module_names is not None else list(all_modules.keys()) - - checksums: dict[str, str] = {} - for name in names: - module = all_modules.get(name) - if module is None: - checksums[name] = "not_found" - continue - checksums[name] = compute_weights_checksum( - iter_materialized_weights(module) - ) - return checksums - OOM_MSG = """ OOM detected. Possible solutions: diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index d46e76f2f..ce380808d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -20,6 +20,8 @@ from sglang.multimodal_gen.runtime.distributed import get_world_group from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( GetWeightsChecksumReqInput, UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( GetDisaggStatsReq, @@ -44,6 +46,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( BatchMetricsWindow, OutputBatch, ) +from sglang.multimodal_gen.runtime.post_training.scheduler_post_training_mixin import ( + SchedulerPostTrainingMixin, +) from sglang.multimodal_gen.runtime.server_args import ( PortArgs, ServerArgs, @@ -69,7 +74,7 @@ _MAX_RECV_REQS_PER_POLL = 1024 _BATCH_METRICS_LOG_INTERVAL = 5 -class Scheduler(SchedulerDisaggMixin): +class Scheduler(SchedulerPostTrainingMixin, SchedulerDisaggMixin): """ Runs the main event loop for the rank 0 worker. It listens for external requests via ZMQ and coordinates with other workers. @@ -132,6 +137,10 @@ class Scheduler(SchedulerDisaggMixin): ReleaseRealtimeSessionReq: self._handle_release_realtime_session, GetDisaggStatsReq: self._handle_get_disagg_stats, UpdateWeightFromDiskReqInput: self._handle_update_weights_from_disk, + UpdateWeightFromTensorReqInput: self._handle_update_weights_from_tensor, + UpdateWeightFromTensorCheckerReqInput: ( + self._handle_update_weights_from_tensor_checker + ), GetWeightsChecksumReqInput: self._handle_get_weights_checksum, } @@ -211,25 +220,6 @@ class Scheduler(SchedulerDisaggMixin): req = reqs[0] return self.worker.release_realtime_session(req.session_id) - def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: - """Handle update_weights_from_disk request for RL workflows.""" - req = reqs[0] - success, message = self.worker.update_weights_from_disk( - model_path=req.model_path, - flush_cache=req.flush_cache, - target_modules=req.target_modules, - ) - return OutputBatch( - output={"success": success, "message": message}, - error=None if success else message, - ) - - def _handle_get_weights_checksum(self, reqs: List[Any]) -> OutputBatch: - """Handle get_weights_checksum request.""" - req = reqs[0] - checksums = self.worker.get_weights_checksum(module_names=req.module_names) - return OutputBatch(output=checksums) - @staticmethod def _normalize_generation_reqs(reqs: list[Any]) -> list[Req]: if len(reqs) == 1 and isinstance(reqs[0], list): diff --git a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py new file mode 100644 index 000000000..ee0045d87 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.multimodal_gen.runtime.distributed import get_tp_rank, get_tp_world_size +from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + iter_materialized_weights, +) +from sglang.multimodal_gen.runtime.post_training.tensor_update_checker import ( + TensorUpdateChecker, +) +from sglang.multimodal_gen.runtime.post_training.weights_updater import ( + WeightsUpdater, + get_updatable_modules, +) +from sglang.srt.utils import MultiprocessingSerializer +from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, + ) + + +class GPUWorkerPostTrainingMixin: + def update_weights_from_disk( + self, + model_path: str, + flush_cache: bool = True, + target_modules: list[str] | None = None, + ) -> tuple[bool, str]: + if not self.pipeline: + return False, "Pipeline is not initialized" + + updater = WeightsUpdater(self.pipeline) + success, message = updater.update_weights_from_disk( + model_path, + flush_cache=flush_cache, + target_modules=target_modules, + ) + if success: + self.server_args.model_path = model_path + self.pipeline.model_path = model_path + return success, message + + def update_weights_from_tensor( + self, + req: UpdateWeightFromTensorReqInput, + ) -> tuple[bool, str]: + if not self.pipeline: + return False, "Pipeline is not initialized" + + payload, error = self._select_rank_scoped_payload( + payloads=req.serialized_named_tensors, + field_name="serialized_named_tensors", + ) + if error is not None: + return False, error + + monkey_patch_torch_reductions() + try: + named_tensors = MultiprocessingSerializer.deserialize(payload) + except Exception as e: + return False, f"Failed to deserialize serialized_named_tensors: {e}" + + updater = WeightsUpdater(self.pipeline) + return updater.update_weights_from_tensor( + named_tensors=named_tensors, + load_format=req.load_format, + target_modules=req.target_modules, + ) + + def update_weights_from_tensor_checker( + self, + req: UpdateWeightFromTensorCheckerReqInput, + ) -> tuple[bool, str]: + if not self.pipeline: + return False, "Pipeline is not initialized" + + checker = TensorUpdateChecker(self.pipeline) + result = checker.verify_across_tp( + target_module=req.target_module, + expected_named_tensors_sha256=req.expected_named_tensors_sha256, + tp_rank=get_tp_rank(), + tp_world_size=get_tp_world_size(), + tp_cpu_group=self.tp_cpu_group, + tp_root_rank=self.tp_group.first_rank, + ) + if self.sp_group.world_size == 1: + return result + + import torch + + is_sp_root = self.sp_group.rank_in_group == 0 + gathered_results = [None] * self.sp_group.world_size if is_sp_root else None + torch.distributed.gather_object( + result, + gathered_results, + dst=self.sp_group.first_rank, + group=self.sp_cpu_group, + ) + + final_result = None + if is_sp_root: + failures = [ + (rank, message) + for rank, (success, message) in enumerate(gathered_results) + if not success + ] + if failures: + rank, message = failures[0] + if len(failures) == 1: + final_result = (False, f"SP rank {rank}: {message}") + else: + final_result = ( + False, + f"{len(failures)} SP ranks failed update_weight_from_tensor_checker; " + f"first failure on rank {rank}: {message}", + ) + else: + final_result = result + + final_result_holder = [final_result] + torch.distributed.broadcast_object_list( + final_result_holder, + src=self.sp_group.first_rank, + group=self.sp_cpu_group, + ) + return final_result_holder[0] + + def get_weights_checksum( + self, module_names: list[str] | None = None + ) -> dict[str, str]: + if not self.pipeline: + return {"error": "Pipeline is not initialized"} + + all_modules = get_updatable_modules(self.pipeline) + names = module_names if module_names is not None else list(all_modules.keys()) + + checksums: dict[str, str] = {} + for name in names: + module = all_modules.get(name) + if module is None: + checksums[name] = "not_found" + continue + checksums[name] = compute_weights_checksum( + iter_materialized_weights(module) + ) + return checksums + + def _select_rank_scoped_payload( + self, + payloads: list, + field_name: str, + ) -> tuple[object | None, str | None]: + if not isinstance(payloads, list): + return None, f"{field_name} must be a list" + if not payloads: + return None, f"{field_name} is required" + + tp_world_size = get_tp_world_size() + if len(payloads) not in (1, tp_world_size): + return ( + None, + f"{field_name} size must be 1 or tp_size ({tp_world_size}), " + f"got {len(payloads)}", + ) + + payload_idx = get_tp_rank() if len(payloads) == tp_world_size else 0 + return payloads[payload_idx], None diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py new file mode 100644 index 000000000..025643b09 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import Any, List + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch + + +class SchedulerPostTrainingMixin: + def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + success, message = self.worker.update_weights_from_disk( + model_path=req.model_path, + flush_cache=req.flush_cache, + target_modules=req.target_modules, + ) + return OutputBatch( + output={"success": success, "message": message}, + error=None if success else message, + ) + + def _handle_update_weights_from_tensor(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + success, message = self.worker.update_weights_from_tensor(req) + if self.server_args.tp_size > 1: + import torch + + torch.distributed.barrier(group=self.worker.tp_cpu_group) + return OutputBatch( + output={"success": success, "message": message}, + error=None if success else message, + ) + + def _handle_update_weights_from_tensor_checker( + self, reqs: List[Any] + ) -> OutputBatch: + req = reqs[0] + success, message = self.worker.update_weights_from_tensor_checker(req) + return OutputBatch( + output={"success": success, "message": message}, + error=None if success else message, + ) + + def _handle_get_weights_checksum(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + checksums = self.worker.get_weights_checksum(module_names=req.module_names) + return OutputBatch(output=checksums) diff --git a/python/sglang/multimodal_gen/runtime/post_training/tensor_update_checker.py b/python/sglang/multimodal_gen/runtime/post_training/tensor_update_checker.py new file mode 100644 index 000000000..94eaace55 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/tensor_update_checker.py @@ -0,0 +1,263 @@ +"""Verification helpers for diffusion update_weights_from_tensor workflows.""" + +from __future__ import annotations + +import hashlib +from collections.abc import Iterable + +import torch +from torch.distributed.tensor import DTensor + +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + iter_materialized_weights, +) + +_MAX_DISPLAY_TENSORS = 5 + + +def _materialize_local_tensor(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, DTensor): + tensor = tensor._local_tensor + return tensor.detach().cpu().contiguous() + + +def compute_tensor_sha256(tensor: torch.Tensor) -> str: + tensor = _materialize_local_tensor(tensor) + hasher = hashlib.sha256() + hasher.update(str(tensor.dtype).encode("utf-8")) + hasher.update(repr(tuple(tensor.shape)).encode("utf-8")) + hasher.update(tensor.view(torch.uint8).numpy().tobytes()) + return hasher.hexdigest() + + +def build_named_tensor_sha256( + named_tensors: Iterable[tuple[str, torch.Tensor]], +) -> dict[str, str]: + sha256_by_name: dict[str, str] = {} + for name, tensor in named_tensors: + sha256_by_name[name] = compute_tensor_sha256(tensor) + return sha256_by_name + + +class TensorUpdateChecker: + def __init__(self, pipeline): + self.pipeline = pipeline + + def verify_across_tp( + self, + target_module: str, + expected_named_tensors_sha256: dict[str, str], + tp_rank: int, + tp_world_size: int, + tp_cpu_group, + tp_root_rank: int, + ) -> tuple[bool, str]: + if tp_world_size == 1: + return self.verify( + target_module=target_module, + expected_named_tensors_sha256=expected_named_tensors_sha256, + ) + + module = self.pipeline.get_module(target_module) + if module is None: + return False, f"Module '{target_module}' is not initialized" + + local_named_tensors = dict( + self._iter_module_named_tensors( + module, expected_named_tensors_sha256.keys() + ) + ) + reference_tensors = dict(module.named_parameters()) + reference_tensors.update(dict(module.named_buffers())) + actual_named_tensors_sha256: dict[str, str] | None = ( + {} if tp_rank == 0 else None + ) + + for name, expected_sha256 in expected_named_tensors_sha256.items(): + gathered_tensors: list[torch.Tensor | None] | None = ( + [None] * tp_world_size if tp_rank == 0 else None + ) + torch.distributed.gather_object( + ( + _materialize_local_tensor(local_named_tensors[name]) + if name in local_named_tensors + else None + ), + gathered_tensors, + dst=tp_root_rank, + group=tp_cpu_group, + ) + if tp_rank != 0: + continue + + valid_tensors = [ + tensor for tensor in gathered_tensors if tensor is not None + ] + if len(valid_tensors) != len(gathered_tensors): + continue + + local_sha256s = [compute_tensor_sha256(tensor) for tensor in valid_tensors] + if all(local_sha256 == expected_sha256 for local_sha256 in local_sha256s): + actual_named_tensors_sha256[name] = expected_sha256 + continue + + reference_tensor = reference_tensors.get(name) + candidate_dims: list[int] = [] + if isinstance(reference_tensor, DTensor): + for placement in reference_tensor.placements: + shard_dim = getattr(placement, "dim", None) + if isinstance(shard_dim, int) and shard_dim not in candidate_dims: + candidate_dims.append(shard_dim) + for attr in ("input_dim", "output_dim"): + shard_dim = getattr(reference_tensor, attr, None) + if isinstance(shard_dim, int) and shard_dim not in candidate_dims: + candidate_dims.append(shard_dim) + + reconstructed_sha256 = None + first_tensor = valid_tensors[0] + for shard_dim in candidate_dims: + if first_tensor.ndim == 0: + break + + shard_dim %= first_tensor.ndim + compatible = True + for tensor in valid_tensors[1:]: + if ( + tensor.ndim != first_tensor.ndim + or tensor.dtype != first_tensor.dtype + ): + compatible = False + break + if any( + lhs != rhs + for dim, (lhs, rhs) in enumerate( + zip(first_tensor.shape, tensor.shape) + ) + if dim != shard_dim + ): + compatible = False + break + if not compatible: + continue + + reconstructed = torch.cat(valid_tensors, dim=shard_dim).contiguous() + if compute_tensor_sha256(reconstructed) == expected_sha256: + reconstructed_sha256 = expected_sha256 + break + + actual_named_tensors_sha256[name] = reconstructed_sha256 or local_sha256s[0] + + final_result: tuple[bool, str] | None = None + if tp_rank == 0: + final_result = self._compare_manifests( + target_module=target_module, + expected_named_tensors_sha256=expected_named_tensors_sha256, + actual_named_tensors_sha256=actual_named_tensors_sha256, + ) + if final_result[0]: + final_result = ( + True, + f"Verified module '{target_module}' update across {tp_world_size} TP ranks.", + ) + + final_result_holder = [final_result] + torch.distributed.broadcast_object_list( + final_result_holder, + src=tp_root_rank, + group=tp_cpu_group, + ) + final_result = final_result_holder[0] + assert final_result is not None + return final_result + + def verify( + self, + target_module: str, + expected_named_tensors_sha256: dict[str, str], + ) -> tuple[bool, str]: + module = self.pipeline.get_module(target_module) + if module is None: + return False, f"Module '{target_module}' is not initialized" + + actual_named_tensors_sha256 = build_named_tensor_sha256( + self._iter_module_named_tensors( + module, expected_named_tensors_sha256.keys() + ) + ) + return self._compare_manifests( + target_module=target_module, + expected_named_tensors_sha256=expected_named_tensors_sha256, + actual_named_tensors_sha256=actual_named_tensors_sha256, + ) + + def _iter_module_named_tensors( + self, + module: torch.nn.Module, + expected_names: Iterable[str], + ): + expected_name_set = set(expected_names) + seen_names: set[str] = set() + + for name, tensor in iter_materialized_weights(module): + if name not in expected_name_set: + continue + seen_names.add(name) + yield name, tensor + + for name, tensor in module.named_buffers(): + if name in seen_names or name not in expected_name_set: + continue + seen_names.add(name) + yield name, tensor + + def _compare_manifests( + self, + *, + target_module: str, + expected_named_tensors_sha256: dict[str, str], + actual_named_tensors_sha256: dict[str, str], + ) -> tuple[bool, str]: + missing_names = sorted( + name + for name in expected_named_tensors_sha256 + if name not in actual_named_tensors_sha256 + ) + mismatched_names = sorted( + name + for name, expected_sha256 in expected_named_tensors_sha256.items() + if name in actual_named_tensors_sha256 + and actual_named_tensors_sha256[name] != expected_sha256 + ) + + if missing_names or mismatched_names: + parts: list[str] = [] + if missing_names: + parts.append( + "missing " + f"{len(missing_names)} tensor(s): " + f"{self._format_tensor_names(missing_names)}" + ) + if mismatched_names: + parts.append( + "checksum mismatch for " + f"{len(mismatched_names)} tensor(s): " + f"{self._format_tensor_names(mismatched_names)}" + ) + return ( + False, + f"Module '{target_module}' update weight check failed: " + + "; ".join(parts), + ) + + return ( + True, + f"Verified module '{target_module}' update for " + f"{len(expected_named_tensors_sha256)} tensor(s).", + ) + + def _format_tensor_names(self, names: list[str]) -> str: + displayed = names[:_MAX_DISPLAY_TENSORS] + formatted = ", ".join(displayed) + if len(names) > _MAX_DISPLAY_TENSORS: + formatted += f", ... (+{len(names) - _MAX_DISPLAY_TENSORS} more)" + return formatted diff --git a/python/sglang/multimodal_gen/runtime/loader/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py similarity index 55% rename from python/sglang/multimodal_gen/runtime/loader/weights_updater.py rename to python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 1967fb2d2..dce327c13 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -42,6 +42,7 @@ from __future__ import annotations import gc from pathlib import Path +from typing import Any import torch from torch.distributed.tensor import DTensor, distribute_tensor @@ -49,6 +50,7 @@ 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, @@ -59,8 +61,13 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im 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]: @@ -115,6 +122,9 @@ def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None: 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] @@ -125,30 +135,92 @@ def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None: 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, dict(module.named_parameters())) + load_weights_into_model(remaining, model_params) else: - load_weights_into_model(weights_iter, dict(module.named_parameters())) + load_weights_into_model(weights_iter, model_params) -def load_weights_into_model(weights_iter, model_params: dict) -> None: +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] - if param.shape != loaded_weight.shape: - raise ValueError( - f"Shape mismatch for {name}: model={param.shape}, loaded={loaded_weight.shape}" - ) - if isinstance(param, DTensor): - distributed_weight = distribute_tensor( - loaded_weight.to(param.dtype), - param.device_mesh, - param.placements, - ) - param._local_tensor.copy_(distributed_weight._local_tensor) + weight_loader = getattr(param, "weight_loader", None) + if callable(weight_loader): + weight_loader(param, loaded_weight.to(param.dtype)) else: - param.data.copy_(loaded_weight.to(param.dtype)) + 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: @@ -293,3 +365,142 @@ class WeightsUpdater: 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}")