From b8996a5ab2044283b4f90caddc1c4481e0310f28 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 15:20:06 +0800 Subject: [PATCH] [diffusion] fix: keep large vocab tables in host memory under layerwise offload (#35626) --- .../memory_managers/layerwise_offload.py | 103 +++++++++++++++++- .../runtime/models/encoders/qwen3vl.py | 2 + .../runtime/models/encoders/t5.py | 4 + .../unit/test_host_resident_vocab_table.py | 99 +++++++++++++++++ 4 files changed, 207 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 f85296d51..9ce5f169c 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,98 @@ 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 _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]: + """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(): + 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 + + +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 +365,10 @@ class LayerwiseOffloadManager: # 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() @@ -892,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: @@ -1067,7 +1165,10 @@ class LayerwiseOffloadableModuleMixin: 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/runtime/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py index 2a7f5e565..6a778c5fc 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 @@ class Qwen3VLTextDecoderLayer(nn.Module): 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 99cd23308..3834fb69e 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 @@ class T5Stack(nn.Module): 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 @@ class T5EncoderModel(TextEncoder): 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 new file mode 100644 index 000000000..5ee6b8180 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py @@ -0,0 +1,99 @@ +"""A declared 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" + + +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 _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_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_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 = _Declared() + 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 = _Declared() + 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 = _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_nothing_qualifies(self): + model = _Undeclared() + with patch(THRESHOLD_PATH, 1024): + detached = detach_host_resident_tables(model) + restore_host_resident_tables(detached, "cpu") + assert detached == [] + assert not model.embed._forward_pre_hooks