chore: import upstream snapshot with attribution

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

python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""RL-specific dataclasses used by post-training and rollout paths."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class RolloutSessionData:
+    """Per-batch rollout state created by prepare_rollout(), lives on the batch object.
+
+    Cleared by setting ``batch._rollout_session_data = None``.
+    """
+
+    pipeline_config: Any = None
+    sigma_max: float = 0.0
+    latents_shape: tuple | None = None
+    noise_buffer: torch.Tensor | None = None
+
+    local_log_prob_sum: list[torch.Tensor] = field(default_factory=list)
+    local_log_prob_count: list[torch.Tensor] = field(default_factory=list)
+
+    local_variance_noises: list[torch.Tensor] = field(default_factory=list)
+    local_prev_sample_means: list[torch.Tensor] = field(default_factory=list)
+    local_noise_std_devs: list[torch.Tensor] = field(default_factory=list)
+    local_model_outputs: list[torch.Tensor] = field(default_factory=list)
+
+
+@dataclass
+class RolloutDebugTensors:
+    """Container for rollout debug tensors collected during denoising."""
+
+    rollout_variance_noises: torch.Tensor | None = None
+    rollout_prev_sample_means: torch.Tensor | None = None
+    rollout_noise_std_devs: torch.Tensor | None = None
+    rollout_model_outputs: torch.Tensor | None = None
+
+
+@dataclass
+class RolloutDenoisingEnv:
+    image_kwargs: dict[str, Any] | None = None
+    pos_cond_kwargs: dict[str, Any] | None = None
+    neg_cond_kwargs: dict[str, Any] | None = None
+    guidance: torch.Tensor | None = None
+
+
+@dataclass
+class RolloutDitTrajectory:
+    # [B, T+1, ...]: per-step noisy latents x_{t_0..t_{T-1}} followed by the
+    # final denoised latent x_{t_T} (last scheduler.step output).
+    latents: torch.Tensor | None = None
+    timesteps: torch.Tensor | None = None  # [T]
+
+
+@dataclass
+class RolloutTrajectoryData:
+    rollout_log_probs: torch.Tensor | None = None
+    rollout_debug_tensors: RolloutDebugTensors | None = None
+    denoising_env: RolloutDenoisingEnv | None = None
+    dit_trajectory: RolloutDitTrajectory | None = None

python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py

@@ -0,0 +1,206 @@
+"""Mixin for rollout-related denoising hooks.
+
+Moved out of DenoisingStage to keep the core stage lean.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDenoisingEnv,
+    RolloutDitTrajectory,
+    RolloutTrajectoryData,
+)
+from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import (
+    SchedulerRLMixin,
+)
+from sglang.multimodal_gen.runtime.post_training.sp_utils import (
+    gather_stacked_latents_for_sp,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def _kwargs_to_cpu(d: Any) -> Any:
+    if isinstance(d, torch.Tensor):
+        return d.detach().cpu()
+    if isinstance(d, dict):
+        return {k: _kwargs_to_cpu(v) for k, v in d.items()}
+    if isinstance(d, list):
+        return [_kwargs_to_cpu(v) for v in d]
+    if isinstance(d, tuple):
+        return tuple(_kwargs_to_cpu(v) for v in d)
+    return d
+
+
+class RolloutDenoisingMixin:
+
+    def _maybe_prepare_rollout(self, batch: Req):
+        """Prepare denoising loop for rollout."""
+        if not isinstance(self.scheduler, SchedulerRLMixin):
+            if batch.rollout:
+                raise ValueError(
+                    f"Scheduler {type(self.scheduler)} does not support rollout"
+                )
+            return
+
+        self.scheduler.release_rollout_resources(batch)
+        if batch.rollout:
+            self.scheduler.prepare_rollout(
+                batch=batch,
+                pipeline_config=self.server_args.pipeline_config,
+            )
+
+    def _maybe_collect_rollout_log_probs(self, batch: Req):
+        if not isinstance(self.scheduler, SchedulerRLMixin):
+            if batch.rollout:
+                raise ValueError(
+                    f"Scheduler {type(self.scheduler)} does not support rollout"
+                )
+            return
+
+        if batch.rollout:
+            if batch.rollout_trajectory_data is None:
+                batch.rollout_trajectory_data = RolloutTrajectoryData()
+            batch.rollout_trajectory_data.rollout_log_probs = (
+                self.scheduler.collect_rollout_log_probs(batch)
+            )
+            if batch.rollout_debug_mode:
+                batch.rollout_trajectory_data.rollout_debug_tensors = (
+                    self.scheduler.collect_rollout_debug_tensors(batch)
+                )
+            self.scheduler.release_rollout_resources(batch)
+
+    def _postprocess_rollout_outputs(
+        self,
+        batch: Req,
+        latents: torch.Tensor,
+        num_inference_steps: int,
+        final_timestep: torch.Tensor,
+        server_args: ServerArgs,
+    ) -> None:
+        """Finalize rollout-only outputs.
+
+        Must be called before ``_post_denoising_loop`` so that ``latents`` (the
+        last ``scheduler.step`` output) is still SP-sharded and can be gathered
+        uniformly with the per-step trajectory latents.
+        """
+        self._maybe_collect_rollout_log_probs(batch)
+        # Append final denoised latent as the (T+1)-th entry (step_index=T),
+        # routed through the same filter so rollout_return_step_indices can
+        # include/exclude it.
+        self._maybe_append_dit_trajectory_step(
+            batch=batch,
+            latents=latents,
+            timestep_value=final_timestep,
+            step_index=num_inference_steps,
+        )
+        self._maybe_finalize_denoising_env_collection(
+            batch=batch,
+            pipeline_config=server_args.pipeline_config,
+        )
+
+    def _maybe_init_denoising_env_collection(
+        self,
+        batch,
+        pipeline_config,
+        image_kwargs: dict[str, Any],
+        pos_cond_kwargs: dict[str, Any],
+        neg_cond_kwargs: dict[str, Any],
+        guidance: torch.Tensor | None,
+    ) -> None:
+        collect_env = batch.rollout_return_denoising_env
+        collect_traj = batch.rollout_return_dit_trajectory
+        if not (collect_env or collect_traj):
+            batch._rollout_denoising_env_state = None
+            return
+
+        if collect_env:
+            env = RolloutDenoisingEnv(
+                image_kwargs=_kwargs_to_cpu(image_kwargs),
+                pos_cond_kwargs=_kwargs_to_cpu(pos_cond_kwargs),
+                neg_cond_kwargs=(
+                    _kwargs_to_cpu(neg_cond_kwargs) if neg_cond_kwargs else None
+                ),
+                guidance=guidance.detach().cpu() if guidance is not None else None,
+            )
+            pos_src = pos_cond_kwargs
+            neg_src = neg_cond_kwargs
+        else:
+            env = None
+            pos_src = None
+            neg_src = None
+
+        batch._rollout_denoising_env_state = {
+            "env": env,
+            "step_latents": [],
+            "step_timesteps": [],
+            "pos_cond_kwargs_src": pos_src,
+            "neg_cond_kwargs_src": neg_src,
+        }
+
+    def _maybe_append_dit_trajectory_step(
+        self,
+        batch,
+        latents: torch.Tensor,
+        timestep_value: torch.Tensor,
+        step_index: int,
+    ) -> None:
+        if not batch.rollout or not batch.rollout_return_dit_trajectory:
+            return
+        state = getattr(batch, "_rollout_denoising_env_state", None)
+        if state is None:
+            return
+
+        return_step_indices = getattr(batch, "rollout_return_step_indices", None)
+        if return_step_indices is not None and step_index not in return_step_indices:
+            return
+
+        state["step_latents"].append(latents.detach())
+        state["step_timesteps"].append(timestep_value.detach().cpu())
+
+    def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None:
+        state = getattr(batch, "_rollout_denoising_env_state", None)
+        if state is None:
+            return
+
+        env: RolloutDenoisingEnv | None = state["env"]
+        step_latents: list[torch.Tensor] = state["step_latents"]
+        step_timesteps: list[torch.Tensor] = state["step_timesteps"]
+
+        if batch.rollout_trajectory_data is None:
+            batch.rollout_trajectory_data = RolloutTrajectoryData()
+
+        if step_latents and batch.rollout_return_dit_trajectory:
+            step_latents_tensor = torch.stack(step_latents, dim=1)
+            step_latents_tensor = gather_stacked_latents_for_sp(
+                pipeline_config=pipeline_config,
+                batch=batch,
+                stacked_latents=step_latents_tensor,
+            )
+            batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory(
+                latents=step_latents_tensor.cpu(),
+                timesteps=torch.stack(step_timesteps, dim=0).cpu(),
+            )
+
+        if env is not None and batch.rollout_return_denoising_env:
+            gather_fn = getattr(
+                pipeline_config, "gather_denoising_env_static_for_sp", None
+            )
+
+            pos_src = state.get("pos_cond_kwargs_src")
+            if pos_src is not None and env.pos_cond_kwargs is not None:
+                gathered_pos = gather_fn(batch, pos_src) if gather_fn else pos_src
+                env.pos_cond_kwargs = _kwargs_to_cpu(gathered_pos)
+
+            neg_src = state.get("neg_cond_kwargs_src")
+            if neg_src is not None and env.neg_cond_kwargs is not None:
+                gathered_neg = gather_fn(batch, neg_src) if gather_fn else neg_src
+                env.neg_cond_kwargs = _kwargs_to_cpu(gathered_neg)
+
+            batch.rollout_trajectory_data.denoising_env = env
+
+        batch._rollout_denoising_env_state = None

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)

python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Debug tensor helpers for rollout-enabled schedulers."""
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDebugTensors,
+    RolloutSessionData,
+)
+
+
+class SchedulerRLDebugMixin:
+    @staticmethod
+    def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None:
+        rollout_session_data.local_variance_noises = []
+        rollout_session_data.local_prev_sample_means = []
+        rollout_session_data.local_noise_std_devs = []
+        rollout_session_data.local_model_outputs = []
+
+    def append_local_rollout_debug_tensors(
+        self,
+        batch,
+        *,
+        variance_noise: torch.Tensor,
+        prev_sample_mean: torch.Tensor,
+        noise_std_dev: torch.Tensor,
+        model_output: torch.Tensor,
+    ) -> None:
+        rollout_session_data = batch._rollout_session_data
+        batch_size = variance_noise.shape[0]
+        # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused.
+        rollout_session_data.local_variance_noises.append(variance_noise.clone())
+        rollout_session_data.local_prev_sample_means.append(prev_sample_mean)
+        rollout_session_data.local_noise_std_devs.append(
+            noise_std_dev.expand((batch_size, 1))
+        )
+        rollout_session_data.local_model_outputs.append(model_output)
+
+    def consume_local_rollout_debug_tensors(
+        self,
+        batch,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        rollout_session_data = batch._rollout_session_data
+        variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1)
+        prev_sample_means = torch.stack(
+            rollout_session_data.local_prev_sample_means, dim=1
+        )
+        noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1)
+        model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1)
+        self._reset_rollout_debug_tensors(rollout_session_data)
+        return variance_noises, prev_sample_means, noise_std_devs, model_outputs
+
+    def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors:
+        """
+        Consume rollout debug tensors and merge for all SP ranks.
+
+        Returns rollout debug tensors with shape [B, T, ...].
+        """
+        rollout_session_data = batch._rollout_session_data
+        variance_noises, prev_sample_means, noise_std_devs, model_outputs = (
+            self.consume_local_rollout_debug_tensors(batch)
+        )
+
+        if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False):
+            variance_noises = variance_noises.to(get_local_torch_device())
+            prev_sample_means = prev_sample_means.to(get_local_torch_device())
+            noise_std_devs = noise_std_devs.to(get_local_torch_device())
+            model_outputs = model_outputs.to(get_local_torch_device())
+            pipeline_config = rollout_session_data.pipeline_config
+            bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1]
+
+            # [B, T, ...] -> [B*T, ...]
+            variance_noises_packed = variance_noises.contiguous().reshape(
+                bsz * num_steps, *variance_noises.shape[2:]
+            )
+            prev_sample_means_packed = prev_sample_means.contiguous().reshape(
+                bsz * num_steps, *prev_sample_means.shape[2:]
+            )
+            model_outputs_packed = model_outputs.contiguous().reshape(
+                bsz * num_steps, *model_outputs.shape[2:]
+            )
+
+            # Gather on packed tensors first.
+            variance_noises_packed = pipeline_config.gather_latents_for_sp(
+                variance_noises_packed, batch=batch
+            )
+            prev_sample_means_packed = pipeline_config.gather_latents_for_sp(
+                prev_sample_means_packed, batch=batch
+            )
+            model_outputs_packed = pipeline_config.gather_latents_for_sp(
+                model_outputs_packed, batch=batch
+            )
+
+            # Unpack back to [B, T, ...].
+            variance_noises = variance_noises_packed.reshape(
+                bsz, num_steps, *variance_noises_packed.shape[1:]
+            )
+            prev_sample_means = prev_sample_means_packed.reshape(
+                bsz, num_steps, *prev_sample_means_packed.shape[1:]
+            )
+            model_outputs = model_outputs_packed.reshape(
+                bsz, num_steps, *model_outputs_packed.shape[1:]
+            )
+            # noise_std_devs is same on every device, not a sharded latent tensor.
+
+        return RolloutDebugTensors(
+            rollout_variance_noises=variance_noises.cpu(),
+            rollout_prev_sample_means=prev_sample_means.cpu(),
+            rollout_noise_std_devs=noise_std_devs.cpu(),
+            rollout_model_outputs=model_outputs.cpu(),
+        )

python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flow-matching rollout step utilities for log-prob computation."""
+
+import math
+from typing import Any, Union
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutSessionData,
+)
+from sglang.multimodal_gen.runtime.post_training.scheduler_rl_debug_mixin import (
+    SchedulerRLDebugMixin,
+)
+
+_LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi))
+
+
+class SchedulerRLMixin(SchedulerRLDebugMixin):
+    @staticmethod
+    def _get_rollout_session_data(batch) -> RolloutSessionData:
+        """Return the RolloutSessionData attached to *batch*, or raise if not prepared."""
+        rollout_session_data = getattr(batch, "_rollout_session_data", None)
+        if rollout_session_data is None:
+            raise RuntimeError("prepare_rollout() not called before rollout")
+        return rollout_session_data
+
+    def release_rollout_resources(self, batch) -> None:
+        """Release rollout-owned resources. Call when denoising ends or before a new rollout."""
+        batch._rollout_session_data = None
+
+    def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None:
+        """Enable rollout and set SDE/CPS params. Call once before the denoising loop."""
+        if get_sp_world_size() > 1 and pipeline_config is None:
+            raise RuntimeError(
+                "SP rollout requires pipeline_config to be passed to prepare_rollout()."
+            )
+        batch._rollout_session_data = RolloutSessionData(
+            pipeline_config=pipeline_config,
+            sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(),
+            latents_shape=(
+                tuple(batch.latents.shape) if batch.latents is not None else None
+            ),
+        )
+
+    def already_prepared_rollout(self, batch) -> bool:
+        return getattr(batch, "_rollout_session_data", None) is not None
+
+    def _get_or_create_rollout_noise_buffer(
+        self,
+        rollout_session_data: RolloutSessionData,
+        full_shape: tuple,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Get or create the reusable noise buffer (local or full shape) for rollout."""
+        buffer = rollout_session_data.noise_buffer
+        if (
+            buffer is None
+            or buffer.shape != full_shape
+            or buffer.dtype != dtype
+            or buffer.device != device
+        ):
+            buffer = torch.empty(full_shape, device=device, dtype=dtype)
+            rollout_session_data.noise_buffer = buffer
+        return buffer
+
+    def _rollout_variance_noise(
+        self,
+        batch,
+        model_output: torch.FloatTensor,
+        generator: Union[torch.Generator, list[torch.Generator]],
+    ) -> torch.FloatTensor:
+        """Generate variance noise for rollout. If generator is a list, use generator[i] for the i-th batch item."""
+        assert generator is not None, "Generator must be provided"
+
+        rollout_session_data = self._get_rollout_session_data(batch)
+        device = model_output.device
+        dtype = model_output.dtype
+        local_shape = tuple(model_output.shape)
+
+        B = local_shape[0]
+        if isinstance(generator, torch.Generator):
+            assert B == 1, "Generator must be a list if batch size is not 1"
+            generator = [generator]
+        else:
+            assert (
+                len(generator) == B
+            ), "Generator list must have the same length as batch size"
+
+        buffer = self._get_or_create_rollout_noise_buffer(
+            rollout_session_data, rollout_session_data.latents_shape, device, dtype
+        )
+        for i in range(B):
+            torch.randn(
+                rollout_session_data.latents_shape,
+                out=buffer[i : i + 1],
+                generator=generator[i],
+            )
+
+        sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp(
+            batch=batch, latents=buffer
+        )
+        if tuple(sharded_noise.shape) != local_shape:
+            raise ValueError(
+                "Rollout SP noise shape mismatch after shard. "
+                f"Expected local_shape={local_shape}, got {tuple(sharded_noise.shape)}."
+            )
+        return sharded_noise
+
+    def flow_sde_sampling(
+        self,
+        batch,
+        model_output: torch.FloatTensor,
+        sample: torch.FloatTensor,
+        current_sigma: torch.FloatTensor,
+        next_sigma: torch.FloatTensor,
+        generator: torch.Generator,
+    ) -> torch.Tensor:
+        """Flow rollout step for log-prob / sampling (see FlowGRPO-style references).
+
+        ``rollout_sde_type`` (from batch SamplingParams):
+
+        1. ``"sde"``: Standard stochastic differential equation transition (Gaussian).
+        2. ``"cps"``: Coupled Particle Sampling.
+        3. ``"ode"``: Deterministic ODE step (no diffusion noise).
+        """
+        rollout_session_data = self._get_rollout_session_data(batch)
+        sde_type = batch.rollout_sde_type
+        noise_level = float(batch.rollout_noise_level)
+        log_prob_no_const = batch.rollout_log_prob_no_const
+        debug_mode = bool(getattr(batch, "rollout_debug_mode", False))
+
+        if not log_prob_no_const and sde_type != "ode":
+            assert (
+                noise_level > 0
+            ), "True log-probability computation requires a non-zero noise level."
+
+        dt = next_sigma - current_sigma
+
+        # step_index comes from the denoising-loop counter stashed by
+        # DenoisingStage — scheduler._step_index would differ when
+        # _begin_index != 0 (e.g. partial denoising).
+        sde_step_indices = getattr(batch, "rollout_sde_step_indices", None)
+        loop_step_index = getattr(batch, "_rollout_loop_step_index", None)
+        if (
+            sde_type != "ode"
+            and sde_step_indices is not None
+            and loop_step_index is not None
+            and loop_step_index not in sde_step_indices
+        ):
+            effective_sde_type = "ode"
+        else:
+            effective_sde_type = sde_type
+
+        # sde/cps: cast to fp32 to match flowGRPO semantics and avoid the
+        # 0-dim-fp32 wrapped-scalar promotion demoting log-prob to bf16.
+        # ode: keep dtypes unchanged so rollout(ode) stays bit-exact with
+        # rollout=False (scheduling_flow_match_euler_discrete.step()).
+        # log_prob is computed on the full pre-shard noise buffer so SP ranks
+        # produce identical sums — see collect_rollout_log_probs().
+        if effective_sde_type == "sde":
+            model_output = model_output.float()
+            sample = sample.float()
+            variance_noise = self._rollout_variance_noise(
+                batch, model_output, generator
+            )
+            full_variance_noise = rollout_session_data.noise_buffer
+            std_dev_t = (
+                torch.sqrt(
+                    current_sigma
+                    / (
+                        1
+                        - torch.where(
+                            torch.isclose(current_sigma, current_sigma.new_tensor(1.0)),
+                            rollout_session_data.sigma_max,
+                            current_sigma,
+                        )
+                    )
+                )
+                * noise_level
+            )
+            noise_std_dev = std_dev_t * torch.sqrt(-1 * dt)
+            prev_sample_mean = (
+                sample * (1 + std_dev_t**2 / (2 * current_sigma) * dt)
+                + model_output
+                * (1 + std_dev_t**2 * (1 - current_sigma) / (2 * current_sigma))
+                * dt
+            )
+
+            weighted_variance_noise = variance_noise * noise_std_dev
+            prev_sample = prev_sample_mean + weighted_variance_noise
+            log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2)
+
+        elif effective_sde_type == "cps":
+            model_output = model_output.float()
+            sample = sample.float()
+            variance_noise = self._rollout_variance_noise(
+                batch, model_output, generator
+            )
+            full_variance_noise = rollout_session_data.noise_buffer
+            std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2)
+            noise_std_dev = std_dev_t
+            pred_original_sample = sample - current_sigma * model_output
+            noise_estimate = sample + model_output * (1 - current_sigma)
+            prev_sample_mean = pred_original_sample * (
+                1 - next_sigma
+            ) + noise_estimate * torch.sqrt(next_sigma**2 - std_dev_t**2)
+
+            weighted_variance_noise = variance_noise * noise_std_dev
+            prev_sample = prev_sample_mean + weighted_variance_noise
+            log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2)
+
+        elif effective_sde_type == "ode":
+            prev_sample = sample + dt * model_output
+            prev_sample_mean = prev_sample
+            variance_noise = torch.zeros_like(model_output)
+            noise_std_dev = torch.zeros(
+                (), device=model_output.device, dtype=model_output.dtype
+            )
+            log_prob_no_const_val = torch.zeros(
+                rollout_session_data.latents_shape,
+                device=model_output.device,
+                dtype=torch.float32,
+            )
+            # Only enforce the "no full log-prob with ODE" constraint when the
+            # user explicitly chose ODE globally.
+            if sde_type == "ode":
+                assert (
+                    log_prob_no_const
+                ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation"
+
+        else:
+            raise ValueError(f"Unsupported sde_type: {sde_type}")
+
+        reduce_dims = list(range(1, len(log_prob_no_const_val.shape)))
+        local_elem_count = log_prob_no_const_val.new_full(
+            (log_prob_no_const_val.shape[0],),
+            float(math.prod(log_prob_no_const_val.shape[1:])),
+        )
+
+        if log_prob_no_const or effective_sde_type == "ode":
+            log_prob_local_sum = log_prob_no_const_val.sum(dim=reduce_dims)
+        else:
+            log_prob_local_sum = (
+                log_prob_no_const_val / (2 * (noise_std_dev**2))
+                - torch.log(noise_std_dev)
+                - _LOG_SQRT_2PI
+            ).sum(dim=list(range(1, len(log_prob_no_const_val.shape))))
+
+        if debug_mode:
+            self.append_local_rollout_debug_tensors(
+                batch,
+                variance_noise=variance_noise,
+                prev_sample_mean=prev_sample_mean,
+                noise_std_dev=noise_std_dev,
+                model_output=model_output,
+            )
+
+        self.append_local_rollout_log_probs(batch, log_prob_local_sum, local_elem_count)
+
+        return prev_sample
+
+    def append_local_rollout_log_probs(
+        self, batch, log_prob_sum: torch.Tensor, log_prob_count: torch.Tensor
+    ) -> None:
+        rollout_session_data = self._get_rollout_session_data(batch)
+        rollout_session_data.local_log_prob_sum.append(log_prob_sum)
+        rollout_session_data.local_log_prob_count.append(log_prob_count)
+
+    def consume_local_rollout_log_probs(
+        self, batch
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        rollout_session_data = self._get_rollout_session_data(batch)
+        # [B, T]: batch dim 0, denoising step dim 1
+        values_sum = torch.stack(rollout_session_data.local_log_prob_sum, dim=1)
+        values_count = torch.stack(rollout_session_data.local_log_prob_count, dim=1)
+        rollout_session_data.local_log_prob_sum = []
+        rollout_session_data.local_log_prob_count = []
+        return values_sum, values_count
+
+    def collect_rollout_log_probs(self, batch: Req) -> torch.Tensor | None:
+        """Per-step sums are already computed on the full pre-shard noise
+        buffer inside flow_sde_sampling, so every SP rank holds identical
+        values here and no all-reduce is needed."""
+
+        trajectory_log_prob_sum, trajectory_log_prob_count = (
+            self.consume_local_rollout_log_probs(batch)
+        )
+        rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count
+        return rollout_log_probs_tensor.cpu()

python/sglang/multimodal_gen/runtime/post_training/sp_utils.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sequence Parallel helpers for post-training rollout code."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+    sequence_model_parallel_all_reduce,
+)
+
+
+def should_do_sp_collective(batch) -> bool:
+    return get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False)
+
+
+def gather_stacked_latents_for_sp(
+    pipeline_config,
+    batch,
+    stacked_latents: torch.Tensor,
+) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return stacked_latents
+    if stacked_latents.dim() < 2:
+        return stacked_latents
+    bsz, t_steps = stacked_latents.shape[0], stacked_latents.shape[1]
+    flat_inputs = stacked_latents.flatten(0, 1).contiguous()
+    gathered_flat_inputs = pipeline_config.gather_latents_for_sp(
+        flat_inputs, batch=batch
+    )
+    return gathered_flat_inputs.unflatten(0, (bsz, t_steps))
+
+
+def all_reduce_if_sp_sharded(batch, tensor: torch.Tensor) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return tensor
+    tensor = tensor.to(get_local_torch_device())
+    sequence_model_parallel_all_reduce(tensor)
+    return tensor
+
+
+def all_gather_if_sp_sharded(batch, x: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return x
+    x = x.to(get_local_torch_device()).contiguous()
+    return sequence_model_parallel_all_gather(x, dim=dim)
+
+
+def maybe_trim_sp_rope_seq_for_batch(batch, rope: torch.Tensor) -> torch.Tensor:
+    raw = getattr(batch, "raw_latent_shape", None)
+    if raw is None or len(raw) < 2:
+        return rope
+    target = int(raw[1])
+    if rope.shape[0] > target:
+        return rope[:target]
+    return rope

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

python/sglang/multimodal_gen/runtime/post_training/weights_updater.py

@@ -0,0 +1,520 @@
+"""
+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}")