From f744607567d31b7dfaa3f2cdb53648c1127b5162 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 13:05:14 +0800 Subject: [PATCH] [diffusion] UX: report where a component's weights are (#35618) --- .../component_loaders/component_loader.py | 5 +- .../multimodal_gen/runtime/loader/utils.py | 121 ++++++++++++++++++ .../memory_managers/layerwise_offload.py | 12 +- .../unit/test_component_residency_report.py | 113 ++++++++++++++++ 4 files changed, 248 insertions(+), 3 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_component_residency_report.py diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index ce95535c7..f548edf61 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.layers.attention.selector import ( from sglang.multimodal_gen.runtime.loader.utils import ( _normalize_component_type, component_name_to_loader_cls, + format_component_residency, get_memory_usage_of_component, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( @@ -254,11 +255,11 @@ class ComponentLoader(ABC): model_size = get_memory_usage_of_component(component) or "NA" consumed = gpu_mem_before_loading - current_gpu_mem logger.info( - f"Loaded %s: %s ({source} version). model size: %s GB, consumed GPU mem: %.2f GB, avail GPU mem: %.2f GB", + f"Loaded %s: %s ({source} version). model size: %s GB, %s. avail GPU mem: %.2f GB", component_name, component.__class__.__name__, model_size, - consumed, + format_component_residency(component), current_gpu_mem, ) return component, consumed diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index 5fc3bac23..a70ed6974 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """Utilities for selecting and loading models.""" +import bisect import contextlib import glob import json @@ -333,5 +334,125 @@ def get_memory_usage_of_component(module) -> float | None: return round(usage, 2) +def _read_process_mappings() -> tuple[list[int], list[int], list[bool]] | None: + """Sorted (start, end, is_file_backed) of this process's address space. + + Linux only; returns None where /proc is unavailable, and the caller then + reports host bytes without splitting file-backed from anonymous. + """ + try: + with open("/proc/self/maps") as handle: + rows = [] + for line in handle: + fields = line.split(maxsplit=5) + low, _, high = fields[0].partition("-") + path = fields[5].strip() if len(fields) > 5 else "" + # pseudo-paths like [heap] and [stack] are anonymous + rows.append( + (int(low, 16), int(high, 16), bool(path) and path[0] != "[") + ) + except OSError: + return None + rows.sort() + return [r[0] for r in rows], [r[1] for r in rows], [r[2] for r in rows] + + +def component_residency_bytes(module) -> Dict[str, int]: + """Where a component's weights actually sit, in bytes. + + Four buckets, ordered by what the kernel can do with them: device memory, + pinned host memory (which it cannot reclaim at all), file-backed host + memory (which it can drop without swapping), and anonymous host memory. + + Two caveats. `host_mapped` counts the size of the file mapping, not the + pages currently resident in it -- a mapped safetensors file is faulted in + lazily, so the real footprint is at most this. And pinned is tested first + because CUDA's host allocator sits behind a named mapping, which the + file-backed check alone would misread. + + Layerwise-offloaded weights are absent from parameters()/buffers(): the + module keeps (1,) placeholders while its offload managers own the host + copy, so those managers are walked too. Sizes are taken from the storage + and deduped by it, because one flat host buffer backs many logical weights. + """ + if not isinstance(module, nn.Module): + return {} + + totals = {"vram": 0, "host_pinned": 0, "host_mapped": 0, "host": 0} + seen: set[int] = set() + mappings = _read_process_mappings() + + def is_file_backed(pointer: int) -> bool: + if mappings is None: + return False + starts, ends, backed = mappings + index = bisect.bisect_right(starts, pointer) - 1 + if index < 0 or pointer >= ends[index]: + return False + return backed[index] + + def add(tensor: torch.Tensor) -> None: + try: + storage = tensor.untyped_storage() + pointer = storage.data_ptr() + except Exception: + return + # a zero pointer is an empty offload placeholder, not a weight + if pointer == 0 or pointer in seen: + return + seen.add(pointer) + if tensor.device.type != "cpu": + totals["vram"] += storage.nbytes() + return + try: + pinned = tensor.is_pinned() + except Exception: + pinned = False + if pinned: + bucket = "host_pinned" + elif is_file_backed(pointer): + bucket = "host_mapped" + else: + bucket = "host" + totals[bucket] += storage.nbytes() + + for tensor in module.parameters(): + add(tensor) + for tensor in module.buffers(): + add(tensor) + for manager in getattr(module, "layerwise_offload_managers", None) or []: + iter_cpu_weights = getattr(manager, "iter_cpu_weights", None) + if iter_cpu_weights is None: + continue + for _, tensor in iter_cpu_weights(): + add(tensor) + + return totals + + +def format_component_residency(module) -> str: + """Name the places a component's weights are, skipping the empty ones. + + A component that streams from the host reports no VRAM at rest, which is + the point; saying so beats reporting a zero delta that reads as free. + """ + totals = component_residency_bytes(module) + # `pinned` and `pageable` are the standard CUDA pair, and naming mmap after + # the call says what it is: labels a reader has to guess at defeat the point + # of splitting host bytes in the first place. + labels = ( + ("vram", "vram"), + ("host_pinned", "host pinned"), + ("host_mapped", "host mmap"), + ("host", "host pageable"), + ) + parts = [ + f"{label}: {totals[key] / BYTES_PER_GB:.2f} GB" + for key, label in labels + if totals.get(key) + ] + return ", ".join(parts) if parts else "weights: none" + + # component name -> ComponentLoader class component_name_to_loader_cls: Dict[str, Type[Any]] = {} 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 501bad456..f85296d51 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 @@ -1358,9 +1358,19 @@ def configure_layerwise_offload_modules( configured_component_names.append(component_name) if configured_component_names: + # Report where the weights ended up, not just which components opted + # in. The loader's per-component line runs before this, so it can only + # ever describe the pre-offload placement. + from sglang.multimodal_gen.runtime.loader.utils import ( + format_component_residency, + ) + logger.info( "Enabled layerwise offload for pipeline components: %s", - configured_component_names, + ", ".join( + f"{name} ({format_component_residency(modules[name])})" + for name in configured_component_names + ), ) elif warn_missing: logger.debug("No selected pipeline component enabled layerwise offload") diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency_report.py b/python/sglang/multimodal_gen/test/unit/test_component_residency_report.py new file mode 100644 index 000000000..a9dfa890d --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency_report.py @@ -0,0 +1,113 @@ +"""The loader reports where a component's weights are, not a zero delta.""" + +import pathlib + +import pytest +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.loader.utils import ( + component_residency_bytes, + format_component_residency, +) + + +class _FakeOffloadManager: + """Stands in for LayerwiseOffloadManager's host-side weight store.""" + + def __init__(self, tensors): + self._tensors = tensors + + def iter_cpu_weights(self): + for index, tensor in enumerate(self._tensors): + yield f"w{index}", tensor + + +class _Streamed(nn.Module): + """A module whose real weights live in its managers, not its parameters.""" + + def __init__(self, managers): + super().__init__() + # layerwise offload leaves (1,) placeholders behind + self.placeholder = nn.Parameter(torch.empty(0), requires_grad=False) + self.layerwise_offload_managers = managers + + +class TestComponentResidencyBytes: + def test_host_weights_are_counted(self): + buffer = torch.empty(1024, dtype=torch.float32) + module = _Streamed([_FakeOffloadManager([buffer])]) + totals = component_residency_bytes(module) + assert totals["host"] == 4096 + assert totals["vram"] == 0 + assert totals["host_pinned"] == 0 + + def test_slices_of_one_buffer_are_counted_once(self): + # this is the layerwise layout: one flat buffer, many logical weights + buffer = torch.empty(1024, dtype=torch.float32) + views = [buffer[0:256], buffer[256:512], buffer[512:1024]] + module = _Streamed([_FakeOffloadManager(views)]) + assert component_residency_bytes(module)["host"] == 4096 + + def test_empty_placeholders_are_not_counted(self): + module = _Streamed([]) + assert component_residency_bytes(module) == { + "vram": 0, + "host_pinned": 0, + "host_mapped": 0, + "host": 0, + } + + def test_a_file_backed_tensor_is_separated_from_anonymous(self, tmp_path): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + backing = tmp_path / "weights.bin" + backing.write_bytes(b"\0" * 4096) + mapped = torch.from_file( + str(backing), shared=True, size=1024, dtype=torch.float32 + ) + module = _Streamed([_FakeOffloadManager([mapped])]) + totals = component_residency_bytes(module) + assert totals["host_mapped"] == 4096 + assert totals["host"] == 0 + + def test_pinned_wins_over_the_file_backed_check(self): + if not torch.cuda.is_available(): + pytest.skip("pinning needs CUDA") + # CUDA's host allocator sits behind a named mapping, so the + # file-backed check alone would call this one mapped + pinned = torch.empty(1024, dtype=torch.float32, pin_memory=True) + module = _Streamed([_FakeOffloadManager([pinned])]) + totals = component_residency_bytes(module) + assert totals["host_pinned"] == 4096 + assert totals["host_mapped"] == 0 + + def test_resident_parameters_land_in_vram(self): + if not torch.cuda.is_available(): + pytest.skip("needs a device to report device residency") + module = nn.Linear(64, 64, bias=False).cuda() + totals = component_residency_bytes(module) + assert totals["vram"] == 64 * 64 * module.weight.element_size() + assert totals["host"] == 0 + + def test_pinned_host_weights_are_separated(self): + if not torch.cuda.is_available(): + pytest.skip("pinning needs CUDA") + pinned = torch.empty(512, dtype=torch.float32, pin_memory=True) + module = _Streamed([_FakeOffloadManager([pinned])]) + totals = component_residency_bytes(module) + assert totals["host_pinned"] == 2048 + assert totals["host"] == 0 + + def test_non_module_reports_nothing(self): + assert component_residency_bytes(object()) == {} + + +class TestFormatComponentResidency: + def test_only_non_zero_places_are_named(self): + buffer = torch.empty(int(0.5 * 1024**3 / 4), dtype=torch.float32) + module = _Streamed([_FakeOffloadManager([buffer])]) + assert format_component_residency(module) == "host pageable: 0.50 GB" + + def test_a_component_without_weights_says_so(self): + assert format_component_residency(_Streamed([])) == "weights: none"