feat: add load progress heartbeats to prevent k8s health check kills
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 ✓
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user