[diffusion] CI: guard the anonymous-host budget alongside peak VRAM (#36051)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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
     )

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(),
     )
 
 

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
         }
     }

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
+                    ),
                 }
             )
 

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:

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