From 95f71d09d0fd7ac0c47ae052694fa5275741f6e0 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 11:34:13 +0800 Subject: [PATCH 1/2] [diffusion] fix: keep large vocab tables in host memory under layerwise offload Layerwise offload streams a component's layers and then moves everything left over to the device, on the reasonable assumption that the remainder is small -- norms, a patch embed, a time embed. For a text encoder it is not. umT5-XXL's token table is 256384 x 4096 in fp32, so 3.91 GiB stays resident, which on a 12 GiB card is a third of the device. Residency buys almost nothing there. A table is read by gather, not by GEMM: one row per token, so Wan's 512-token prompt touches 8 MiB of those 3.91 GiB, once per request rather than once per denoising step. Streaming it layer by layer would be worse still -- 3.91 GiB moved to read 8 MiB -- so the table belongs in host memory with the lookup running there and only the result crossing PCIe. Swap tables above 256 MiB for placeholders before the move, restore them afterwards, and hook the module so its input goes to the host and its output comes back. Sharded tables are left alone: they are already divided by the world size, and their output feeds an all-reduce that expects a device tensor. Measured on an RTX 3060 12 GiB, Wan2.1-T2V-1.3B, layerwise offload of dit/text_encoder/vae, 320x576x9f, 8 steps, seed 1234, three runs each: peak VRAM 8540 MiB -> 4447 MiB (-4.0 GiB, -48%) TextEncodingStage 3374.9 ms -> 3384.4 ms (+0.28%, ranges overlap) total request 13838.5 ms -> 13815.8 ms Bit-exact: all six runs produce one SHA-256. --- .../memory_managers/layerwise_offload.py | 83 ++++++++++++++++++- .../unit/test_host_resident_vocab_table.py | 73 ++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py 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 501bad456783..4ba303370f70 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 @@ -81,6 +81,82 @@ def compute_streamed_layers( # Adapted from skywork AI Infra diffusion optimize +# Below this a table is not worth a per-request round trip; above it the ratio +# of table size to rows actually read makes residency clearly wasteful. +HOST_RESIDENT_TABLE_MIN_BYTES = 256 * 1024**2 + + +def _host_resident_tables(model: torch.nn.Module) -> List[torch.nn.Module]: + """Vocab tables large enough that holding them on the device is waste. + + A table is read by gather, not by GEMM: one row per token, so a 512-token + prompt touches 8 MiB of umT5-XXL's 3.91 GiB table. Streaming it layer by + layer would be worse than resident -- 3.91 GiB moved to read 8 MiB -- so it + belongs in host memory with the lookup running there. + """ + tables = [] + for module in model.modules(): + weight = getattr(module, "weight", None) + if weight is None or not hasattr(weight, "dim") or weight.dim() != 2: + continue + if getattr(module, "num_embeddings", None) is None: + continue + # A sharded table is already divided by the world size, and its output + # feeds an all-reduce that expects a device tensor. + if getattr(module, "tp_size", 1) != 1: + continue + if weight.numel() * weight.element_size() < HOST_RESIDENT_TABLE_MIN_BYTES: + continue + tables.append(module) + return tables + + +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.""" + detached = [] + for module in _host_resident_tables(model): + weight = module.weight + detached.append((module, weight.data)) + weight.data = torch.empty(0, dtype=weight.dtype, device=weight.device) + return detached + + +def restore_host_resident_tables( + detached: List[Tuple[torch.nn.Module, torch.Tensor]], + device: torch.device | str, +) -> None: + for module, data in detached: + module.weight.data = data + _install_host_gather_hooks(module, device) + logger.info( + "Keeping %s (%.2f GiB) in host memory: a gather reads one row per " + "token, so residency buys almost nothing.", + type(module).__name__, + data.numel() * data.element_size() / (1024**3), + ) + + +def _install_host_gather_hooks( + module: torch.nn.Module, device: torch.device | str +) -> None: + """Run this module's gather on the host, move only the result.""" + + def _inputs_to_host(_module, args, kwargs): + if not args or not torch.is_tensor(args[0]): + return None + return (args[0].to("cpu"),) + args[1:], kwargs + + def _output_to_device(_module, _args, output): + if not torch.is_tensor(output): + return output + return output.to(device, non_blocking=True) + + module.register_forward_pre_hook(_inputs_to_host, with_kwargs=True) + module.register_forward_hook(_output_to_device) + + class LayerwiseOffloadManager: """A lightweight layerwise CPU offload manager. @@ -273,8 +349,10 @@ def _initialize(self) -> None: # Keep non-layer parameters resident on GPU. Layer tensors have already # been replaced by tiny device placeholders, so this does not reload the # offloaded layer weights. + host_resident = detach_host_resident_tables(self.model) if not self._has_dtensor_weights: self.model.to(self.device) + restore_host_resident_tables(host_resident, self.device) self._finalize_initialization() @@ -1067,7 +1145,10 @@ def configure_layerwise_offload(self, server_args: ServerArgs): if enabled_managers and not any( manager._has_dtensor_weights for manager in enabled_managers ): - self.to(enabled_managers[0].device) + device = enabled_managers[0].device + host_resident = detach_host_resident_tables(self) + self.to(device) + restore_host_resident_tables(host_resident, device) for manager in enabled_managers: manager._finalize_initialization() 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 new file mode 100644 index 000000000000..78734b8e0e09 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py @@ -0,0 +1,73 @@ +"""A large vocab table stays in host memory; the gather runs there.""" + +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + _host_resident_tables, + detach_host_resident_tables, + restore_host_resident_tables, +) + +THRESHOLD_PATH = f"{layerwise_offload.__name__}.HOST_RESIDENT_TABLE_MIN_BYTES" + + +def _model(num_embeddings: int = 4096, dim: int = 64) -> torch.nn.Module: + model = torch.nn.Module() + model.embed = torch.nn.Embedding(num_embeddings, dim) + model.proj = torch.nn.Linear(dim, dim) + return model + + +class TestHostResidentTableSelection: + def test_a_large_table_is_selected(self): + model = _model() + with patch(THRESHOLD_PATH, 1024): + assert _host_resident_tables(model) == [model.embed] + + def test_a_small_table_is_left_alone(self): + model = _model() + with patch(THRESHOLD_PATH, 1 << 40): + assert _host_resident_tables(model) == [] + + def test_a_plain_weight_matrix_is_left_alone(self): + model = torch.nn.Module() + model.proj = torch.nn.Linear(512, 512) + with patch(THRESHOLD_PATH, 1024): + assert _host_resident_tables(model) == [] + + def test_a_sharded_table_is_left_alone(self): + model = _model() + model.embed.tp_size = 2 + with patch(THRESHOLD_PATH, 1024): + assert _host_resident_tables(model) == [] + + +class TestDetachAndRestore: + def test_the_weight_survives_a_move_that_skips_it(self): + model = _model() + original = model.embed.weight.data.clone() + with patch(THRESHOLD_PATH, 1024): + detached = detach_host_resident_tables(model) + assert model.embed.weight.numel() == 0 + model.to("cpu") + restore_host_resident_tables(detached, "cpu") + assert torch.equal(model.embed.weight.data, original) + + def test_the_gather_matches_a_plain_lookup(self): + model = _model() + ids = torch.tensor([[1, 2, 3], [4, 5, 6]]) + expected = torch.nn.functional.embedding(ids, model.embed.weight.data.clone()) + with patch(THRESHOLD_PATH, 1024): + restore_host_resident_tables(detach_host_resident_tables(model), "cpu") + assert torch.equal(model.embed(ids), expected) + + def test_nothing_is_hooked_when_no_table_qualifies(self): + model = _model() + with patch(THRESHOLD_PATH, 1 << 40): + detached = detach_host_resident_tables(model) + restore_host_resident_tables(detached, "cpu") + assert detached == [] + assert not model.embed._forward_pre_hooks From dce49f718dd221aa4814fb24427ff0491dfdac6c Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 14:18:30 +0800 Subject: [PATCH 2/2] [diffusion] fix: make the host-resident vocab table opt-in per model Discovering tables by shape was wrong. The bridge is a forward hook, so it only covers the table's own __call__; a model that also reads the weight directly sees a host tensor mid-graph. sana_wm encodes through a Diffusers-backed HF Gemma3ForConditionalGeneration, whose weight is not reached only by its gather, and CI failed with "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu". Take a declaration instead: `host_resident_table_names` on the mixin, next to `layer_names`, listing dotted paths a model vouches for. Nothing is host resident unless a model opts in, so a third-party backbone is never touched. Declared where the table is reached only through its forward and the saving was measured: T5EncoderModel and UMT5EncoderModel (`shared`, 3.91 GiB on Wan2.1's umT5-XXL) and Qwen3VLTextModel (`embed_tokens`, 1.45 GiB per rank on MiniMax-H3). --- .../memory_managers/layerwise_offload.py | 46 ++++++++---- .../runtime/models/encoders/qwen3vl.py | 2 + .../runtime/models/encoders/t5.py | 4 ++ .../unit/test_host_resident_vocab_table.py | 70 +++++++++++++------ 4 files changed, 87 insertions(+), 35 deletions(-) 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 4ba303370f70..5d1a0e0d2264 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 @@ -86,28 +86,44 @@ def compute_streamed_layers( HOST_RESIDENT_TABLE_MIN_BYTES = 256 * 1024**2 +def _resolve_submodule(root: torch.nn.Module, path: str) -> torch.nn.Module | None: + current: Any = root + for part in path.split("."): + current = getattr(current, part, None) + if current is None: + return None + return current if isinstance(current, torch.nn.Module) else None + + def _host_resident_tables(model: torch.nn.Module) -> List[torch.nn.Module]: - """Vocab tables large enough that holding them on the device is waste. + """Declared vocab tables large enough that device residency is waste. A table is read by gather, not by GEMM: one row per token, so a 512-token prompt touches 8 MiB of umT5-XXL's 3.91 GiB table. Streaming it layer by layer would be worse than resident -- 3.91 GiB moved to read 8 MiB -- so it belongs in host memory with the lookup running there. + + Opt-in per model rather than discovered by shape. The bridge is a forward + hook, so it only covers the table's own ``__call__``; a model that also + reads the weight directly -- a tied ``lm_head``, a functional gather inside + a third-party backbone -- would see a host tensor mid-graph. Only a model + whose table is reached solely through its forward may list it. """ tables = [] for module in model.modules(): - weight = getattr(module, "weight", None) - if weight is None or not hasattr(weight, "dim") or weight.dim() != 2: - continue - if getattr(module, "num_embeddings", None) is None: - continue - # A sharded table is already divided by the world size, and its output - # feeds an all-reduce that expects a device tensor. - if getattr(module, "tp_size", 1) != 1: - continue - if weight.numel() * weight.element_size() < HOST_RESIDENT_TABLE_MIN_BYTES: - continue - tables.append(module) + for path in getattr(module, "host_resident_table_names", ()) or (): + table = _resolve_submodule(module, path) + weight = getattr(table, "weight", None) + if weight is None or not hasattr(weight, "dim") or weight.dim() != 2: + continue + # A sharded table is already divided by the world size, and its + # output feeds an all-reduce that expects a device tensor. + if getattr(table, "tp_size", 1) != 1: + continue + if weight.numel() * weight.element_size() < HOST_RESIDENT_TABLE_MIN_BYTES: + continue + if table not in tables: + tables.append(table) return tables @@ -970,6 +986,10 @@ class LayerwiseOffloadableModuleMixin: # The list of names of this module's layer/block ModuleList or Sequential attributes. layer_names: List[str] = [] + + # Dotted paths to gather-only vocab tables that may stay in host memory + # under layerwise offload. See _host_resident_tables for what qualifies. + host_resident_table_names: List[str] = [] layerwise_offload_managers: list[LayerwiseOffloadManager] = [] def _capture_mps_cpu_non_layer_weights(self) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py index 2a7f5e565673..6a778c5fcb04 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py @@ -481,6 +481,8 @@ def forward( class Qwen3VLTextModel(nn.Module): + # used only as `self.embed_tokens(input_ids)`; no tied output head here + host_resident_table_names = ["embed_tokens"] config: Qwen3VLTextConfig _no_split_modules = ["Qwen3VLTextDecoderLayer"] diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py index 99cd23308a1a..3834fb69ec4a 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py @@ -568,6 +568,8 @@ def forward( class T5EncoderModel(TextEncoder): + # encoder-only: no tied lm_head, the table is reached only by its gather + host_resident_table_names = ["shared"] # dp measured here: 1.9x on the encode stage at batch 2/4/8 # (2xH100, T5-XXL width), max_abs_diff=0 vs replicated supports_dp_encode = True @@ -660,6 +662,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class UMT5EncoderModel(TextEncoder): + # encoder-only: no tied lm_head, the table is reached only by its gather + host_resident_table_names = ["shared"] # dp measured here: 1.9x on the encode stage at batch 2/4/8 # (2xH100, T5-XXL width), max_abs_diff=0 vs replicated supports_dp_encode = True 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 78734b8e0e09..5ee6b8180170 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 @@ -1,4 +1,4 @@ -"""A large vocab table stays in host memory; the gather runs there.""" +"""A declared vocab table stays in host memory; the gather runs there.""" from unittest.mock import patch @@ -14,32 +14,58 @@ THRESHOLD_PATH = f"{layerwise_offload.__name__}.HOST_RESIDENT_TABLE_MIN_BYTES" -def _model(num_embeddings: int = 4096, dim: int = 64) -> torch.nn.Module: - model = torch.nn.Module() - model.embed = torch.nn.Embedding(num_embeddings, dim) - model.proj = torch.nn.Linear(dim, dim) - return model +class _Declared(torch.nn.Module): + host_resident_table_names = ["embed"] + def __init__(self, num_embeddings: int = 4096, dim: int = 64): + super().__init__() + self.embed = torch.nn.Embedding(num_embeddings, dim) + self.proj = torch.nn.Linear(dim, dim) -class TestHostResidentTableSelection: - def test_a_large_table_is_selected(self): - model = _model() + +class _Undeclared(torch.nn.Module): + def __init__(self, num_embeddings: int = 4096, dim: int = 64): + super().__init__() + self.embed = torch.nn.Embedding(num_embeddings, dim) + + +class _Nested(torch.nn.Module): + host_resident_table_names = ["language_model.embed"] + + def __init__(self): + super().__init__() + self.language_model = _Undeclared() + + +class TestSelection: + def test_a_declared_table_is_selected(self): + model = _Declared() with patch(THRESHOLD_PATH, 1024): assert _host_resident_tables(model) == [model.embed] - def test_a_small_table_is_left_alone(self): - model = _model() - with patch(THRESHOLD_PATH, 1 << 40): - assert _host_resident_tables(model) == [] + def test_an_undeclared_table_is_left_alone(self): + # the regression guard: a third-party backbone may read the weight + # outside forward, and a forward hook would not cover that + with patch(THRESHOLD_PATH, 1024): + assert _host_resident_tables(_Undeclared()) == [] + + def test_a_dotted_path_resolves(self): + model = _Nested() + with patch(THRESHOLD_PATH, 1024): + assert _host_resident_tables(model) == [model.language_model.embed] - def test_a_plain_weight_matrix_is_left_alone(self): - model = torch.nn.Module() - model.proj = torch.nn.Linear(512, 512) + def test_a_missing_declared_path_is_skipped(self): + model = _Declared() + model.host_resident_table_names = ["not_there"] with patch(THRESHOLD_PATH, 1024): assert _host_resident_tables(model) == [] + def test_a_small_table_is_left_alone(self): + with patch(THRESHOLD_PATH, 1 << 40): + assert _host_resident_tables(_Declared()) == [] + def test_a_sharded_table_is_left_alone(self): - model = _model() + model = _Declared() model.embed.tp_size = 2 with patch(THRESHOLD_PATH, 1024): assert _host_resident_tables(model) == [] @@ -47,7 +73,7 @@ def test_a_sharded_table_is_left_alone(self): class TestDetachAndRestore: def test_the_weight_survives_a_move_that_skips_it(self): - model = _model() + model = _Declared() original = model.embed.weight.data.clone() with patch(THRESHOLD_PATH, 1024): detached = detach_host_resident_tables(model) @@ -57,16 +83,16 @@ def test_the_weight_survives_a_move_that_skips_it(self): assert torch.equal(model.embed.weight.data, original) def test_the_gather_matches_a_plain_lookup(self): - model = _model() + model = _Declared() ids = torch.tensor([[1, 2, 3], [4, 5, 6]]) expected = torch.nn.functional.embedding(ids, model.embed.weight.data.clone()) with patch(THRESHOLD_PATH, 1024): restore_host_resident_tables(detach_host_resident_tables(model), "cpu") assert torch.equal(model.embed(ids), expected) - def test_nothing_is_hooked_when_no_table_qualifies(self): - model = _model() - with patch(THRESHOLD_PATH, 1 << 40): + def test_nothing_is_hooked_when_nothing_qualifies(self): + model = _Undeclared() + with patch(THRESHOLD_PATH, 1024): detached = detach_host_resident_tables(model) restore_host_resident_tables(detached, "cpu") assert detached == []