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 29670a97c..24a1d1ff9 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 @@ -150,11 +150,16 @@ def _host_resident_tables(model: torch.nn.Module) -> List[torch.nn.Module]: def detach_host_resident_tables( model: torch.nn.Module, ) -> List[Tuple[torch.nn.Module, torch.Tensor]]: - """Swap large vocab tables for placeholders so a `.to(device)` skips them.""" + """Park large vocab tables on the host so a `.to(device)` skips them.""" detached = [] for module in _host_resident_tables(model): weight = module.weight - detached.append((module, weight.data)) + # Most loaders leave the table on the host, but model-owned loading + # paths may already have placed it on the accelerator. The input hook + # below always sends indices to the host, so retaining accelerator data + # here would restore a CUDA weight and create a CPU-index/CUDA-weight + # mismatch in the embedding gather. + detached.append((module, weight.data.to("cpu"))) weight.data = torch.empty(0, dtype=weight.dtype, device=weight.device) return detached diff --git a/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py index 5ee6b8180..326010e57 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py @@ -2,6 +2,7 @@ from unittest.mock import patch +import pytest import torch from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload @@ -97,3 +98,22 @@ class TestDetachAndRestore: restore_host_resident_tables(detached, "cpu") assert detached == [] assert not model.embed._forward_pre_hooks + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_a_device_resident_table_is_parked_on_the_host(self): + model = _Declared().to("cuda") + ids = torch.tensor([[1, 2, 3], [4, 5, 6]], device="cuda") + with torch.no_grad(): + expected = model.embed(ids) + + with patch(THRESHOLD_PATH, 1024): + detached = detach_host_resident_tables(model) + assert model.embed.weight.numel() == 0 + model.to("cuda") + restore_host_resident_tables(detached, "cuda") + + with torch.no_grad(): + actual = model.embed(ids) + assert model.embed.weight.device.type == "cpu" + assert actual.device.type == "cuda" + assert torch.equal(actual, expected)