diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index b7474873b..f0d5809ad 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/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), diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 8f68fc35d..aaf4e69b3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/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")