chore(m11): offload-audit GPU memory probes for vast smoke runs

docker/patch/latest/sglang.patch

@@ -637,7 +637,7 @@ index 293a84350..68911c433 100644
 @@ -137,6 +148,13 @@ class SchedulerUpdateWeightsMixin:
              self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
              self.flush_cache()
- 
+
 +            if self.disaggregation_mode == DisaggregationMode.DECODE:
 +                if hasattr(self, "disagg_decode_prealloc_queue"):
 +                    self.disagg_decode_prealloc_queue.release_memory_occupation()

miles/backends/megatron_utils/actor.py

@@ -254,7 +254,11 @@ def wake_up(self) -> None:
 
         clear_memory()
         if not skip_tms:
+            from miles.utils.memory_utils import log_nontorch
+
+            log_nontorch("before reload_process_groups")
             reload_process_groups()
+            log_nontorch("after reload_process_groups")
         print_memory("after wake_up model")
 
     def _switch_model(self, target_tag: str) -> None:
@@ -679,6 +683,9 @@ def build_cpu_bucket_cache(self, step: int) -> int:
         from .update_weight.cpu_bucket_cache import BucketEntry
         from .update_weight.hf_weight_iterator_base import HfWeightIteratorBase
 
+        from miles.utils.memory_utils import log_nontorch
+
+        log_nontorch(f"before build_cpu_bucket_cache step={step}")
         cache = self._ensure_cpu_bucket_cache()
         is_owner = self._is_cache_owner_rank()
 
@@ -760,6 +767,7 @@ def build_cpu_bucket_cache(self, step: int) -> int:
                 )
 
             cache.put_step(int(step), buckets)
+        log_nontorch(f"after build_cpu_bucket_cache step={step}")
         return int(step)
 
     def run_sync_session(self, plan) -> int:

miles/backends/megatron_utils/model.py

@@ -449,6 +449,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
         return output_tensor, partial(loss_function, args, batch, num_microbatches, apply_megatron_loss_scaling=True)
 
     # Forward pass.
+    from miles.utils.memory_utils import log_nontorch as _log_nontorch
+
+    _log_nontorch("before forward_backward")
     forward_backward_func = get_forward_backward_func()
     losses_reduced = forward_backward_func(
         forward_step_func=forward_step,
@@ -461,6 +464,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
         forward_only=False,
     )
 
+    _log_nontorch("after forward_backward")
     valid_step = True
     grad_norm = 0.0
     if (not disable_optimizer) and (not getattr(args, "check_for_nan_in_loss_and_grad", True)):
@@ -489,6 +493,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
         # Update learning rate.
         assert update_successful
         opt_param_scheduler.step(increment=args.global_batch_size)
+        _log_nontorch("after optimizer.step")
 
     # release grad
     for model_chunk in model:

miles/backends/sglang_utils/sglang_engine.py

@@ -616,13 +616,42 @@ def unload_lora_adapter(self, lora_name: str):
             {"lora_name": lora_name},
         )
 
+    def _log_whole_gpu(self, label: str) -> None:
+        """SGL offload audit: whole-GPU used (nvidia-smi) for this engine's
+        visible GPUs — before/after release/resume shows how much physical
+        memory the engine actually returned, independent of PID-namespace
+        issues that break per-process attribution on some containers."""
+        import subprocess
+
+        try:
+            out = subprocess.check_output(
+                ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"],
+                text=True, timeout=10,
+            )
+            rows = dict(line.split(", ") for line in out.strip().splitlines())
+            cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+            visible = [x.strip() for x in cvd.split(",") if x.strip()] or sorted(rows)
+            usage = {g: f"{rows.get(g)} MiB" for g in visible}
+            # print (not logger): the engine actor process has no logging
+            # handler configured, so logger.info is silently dropped; Ray
+            # forwards actor stdout unconditionally.
+            print(
+                f"[SGL-OFFLOAD-AUDIT] {label} engine={self.server_host}:{self.server_port} gpus={usage}",
+                flush=True,
+            )
+        except Exception as exc:  # noqa: BLE001
+            print(f"[SGL-OFFLOAD-AUDIT] {label}: probe failed {exc!r}", flush=True)
+
     def release_memory_occupation(self, tags: list[str] = None):
         """Release memory occupation. Available tags: weights, kv_cache."""
+        self._log_whole_gpu("before release_memory_occupation")
         self.flush_cache()
-        return self._make_request(
+        result = self._make_request(
             "release_memory_occupation",
             {"tags": tags},
         )
+        self._log_whole_gpu("after release_memory_occupation")
+        return result
 
     # ------------------------------------------------------------------
     # F1 RLix-mode sleep/wake helpers (used by F2 RolloutManager.shrink_engines)
@@ -825,10 +854,13 @@ def resume_memory_occupation(self, tags: list[str] = None):
         """
         Available tags for multi-stage resume: weights, kv_cache
         """
-        return self._make_request(
+        self._log_whole_gpu("before resume_memory_occupation")
+        result = self._make_request(
             "resume_memory_occupation",
             {"tags": tags},
         )
+        self._log_whole_gpu("after resume_memory_occupation")
+        return result
 
     def check_weights(self, action: str):
         return self._make_request("weights_checker", {"action": action})

miles/utils/memory_utils.py

@@ -18,8 +18,17 @@ def clear_memory(clear_host_memory: bool = False):
 def available_memory():
     device = torch.cuda.current_device()
     free, total = torch.cuda.mem_get_info(device)
+    # "gpu" is the LOCAL index inside this actor's CUDA_VISIBLE_DEVICES
+    # slice; report the physical id too so per-GPU logs are unambiguous
+    # when several actors each see their own GPU as cuda:0.
+    import os as _os
+
+    cvd = _os.environ.get("CUDA_VISIBLE_DEVICES", "")
+    visible = [x for x in cvd.split(",") if x.strip()]
+    physical = visible[device] if device < len(visible) else str(device)
     return {
         "gpu": str(device),
+        "physical_gpu": physical,
         "total_GB": _byte_to_gb(total),
         "free_GB": _byte_to_gb(free),
         "used_GB": _byte_to_gb(total - free),
@@ -32,6 +41,26 @@ def _byte_to_gb(n: int):
     return round(n / (1024**3), 2)
 
 
+def log_nontorch(label: str):
+    """Escape-audit probe: log the gap between whole-GPU physical usage and
+    this process's torch allocator — growth in that gap across a phase means
+    the phase allocated GPU memory OUTSIDE torch (context/module loading,
+    NCCL, driver pools) or outside the tms interesting region. Deltas between
+    consecutive probes attribute the ~2.5 GB unpausable tail phase by phase;
+    same-GPU co-tenants (the sleeping engine) are constant and cancel out.
+    """
+    device = torch.cuda.current_device()
+    torch.cuda.synchronize()
+    free, total = torch.cuda.mem_get_info(device)
+    used = _byte_to_gb(total - free)
+    reserved = _byte_to_gb(torch.cuda.memory_reserved(device))
+    info = available_memory()
+    logger.info(
+        f"[NONTORCH-AUDIT] {label}: physical_gpu={info['physical_gpu']} "
+        f"whole_used={used} reserved={reserved} non_torch={round(used - reserved, 2)}"
+    )
+
+
 def print_memory(msg, clear_before_print: bool = False):
     if clear_before_print:
         clear_memory()