[diffusion] fix: stop the mapped-weight store from holding the parameter itself (#35813)

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -468,7 +468,15 @@ class LayerwiseOffloadManager:
                         # Already a view into the checkpoint. Copying it would
                         # add a second copy of bytes the page cache holds
                         # anyway, and that copy is what does not fit.
-                        self._mapped_cpu_weights[layer_idx][name] = local_weight
+                        # `_to_local_tensor` hands back the parameter itself
+                        # for anything that is not a DTensor, so storing it
+                        # directly stores the parameter -- and `weight.data`
+                        # below swaps that same object's storage for a (1,)
+                        # placeholder, leaving the placeholder in the store.
+                        # Keep an independent tensor over the mapped storage.
+                        self._mapped_cpu_weights[layer_idx][
+                            name
+                        ] = local_weight.detach().view_as(local_weight)
                         self._weight_metadata[layer_idx][name] = {
                             "dtype": local_weight.dtype,
                             "shape": tuple(local_weight.shape),

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -1300,6 +1300,37 @@ def test_a_mapped_weight_is_not_written_back(tmp_path, monkeypatch):
     assert torch.equal(before, after), "writeback must not touch the checkpoint"
 
 
+def test_the_mapped_store_survives_the_placeholder(tmp_path, monkeypatch):
+    """The store must not hold the parameter it is about to have overwritten.
+
+    `_to_local_tensor` returns the parameter itself for anything that is not a
+    DTensor, so storing it directly stores the parameter. `Tensor.data = ...`
+    then swaps that object's storage in place rather than rebinding a name, so
+    the store is left holding the `(1,)` placeholder. Nothing raises: the reload
+    does `gpu_tensor.copy_(cpu_tensor)`, and a one-element source broadcasts
+    into the full shape, so the layer is silently reconstructed from one value.
+    """
+    if not pathlib.Path("/proc/self/maps").exists():
+        pytest.skip("needs /proc to tell a mapping from anonymous memory")
+    manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001)
+    stored = manager._mapped_cpu_weights[0]
+    assert stored, "expected the weight to stay mapped"
+
+    parameters = dict(manager.model.named_parameters())
+    for name, tensor in stored.items():
+        assert tensor.numel() > 1, (
+            f"{name} holds {tensor.numel()} element(s): the store is holding the "
+            "placeholder that was assigned to the parameter, not the weight"
+        )
+        assert tensor is not parameters[name], (
+            f"{name} in the store is the parameter object itself, so assigning "
+            "to the parameter's .data will overwrite the store"
+        )
+    assert manager._mapped_bytes == sum(
+        t.numel() * t.element_size() for t in stored.values()
+    ), "the byte counter and the store must describe the same weights"
+
+
 def test_mapped_weights_are_visible_to_checksums(tmp_path, monkeypatch):
     if not pathlib.Path("/proc/self/maps").exists():
         pytest.skip("needs /proc to tell a mapping from anonymous memory")