[Feature] Implement update_weights_from_disk for SGLang-D (Diffusion … (#18306)

Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>

python/sglang/multimodal_gen/test/run_suite.py

@@ -10,6 +10,7 @@ Example:
 
 import argparse
 import os
+import random
 import subprocess
 import sys
 from pathlib import Path
@@ -20,6 +21,13 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
+_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = "test_update_weights_from_disk.py"
+_UPDATE_WEIGHTS_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
+_UPDATE_WEIGHTS_MODEL_PAIR_IDS = (
+    "FLUX.2-klein-base-4B",
+    "Qwen-Image",
+)
+
 SUITES = {
     "1-gpu": [
         "test_server_a.py",
@@ -29,6 +37,7 @@ SUITES = {
         "../cli/test_generate_t2i_perf.py",
         # unit tests (no server needed)
         "../test_sampling_params_validate.py",
+        "test_update_weights_from_disk.py",
         # add new 1-gpu test files here
     ],
     "2-gpu": [
@@ -225,6 +234,27 @@ def run_pytest(files, filter_expr=None):
     return returncode
 
 
+def _is_in_ci() -> bool:
+    return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("1", "true", "yes", "on")
+
+
+def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None:
+    if not _is_in_ci():
+        return
+    if _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE not in suite_files_rel:
+        return
+    if os.environ.get(_UPDATE_WEIGHTS_MODEL_PAIR_ENV):
+        print(
+            f"Using preset {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}="
+            f"{os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV]}"
+        )
+        return
+
+    selected_pair = random.choice(_UPDATE_WEIGHTS_MODEL_PAIR_IDS)
+    os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV] = selected_pair
+    print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run")
+
+
 def main():
     args = parse_args()
 
@@ -239,6 +269,7 @@ def main():
 
     # 2. get files from suite
     suite_files_rel = SUITES[args.suite]
+    _maybe_pin_update_weights_model_pair(suite_files_rel)
 
     suite_files_abs = []
     for f_rel in suite_files_rel:

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

@@ -0,0 +1,667 @@
+"""Tests for diffusion `update_weights_from_disk`.
+
+This module verifies the ability to update model weights in place without restarting
+the server, which is critical for RL workflows and iterative fine-tuning scenarios.
+
+Author:
+
+Menyang Liu, https://github.com/dreamyang-liu
+Chenyang Zhao, https://github.com/zhaochenyang20
+
+We use two model pairs for testing (base model / instruct model pairs):
+
+- FLUX.2-klein-base-4B / FLUX.2-klein-4B
+- Qwen/Qwen-Image / Qwen/Qwen-Image-2512
+
+These model pairs share the same architecture but differ in transformer
+weights. The basic testing logic is to refit the instruct model into the
+base model and verify the checksum of the transformer weights are the same,
+which simulates the real-world RL scenario. However, since these two model
+pairs only differ in transformer weights, and we want to verify update a
+specific module with update_weights_from_disk API, we need to create a perturbed
+instruct model that adds noise to the vae weights. In this sense, the instruct
+model differs from the base model in vae and transformer weights, the text
+encoder are still the same.
+
+To strictly verify the correctness of the refit API, we compare the checksum in
+SHA-256 on the disk and the server.
+
+NOTE and TODO: In the refit a specific module test, we randomly select one module
+from the transformer and vae to refit the server and keep other modules the same.
+As described above, the vae's weights are perturbed. If we select the vae to be the
+target module, ideally speaking, we should assert that the refitted vae's checksum
+is the same as directly computed from the perturbed vae weights in the disk. However,
+since the there is complex weight-name remapping and QKV merge during model loading,
+it is not easy to compare the server-disk checksum for vae and text encoder directly.
+Therefore, if the target module is vae, we only verify that the refitted vae's checksum
+is different from the base model's vae's checksum.
+
+It should be good issue to solve for the community to adds comparison the server-disk
+checksum for vae and text encoder in this test.
+
+=============================================================================
+
+Test organization:
+
+7 test cases in 2 classes;
+two model pairs are tested locally, one in CI.
+
+=============================================================================
+
+Class 1: TestUpdateWeightsFromDisk                  (6 tests) — API contract, checksum & rollback
+Class 2: TestUpdateWeightsFromDiskWithOffload       (1 test) — Offload-aware update + checksum
+
+-----------------------------------------------------------------------------
+
+Class 1: TestUpdateWeightsFromDisk
+
+Validate the update_weights_from_disk API contract, request/response shape,
+error handling, checksum verification, and corrupted-weight rollback.
+
+All tests share one class-scoped server (same process, same in-memory weights).
+Tests that require "base model then update" should be explicitly reset to
+base model first so behavior is order-independent and updates are real
+(base -> perturbed), not no-ops (perturbed -> perturbed).
+
+  • test_update_weights_from_disk_default
+
+    base model -> perturbed model with flush_cache=True.
+    Verifies after-update transformer checksum == perturbed model's
+    transformer disk checksum
+
+
+  • test_update_weights_specific_modules
+
+    base -> perturbed with flush_cache=False.  Randomly selects one module
+    from _DIFFERING_MODULES (transformer and vae) as target_modules, updates
+    only that module. Verifies that:
+    (1) targeted module's in-memory checksum changed;
+    (2) non-targeted modules' in-memory checksums are unchanged.
+
+  • test_update_weights_nonexistent_model
+
+    model_path set to a non-existent path; must fail (400, success=False).
+
+    Ensure server is healthy after failed update and server's transformer
+    checksums equal base model's transformer disk checksum.
+
+  • test_update_weights_missing_model_path
+
+    Request body empty (no model_path); must fail (400, success=False).
+
+    Ensure server is healthy after failed update and server's transformer
+    checksums equal base model's transformer disk checksum.
+
+  • test_update_weights_nonexistent_module
+
+    target_modules=["nonexistent_module"]; must fail (400, success=False).
+
+    Verify server is healthy after failed update and server's checksums
+    equal base model's transformer disk checksum.
+
+  • test_corrupted_weights_rollback
+
+    All-or-nothing rollback: We first refit the server from base model ->
+    perturbed model. We manually truncate the vae weights of the base
+    model to get a corrupted model. We then call the refit to update
+    the server from the perturbed model -> corrupted model. Verify that:
+
+    1. The update fails due to truncated vae, server should roll back to the
+    perturbed model, i.e., server's transformer weights == perturbed model's
+    transformer weights != base model's transformer weights.
+
+    2. After the rollback, server's vae weights == perturbed model's vae
+    weights != base model's vae weights.
+
+    3. After the rollback, server's text encoder weights == base model's
+    text encoder weights == perturbed model's text encoder weights.
+
+-----------------------------------------------------------------------------
+
+Class 2: TestUpdateWeightsFromDiskWithOffload
+
+
+Ensure weight updates and checksum verification work when layerwise offload is enabled
+(--dit-layerwise-offload). With offload, parameters live in CPU buffers and only left
+small torch.empty((1,)) as placeholders on GPU; the updater must write into CPU buffers
+and update prefetched GPU tensors without shape mismatch.
+
+  • test_update_weights_with_offload_enabled
+
+    Server with --dit-layerwise-offload (base). Load perturbed checkpoint;
+    must succeed (200, success=True), no "Shape mismatch". server's transformer checksum
+    matches perturbed model's transformer disk checksum.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+import random
+import shutil
+import tempfile
+import threading
+from collections.abc import Callable
+
+import pytest
+import requests
+from safetensors.torch import load_file, save_file
+
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    compute_weights_checksum,
+    safetensors_weights_iterator,
+)
+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.multimodal_gen.test.server.test_server_utils import (
+    ServerManager,
+)
+from sglang.multimodal_gen.test.test_utils import get_dynamic_server_port, is_in_ci
+
+logger = init_logger(__name__)
+
+
+_TRANSFORMER_MODULE = "transformer"
+_VAE_MODULE = "vae"
+_TEXT_ENCODER_MODULE_PREFIX = "text_encoder"
+
+
+# Modules whose weights differ between the base model and the perturbed
+# perturbed checkpoint
+_DIFFERING_MODULES: list[str] = [_TRANSFORMER_MODULE, _VAE_MODULE]
+
+_ALL_MODEL_PAIRS: list[tuple[str, str]] = [
+    (
+        "black-forest-labs/FLUX.2-klein-base-4B",
+        "black-forest-labs/FLUX.2-klein-4B",
+    ),
+    (
+        "Qwen/Qwen-Image",
+        "Qwen/Qwen-Image-2512",
+    ),
+]
+
+
+_CI_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
+
+
+def _resolve_active_model_pairs() -> list[tuple[str, str]]:
+    if not is_in_ci():
+        return _ALL_MODEL_PAIRS
+
+    pair_by_id = {pair[0].split("/")[-1]: pair for pair in _ALL_MODEL_PAIRS}
+    selected_pair_id = os.environ.get(_CI_MODEL_PAIR_ENV)
+    if selected_pair_id is None:
+        return [random.choice(_ALL_MODEL_PAIRS)]
+
+    selected_pair = pair_by_id.get(selected_pair_id)
+    if selected_pair is None:
+        valid_ids = ", ".join(sorted(pair_by_id))
+        raise ValueError(
+            f"Invalid {_CI_MODEL_PAIR_ENV}={selected_pair_id!r}. "
+            f"Expected one of: {valid_ids}."
+        )
+    return [selected_pair]
+
+
+_ACTIVE_MODEL_PAIRS = _resolve_active_model_pairs()
+_PAIR_IDS = [p[0].split("/")[-1] for p in _ACTIVE_MODEL_PAIRS]
+
+
+@functools.lru_cache(maxsize=None)
+def _compute_checksum_from_disk(model_path: str, module_name: str) -> str:
+    """Compute SHA-256 checksum from safetensors files on disk.
+
+    Uses the same compute_weights_checksum function as the server,
+    so the checksums are directly comparable.
+
+    Results are cached (keyed on model_path and module_name) because the
+    same disk checksum is requested multiple times across tests.
+    """
+    local_path = maybe_download_model(model_path)
+    weights_dir = os.path.join(local_path, module_name)
+    assert os.path.exists(
+        weights_dir
+    ), f"No weights dir for {module_name} in {local_path}"
+
+    safetensors_files = _list_safetensors_files(weights_dir)
+    assert safetensors_files, f"No safetensors files in {weights_dir}"
+
+    return compute_weights_checksum(safetensors_weights_iterator(safetensors_files))
+
+
+def _clone_model_with_modified_module(
+    src_model: str,
+    dst_model: str,
+    target_module: str,
+    transform_safetensor: Callable[[str, str], None],
+) -> None:
+    # Symlink root-level files (model_index.json, etc.).
+    for fname in os.listdir(src_model):
+        src_path = os.path.join(src_model, fname)
+        dst_path = os.path.join(dst_model, fname)
+        if os.path.isfile(src_path) and not os.path.exists(dst_path):
+            os.symlink(src_path, dst_path)
+
+    for module_dir in sorted(os.listdir(src_model)):
+        src_dir = os.path.join(src_model, module_dir)
+        dst_dir = os.path.join(dst_model, module_dir)
+        if not os.path.isdir(src_dir):
+            continue
+
+        if module_dir != target_module:
+            if not os.path.exists(dst_dir):
+                os.symlink(src_dir, dst_dir)
+            continue
+
+        os.makedirs(dst_dir, exist_ok=True)
+        transformed = False
+        for fname in sorted(os.listdir(src_dir)):
+            src_file = os.path.join(src_dir, fname)
+            dst_file = os.path.join(dst_dir, fname)
+            if not os.path.isfile(src_file):
+                continue
+
+            if not fname.endswith(".safetensors") or transformed:
+                if not os.path.exists(dst_file):
+                    os.symlink(src_file, dst_file)
+                continue
+
+            transform_safetensor(src_file, dst_file)
+            transformed = True
+
+
+def _truncate_safetensor(src_file: str, dst_file: str) -> None:
+    shutil.copy2(src_file, dst_file)
+    size = os.path.getsize(dst_file)
+    with open(dst_file, "r+b") as f:
+        f.truncate(size - 2)
+    logger.info(
+        "Created corrupted safetensors: %s (%d -> %d bytes)",
+        dst_file,
+        size,
+        size - 2,
+    )
+
+
+def _perturb_safetensor(src_file: str, dst_file: str) -> None:
+
+    tensors = load_file(src_file)
+    perturbed = {
+        k: (t + 0.01 if t.is_floating_point() else t) for k, t in tensors.items()
+    }
+    save_file(perturbed, dst_file)
+    logger.info("Created perturbed safetensors: %s", dst_file)
+
+
+class _UpdateWeightsApiMixin:
+    def _update_weights(
+        self,
+        base_url: str,
+        model_path: str,
+        flush_cache: bool = True,
+        target_modules: list[str] | None = None,
+        timeout: int = 300,
+    ) -> tuple[dict, int]:
+        payload = {"model_path": model_path, "flush_cache": flush_cache}
+        if target_modules is not None:
+            payload["target_modules"] = target_modules
+        response = requests.post(
+            f"{base_url}/update_weights_from_disk",
+            json=payload,
+            timeout=timeout,
+        )
+        return response.json(), response.status_code
+
+    def _get_weights_checksum(
+        self,
+        base_url: str,
+        module_names: list[str] | None = None,
+        timeout: int = 300,
+    ) -> dict:
+        payload = {}
+        if module_names is not None:
+            payload["module_names"] = module_names
+        response = requests.post(
+            f"{base_url}/get_weights_checksum",
+            json=payload,
+            timeout=timeout,
+        )
+        assert (
+            response.status_code == 200
+        ), f"get_weights_checksum failed: {response.status_code} {response.text}"
+        return response.json()
+
+    def _assert_server_matches_model(
+        self,
+        base_url: str,
+        expected_model: str,
+    ) -> None:
+        server_checksums = self._get_weights_checksum(
+            base_url, module_names=[_TRANSFORMER_MODULE]
+        )
+        expected_cs = _compute_checksum_from_disk(expected_model, _TRANSFORMER_MODULE)
+        server_cs = server_checksums.get(_TRANSFORMER_MODULE)
+        assert server_cs == expected_cs, (
+            f"Checksum mismatch on '{_TRANSFORMER_MODULE}'\n"
+            f"  expected({expected_model}): {expected_cs}\n"
+            f"  server: {server_cs}"
+        )
+
+
+class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
+
+    @pytest.fixture(
+        scope="class",
+        params=_ACTIVE_MODEL_PAIRS,
+        ids=_PAIR_IDS,
+    )
+    def diffusion_server_no_offload(self, request):
+        default_model, source_model = request.param
+        port = get_dynamic_server_port()
+        wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600"))
+
+        manager = ServerManager(
+            model=default_model,
+            port=port,
+            wait_deadline=wait_deadline,
+            extra_args="--num-gpus 1",
+        )
+
+        # Ensure models are local before spawning threads that need the paths.
+        local_default = maybe_download_model(default_model)
+        local_source = maybe_download_model(source_model)
+
+        perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_")
+        corrupted_vae_model_dir = tempfile.mkdtemp(prefix="sglang_corrupted_")
+
+        # Run all disk I/O in background while the server boots.
+        bg_threads = [
+            threading.Thread(
+                target=_compute_checksum_from_disk, args=(default_model, module)
+            )
+            for module in _DIFFERING_MODULES
+        ] + [
+            threading.Thread(
+                target=_clone_model_with_modified_module,
+                args=(
+                    local_source,
+                    perturbed_vae_model_dir,
+                    _VAE_MODULE,
+                    _perturb_safetensor,
+                ),
+            ),
+            threading.Thread(
+                target=_clone_model_with_modified_module,
+                args=(
+                    local_default,
+                    corrupted_vae_model_dir,
+                    _VAE_MODULE,
+                    _truncate_safetensor,
+                ),
+            ),
+        ]
+        for t in bg_threads:
+            t.start()
+
+        ctx = manager.start()
+        for t in bg_threads:
+            t.join()
+
+        # Sanity: all _DIFFERING_MODULES should differ between base and perturbed.
+        for module in _DIFFERING_MODULES:
+            assert _compute_checksum_from_disk(
+                default_model, module
+            ) != _compute_checksum_from_disk(perturbed_vae_model_dir, module), (
+                f"Assumption violated: {module} should differ between "
+                f"{default_model} and {perturbed_vae_model_dir}"
+            )
+
+        try:
+            yield ctx, default_model, perturbed_vae_model_dir, corrupted_vae_model_dir
+        finally:
+            ctx.cleanup()
+            shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True)
+            shutil.rmtree(corrupted_vae_model_dir, ignore_errors=True)
+
+    def test_update_weights_from_disk_default(self, diffusion_server_no_offload):
+        """Default update (target_modules=None, flush_cache=True): all changed modules updated."""
+        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        self._update_weights(base_url, default_model, flush_cache=True)
+
+        result, status_code = self._update_weights(
+            base_url, perturbed_model_dir, flush_cache=True
+        )
+        assert status_code == 200
+        assert result.get("success", False), f"Update failed: {result.get('message')}"
+
+        self._assert_server_matches_model(base_url, perturbed_model_dir)
+
+    def test_update_weights_specific_modules(self, diffusion_server_no_offload):
+        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        # Reset server to default_model.
+        self._update_weights(base_url, default_model)
+        before_checksums = self._get_weights_checksum(
+            base_url, module_names=_DIFFERING_MODULES
+        )
+
+        target_modules = [random.choice(_DIFFERING_MODULES)]
+        result, status_code = self._update_weights(
+            base_url,
+            perturbed_model_dir,
+            target_modules=target_modules,
+            flush_cache=False,
+        )
+        assert status_code == 200, f"Update failed: {result}"
+        assert result.get("success", False), f"Update failed: {result.get('message')}"
+
+        after_checksums = self._get_weights_checksum(
+            base_url, module_names=_DIFFERING_MODULES
+        )
+
+        # Targeted module should have changed.
+        for name in target_modules:
+            assert after_checksums.get(name) != before_checksums.get(name), (
+                f"Targeted module '{name}' checksum should change after update\n"
+                f"  before: {before_checksums.get(name)}\n"
+                f"  after:  {after_checksums.get(name)}"
+            )
+
+        # Non-targeted modules should be unchanged.
+        for name, cs in after_checksums.items():
+            if name in target_modules or cs == "not_found":
+                continue
+            assert cs == before_checksums.get(name), (
+                f"Non-targeted module '{name}' should be unchanged\n"
+                f"  before: {before_checksums.get(name)}\n"
+                f"  after:  {cs}"
+            )
+
+    def test_update_weights_nonexistent_model(self, diffusion_server_no_offload):
+        """Nonexistent model path must fail (400). Server healthy, checksums == base disk."""
+        ctx, default_model, _, _ = diffusion_server_no_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        self._update_weights(base_url, default_model)
+
+        result, status_code = self._update_weights(
+            base_url,
+            "/nonexistent/path/to/model",
+            timeout=60,
+        )
+        logger.info(f"Update result for nonexistent model: {result}")
+
+        assert status_code == 400, f"Expected 400, got {status_code}"
+        assert not result.get("success", True), "Should fail for nonexistent model"
+        self._assert_server_matches_model(base_url, default_model)
+
+    def test_update_weights_missing_model_path(self, diffusion_server_no_offload):
+        """Request without model_path must fail (400). Server healthy, checksums == base disk."""
+        ctx, default_model, _, _ = diffusion_server_no_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        self._update_weights(base_url, default_model)
+
+        response = requests.post(
+            f"{base_url}/update_weights_from_disk",
+            json={},
+            timeout=30,
+        )
+
+        assert response.status_code == 400, f"Expected 400, got {response.status_code}"
+        result = response.json()
+        assert not result.get("success", True), "Should fail when model_path is missing"
+        self._assert_server_matches_model(base_url, default_model)
+
+    def test_update_weights_nonexistent_module(self, diffusion_server_no_offload):
+        """Nonexistent module must fail (400). Server healthy, checksums == base disk."""
+        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        self._update_weights(base_url, default_model)
+
+        result, status_code = self._update_weights(
+            base_url,
+            perturbed_model_dir,
+            target_modules=["nonexistent_module"],
+            timeout=60,
+        )
+        logger.info(f"Update nonexistent module result: {result}")
+
+        assert status_code == 400, f"Expected 400, got {status_code}"
+        assert not result.get("success", True), "Should fail for nonexistent module"
+        assert "not found in pipeline" in result.get("message", "")
+        self._assert_server_matches_model(base_url, default_model)
+
+    def test_corrupted_weights_rollback(self, diffusion_server_no_offload):
+        ctx, default_model, perturbed_model_dir, corrupted_vae_model_dir = (
+            diffusion_server_no_offload
+        )
+        base_url = f"http://localhost:{ctx.port}"
+
+        # base → perturbed
+        self._update_weights(base_url, default_model)
+        base_checksums = self._get_weights_checksum(base_url)
+
+        result, status_code = self._update_weights(base_url, perturbed_model_dir)
+        assert status_code == 200 and result.get("success")
+        perturbed_checksums = self._get_weights_checksum(base_url)
+
+        text_encoder_modules = sorted(
+            name
+            for name in perturbed_checksums
+            if _TEXT_ENCODER_MODULE_PREFIX in name
+            and perturbed_checksums.get(name) != "not_found"
+            and base_checksums.get(name) != "not_found"
+        )
+        assert (
+            text_encoder_modules
+        ), "Expected at least one text encoder module checksum"
+
+        # perturbed → corrupted (should fail and rollback)
+        rollback_targets = [_TRANSFORMER_MODULE, _VAE_MODULE]
+        result, status_code = self._update_weights(
+            base_url,
+            corrupted_vae_model_dir,
+            target_modules=rollback_targets,
+        )
+        assert (
+            status_code == 400
+        ), f"Expected 400 on corrupted weights, got {status_code}"
+        assert not result.get("success", True)
+        message = result.get("message", "")
+        assert "rolled back" in message.lower()
+        # The updater reports the first failing module in the error message.
+        # With ordered target_modules=[transformer, vae], this makes the
+        # failure point explicit: transformer is processed first, then vae fails.
+        assert (
+            "Failed to update module 'vae'" in message
+        ), f"Expected vae to be the explicit failure point, got: {message}"
+        rolled_back_checksums = self._get_weights_checksum(base_url)
+
+        # 1) transformer: server == perturbed != base
+        transformer_base = base_checksums.get(_TRANSFORMER_MODULE)
+        transformer_perturbed = perturbed_checksums.get(_TRANSFORMER_MODULE)
+        transformer_rolled_back = rolled_back_checksums.get(_TRANSFORMER_MODULE)
+        assert transformer_rolled_back == transformer_perturbed
+        assert transformer_rolled_back != transformer_base
+
+        # 2) vae: server == perturbed != base
+        vae_base = base_checksums.get(_VAE_MODULE)
+        vae_perturbed = perturbed_checksums.get(_VAE_MODULE)
+        vae_rolled_back = rolled_back_checksums.get(_VAE_MODULE)
+        assert vae_rolled_back == vae_perturbed
+        assert vae_rolled_back != vae_base
+
+        # 3) text encoder(s): server == base == perturbed
+        for name in text_encoder_modules:
+            assert rolled_back_checksums.get(name) == perturbed_checksums.get(
+                name
+            ), f"Text encoder module '{name}' should stay equal to perturbed"
+            assert rolled_back_checksums.get(name) == base_checksums.get(
+                name
+            ), f"Text encoder module '{name}' should stay equal to base"
+
+
+class TestUpdateWeightsFromDiskWithOffload(_UpdateWeightsApiMixin):
+    """Test update_weights_from_disk with layerwise offload enabled."""
+
+    @pytest.fixture(scope="class", params=_ACTIVE_MODEL_PAIRS, ids=_PAIR_IDS)
+    def diffusion_server_with_offload(self, request):
+        default_model, source_model = request.param
+        port = get_dynamic_server_port()
+        wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600"))
+
+        local_source = maybe_download_model(source_model)
+        perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_")
+
+        clone_thread = threading.Thread(
+            target=_clone_model_with_modified_module,
+            args=(
+                local_source,
+                perturbed_vae_model_dir,
+                _VAE_MODULE,
+                _perturb_safetensor,
+            ),
+        )
+        clone_thread.start()
+
+        manager = ServerManager(
+            model=default_model,
+            port=port,
+            wait_deadline=wait_deadline,
+            extra_args="--num-gpus 1 --dit-layerwise-offload true",
+        )
+
+        ctx = manager.start()
+        clone_thread.join()
+
+        try:
+            yield ctx, default_model, perturbed_vae_model_dir
+        finally:
+            ctx.cleanup()
+            shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True)
+
+    def test_update_weights_with_offload_enabled(self, diffusion_server_with_offload):
+        ctx, _, perturbed_model_dir = diffusion_server_with_offload
+        base_url = f"http://localhost:{ctx.port}"
+
+        result, status_code = self._update_weights(base_url, perturbed_model_dir)
+        assert status_code == 200, f"Expected 200, got {status_code}"
+        assert result.get("success", False), f"Update failed: {result.get('message')}"
+
+        message = result.get("message", "")
+        assert "Shape mismatch" not in message, f"Shape mismatch detected: {message}"
+
+        self._assert_server_matches_model(base_url, perturbed_model_dir)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -28,6 +28,8 @@ from typing import Sequence
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
 
+DEFAULT_SMALL_MODEL = "Tongyi-MAI/Z-Image-Turbo"
+
 
 @dataclass
 class ToleranceConfig:
@@ -339,8 +341,6 @@ TURBOWAN_I2V_sampling_params = DiffusionSamplingParams(
     fps=4,
 )
 
-DEFAULT_SMALL_MODEL = "Tongyi-MAI/Z-Image-Turbo"
-
 # All test cases with clean default values
 # To test different models, simply add more DiffusionCase entries
 ONE_GPU_CASES_A: list[DiffusionTestCase] = [