diag: add CUDA sync after mega_moe finalize + forward to catch errors

This commit is contained in:
2026-05-11 19:51:44 +00:00
parent 35f6b66678
commit 48b905406a

View File

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