[diffusion] UX: report where a component's weights are (#35618)

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

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]] = {}

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")

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"