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:
2026-05-16 05:51:35 +00:00
parent e5370140cb
commit a569612df5

View File

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