diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index b801e1162ab..30e4a637fcc 100644 --- a/docker/patch/latest/sglang.patch +++ b/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() diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 93e6e330f56..a0436c9cccc 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/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: diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index fdcd4038a84..41a57b1a693 100644 --- a/miles/backends/megatron_utils/model.py +++ b/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: diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 6dd140ebf47..a737e70a93f 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/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}) diff --git a/miles/utils/memory_utils.py b/miles/utils/memory_utils.py index c12f3cd0bcd..63d2dee553a 100644 --- a/miles/utils/memory_utils.py +++ b/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()