diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index a7aaa7d13..a409abb6f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -240,6 +240,15 @@ def safetensors_weights_iterator( if use_runai_model_streamer else FALLBACK_READER.name ) + elif to_cpu: + # A host-bound load keeps the checkpoint mapping: mapped pages are the + # zero-copy optimum there, and everything downstream that budgets host + # memory (layerwise offload, pinning, the mapped-weight gate) assumes + # them. The streamer materializes anonymous copies instead -- measured + # as the whole 61.7 GB DiT landing in host anon on the 5090 CI runner + # -- and its strengths (direct-to-GPU, remote streaming) do not apply + # to a local file headed for the CPU. + requested = FALLBACK_READER.name backend = select_weight_reader( requested=requested, needs_key_filter=key_filter is not None ) diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 0c5ebaf3b..297a06951 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -5,6 +5,7 @@ import logging import os import subprocess import sys +import threading import time from datetime import datetime from functools import lru_cache @@ -34,6 +35,10 @@ class MemorySnapshot: reserved_mb: float # current reserved memory (actual VRAM) peak_allocated_mb: float # peak allocated since last reset peak_reserved_mb: float # peak reserved since last reset + # Peak anonymous host memory (RssAnon) sampled by the worker. Anonymous, + # not RSS: file-backed pages the kernel can drop are not a budget cost. + # 0.0 where no sampler ran (non-Linux, or an old record). + peak_host_anon_mb: float = 0.0 def to_dict(self) -> Dict[str, Any]: return { @@ -41,6 +46,7 @@ class MemorySnapshot: "reserved_mb": round(self.reserved_mb, 2), "peak_allocated_mb": round(self.peak_allocated_mb, 2), "peak_reserved_mb": round(self.peak_reserved_mb, 2), + "peak_host_anon_mb": round(self.peak_host_anon_mb, 2), } @@ -125,6 +131,58 @@ def get_git_commit_hash() -> str: return "N/A" +class _HostAnonSampler: + """Tracks this process's peak anonymous host memory (RssAnon). + + The kernel keeps a high-water mark for RSS (VmHWM) but none for the + anonymous share, and the anonymous share is the budget cost: file-backed + pages are droppable and come back on their own. A 1 s sampling thread is + enough resolution for weight-sized (GiB, seconds-long) growth. + """ + + def __init__(self) -> None: + self._peak_kb = 0 + self._started = False + self._lock = threading.Lock() + + def _read_kb(self) -> int: + try: + with open("/proc/self/status") as handle: + for line in handle: + if line.startswith("RssAnon:"): + return int(line.split()[1]) + except OSError: + pass + return 0 + + def _run(self) -> None: + while True: + value = self._read_kb() + if value > self._peak_kb: + self._peak_kb = value + time.sleep(1.0) + + def _ensure_started(self) -> None: + if self._started: + return + with self._lock: + if self._started: + return + self._started = True + if self._read_kb() == 0: + return # no /proc on this platform; peak stays 0 + threading.Thread( + target=self._run, name="host-anon-sampler", daemon=True + ).start() + + def peak_mb(self) -> float: + self._ensure_started() + return max(self._peak_kb, self._read_kb()) / 1024.0 + + +_host_anon_sampler = _HostAnonSampler() + + def capture_memory_snapshot() -> MemorySnapshot: if not torch.get_device_module().is_available(): return MemorySnapshot( @@ -132,6 +190,7 @@ def capture_memory_snapshot() -> MemorySnapshot: reserved_mb=0.0, peak_allocated_mb=0.0, peak_reserved_mb=0.0, + peak_host_anon_mb=_host_anon_sampler.peak_mb(), ) if current_platform.is_mps(): @@ -142,6 +201,7 @@ def capture_memory_snapshot() -> MemorySnapshot: reserved_mb=reserved / (1024**2), peak_allocated_mb=allocated / (1024**2), peak_reserved_mb=reserved / (1024**2), + peak_host_anon_mb=_host_anon_sampler.peak_mb(), ) allocated = torch.get_device_module().memory_allocated() @@ -154,6 +214,7 @@ def capture_memory_snapshot() -> MemorySnapshot: reserved_mb=reserved / (1024**2), peak_allocated_mb=peak_allocated / (1024**2), peak_reserved_mb=peak_reserved / (1024**2), + peak_host_anon_mb=_host_anon_sampler.peak_mb(), ) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json index 43574f50b..bd27a971a 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -155,6 +155,8 @@ "expected_median_denoise_ms": 11000.0, "load_peak_vram_mb": 6000.0, "runtime_peak_vram_mb": 12288.0, + "load_peak_host_anon_mb": 32768.0, + "runtime_peak_host_anon_mb": 32768.0, "estimated_full_test_time_s": 1050.0 } } diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 3a42fab57..eea8b0b88 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -467,8 +467,13 @@ class DiffusionServerBase: expected_load_peak_vram_mb, expected_runtime_peak_vram_mb, ) + validator.validate_peak_host_anon( + summary, + scenario.load_peak_host_anon_mb, + scenario.runtime_peak_host_anon_mb, + ) except AssertionError as e: - logger.error(f"Peak VRAM validation failed for {case.id}:\n{e}") + logger.error(f"Peak memory validation failed for {case.id}:\n{e}") self._dump_baseline_for_testcase(case, summary, missing_scenario) raise @@ -664,6 +669,10 @@ class DiffusionServerBase: { "load_peak_vram_mb": round(summary.load_peak_vram_mb, 2), "runtime_peak_vram_mb": round(summary.runtime_peak_vram_mb, 2), + "load_peak_host_anon_mb": round(summary.load_peak_host_anon_mb, 2), + "runtime_peak_host_anon_mb": round( + summary.runtime_peak_host_anon_mb, 2 + ), } ) diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 6db15c69a..0a4b30346 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -610,6 +610,46 @@ class PerformanceValidator: unit=" MiB", ) + def validate_peak_host_anon( + self, + summary: PerformanceSummary, + expected_load_mb: float | None, + expected_runtime_mb: float | None, + ) -> None: + """Anonymous-host budget: peaks must stay at or under the baseline. + + Skipped wholesale when the baseline carries no host figures (older + scenarios) or the record has none (non-Linux, or a server predating + the sampler) -- the VRAM checks do not imply anything about the host, + as the LoRA-merge blow-up showed: VRAM green, host budget gone. + """ + if expected_load_mb is None and expected_runtime_mb is None: + return + if summary.runtime_peak_host_anon_mb <= 0: + logger.warning( + "Host-anon baseline present but the record has no host peaks; " + "skipping the host budget check" + ) + return + if expected_load_mb is not None: + self._assert_le( + "Load Peak Host Anon", + summary.load_peak_host_anon_mb, + expected_load_mb, + self.tolerances.host_anon, + min_abs_tolerance=256.0, + unit=" MiB", + ) + if expected_runtime_mb is not None: + self._assert_le( + "Runtime Peak Host Anon", + summary.runtime_peak_host_anon_mb, + expected_runtime_mb, + self.tolerances.host_anon, + min_abs_tolerance=256.0, + unit=" MiB", + ) + def validate( self, perf_record: RequestPerfRecord, *args, **kwargs ) -> PerformanceSummary: diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 1bee1dd8c..6c9fcac93 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -48,6 +48,7 @@ class ToleranceConfig: denoise_agg: float load_peak_vram: float = 0.01 runtime_peak_vram: float = 0.02 + host_anon: float = 0.02 @classmethod def load_profile(cls, all_tolerances: dict, profile_name: str) -> ToleranceConfig: @@ -100,6 +101,7 @@ class ToleranceConfig: tol_data.get("runtime_peak_vram", 0.02), ) ), + host_anon=float(tol_data.get("host_anon", 0.02)), ) @@ -115,6 +117,9 @@ class ScenarioConfig: estimated_full_test_time_s: float | None = None load_peak_vram_mb: float | None = None runtime_peak_vram_mb: float | None = None + # Anonymous-host budget caps; None skips the check (older baselines). + load_peak_host_anon_mb: float | None = None + runtime_peak_host_anon_mb: float | None = None @classmethod def from_dict(cls, cfg: dict[str, Any]) -> ScenarioConfig: @@ -131,6 +136,8 @@ class ScenarioConfig: estimated_full_test_time_s=optional_float("estimated_full_test_time_s"), load_peak_vram_mb=optional_float("load_peak_vram_mb"), runtime_peak_vram_mb=optional_float("runtime_peak_vram_mb"), + load_peak_host_anon_mb=optional_float("load_peak_host_anon_mb"), + runtime_peak_host_anon_mb=optional_float("runtime_peak_host_anon_mb"), ) @@ -438,6 +445,8 @@ class PerformanceSummary: all_denoise_steps: dict[int, float] load_peak_vram_mb: float = 0.0 runtime_peak_vram_mb: float = 0.0 + load_peak_host_anon_mb: float = 0.0 + runtime_peak_host_anon_mb: float = 0.0 frames_per_second: float | None = None total_frames: int | None = None avg_frame_time_ms: float | None = None @@ -473,6 +482,14 @@ class PerformanceSummary: runtime_peak_vram_mb = float( record.memory_snapshots.get("runtime_peak", {}).get("peak_reserved_mb", 0.0) ) + load_peak_host_anon_mb = float( + record.memory_snapshots.get("load_peak", {}).get("peak_host_anon_mb", 0.0) + ) + runtime_peak_host_anon_mb = float( + record.memory_snapshots.get("runtime_peak", {}).get( + "peak_host_anon_mb", 0.0 + ) + ) return PerformanceSummary( e2e_ms=e2e_ms, @@ -484,6 +501,8 @@ class PerformanceSummary: all_denoise_steps=per_step, load_peak_vram_mb=load_peak_vram_mb, runtime_peak_vram_mb=runtime_peak_vram_mb, + load_peak_host_anon_mb=load_peak_host_anon_mb, + runtime_peak_host_anon_mb=runtime_peak_host_anon_mb, )