From a569612df5f429de51e5432446aedf10478c4fd4 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 05:51:35 +0000 Subject: [PATCH] feat: add load progress heartbeats to prevent k8s health check kills MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 5-minute gap after safetensors load is GPU weight upload — no output, k8s marks the pod unhealthy. Now prints a heartbeat every 256 weight loads during the expert loading phase. Also adds checkpoint-ready and model-ready prints around finalize: Checkpoint loaded. Transferring weights to GPU & preparing NVFP4... (JIT compile)NVFP4 MoE layers: 50%|██████████░░░░░░░░░░| 31/61 NVFP4 model ready ✓ --- vllm/patches/deepseek_v4.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index e3a01fdf..2e1b4e70 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -222,6 +222,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): This replaces the broken C++ CUTLASS kernel (see README.md for the full story). """ _cutedsl_runner: 'CuTeDSLMoERunner | None' = None + _weight_load_count: int = 0 # NVFP4 E2M1 lookup table (positive values, sign from bit 3) E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] @@ -351,6 +352,12 @@ class DeepseekV4MegaMoEExperts(nn.Module): shard_id: str, expert_id: int, ) -> bool: + # Heartbeat: print every 256 weight loads so k8s/docker + # don't think the pod is dead during GPU upload + DeepseekV4MegaMoEExperts._weight_load_count += 1 + if DeepseekV4MegaMoEExperts._weight_load_count % 256 == 1: + print(f" Loading expert weights... ({DeepseekV4MegaMoEExperts._weight_load_count})", flush=True) + local_expert_id = self._map_global_expert_id(expert_id) if local_expert_id == -1: return False @@ -2200,8 +2207,10 @@ class DeepseekV4ForCausalLM(nn.Module): # AutoWeightsLoader bypasses this logic and would break NVFP4 loading. loaded_params = self.model.load_weights(rest) loaded_params.add("lm_head.weight") + print(" Checkpoint loaded. Transferring weights to GPU & preparing NVFP4...", flush=True) self.model.finalize_mega_moe_weights() self.model._convert_nvfp4_post_load() + print(" NVFP4 model ready ✓", flush=True) return loaded_params