From 48b905406a6153437efc3599e7e2264a5c1cf952 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 11 May 2026 19:51:44 +0000 Subject: [PATCH] diag: add CUDA sync after mega_moe finalize + forward to catch errors --- patches/deepseek_v4.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/patches/deepseek_v4.py b/patches/deepseek_v4.py index f204a9b..8878c4d 100644 --- a/patches/deepseek_v4.py +++ b/patches/deepseek_v4.py @@ -887,6 +887,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): activation_clamp=activation_clamp, fast_math=fast_math, ) + # Sync to catch mega_moe CUDA errors immediately + torch.cuda.synchronize() DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] @@ -2481,7 +2483,12 @@ class DeepseekV4ForCausalLM(nn.Module): loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self.model.finalize_mega_moe_weights() + # Sync after mega_moe weight transform to catch CUDA errors early + torch.cuda.synchronize() + print("[NVFP4] mega_moe finalize_weights done, CUDA OK") self.model._convert_nvfp4_post_load() + torch.cuda.synchronize() + print("[NVFP4] post-load conversion done, CUDA OK") return loaded_params def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: