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 24a1d1ff9..67644a22c 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 @@ -1532,6 +1532,9 @@ class LayerwiseOffloadableModuleMixin: holds = sum(p.numel() * p.element_size() for _, p in resident) if holds <= self._device_headroom_bytes() * PARK_SIGNIFICANCE: # There is room. Give back any host copies rather than hold them. + # Restore any parameters still bound to placeholders first so we + # do not leave weights as (1,) tensors after clearing the copies. + self.restore_non_layer_weights() self._parked_non_layer_weights.clear() return 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 2784ab1bd..64db72b0a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -1902,6 +1902,36 @@ def test_host_copies_are_given_back_when_room_appears(monkeypatch): assert not comp._parked_non_layer_weights, "host copies should be released" +def test_releasing_host_copies_restores_placeholders(monkeypatch): + """When room appears, release host copies but restore parameters from placeholders.""" + local_device = layerwise_offload_mod.current_platform.get_local_torch_device() + if local_device.type == "cpu": + pytest.skip("parking targets device-resident parameters") + + # The parking path is for the non-MPS accelerator path, so bypass the + # MPS-specific early return while still using the real local device. + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "is_mps", lambda: False) + comp = _ParkableResidentComponent(4) + comp.configure_layerwise_offload(_server_args(performance_mode="memory")) + assert comp.non_layer.device.type == local_device.type + _headroom(monkeypatch, 0) + comp.park_non_layer_weights() + assert comp._parked_non_layer_weights + comp.restore_non_layer_weights() + + _headroom(monkeypatch, 400) + comp.park_non_layer_weights() + assert not comp._parked_non_layer_weights + # Only the non-layer parameter must be restored; streamed layer weights are + # intentionally left as placeholders by the layerwise manager. + assert comp.non_layer.shape != (1,), ( + "non_layer left as a (1,) placeholder after host copies were released" + ) + + def test_park_placeholders_are_shared(monkeypatch): """One stand-in per (device, dtype), not one allocation per parked weight.""" comp = _ParkableResidentComponent(4)