diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 94b2e7f3a..7a5553f4b 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -45,7 +45,6 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import get_moe_runner_backend from sglang.srt.layers.quantization import Fp8Config from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( @@ -79,44 +78,116 @@ def _log_eagle_accept_cp_draft_hidden_debug(key: str, message: str, *args): ) +def _load_safetensors_index_weight_map( + model_dir: str, revision: Optional[str] +) -> Optional[dict]: + """Return the safetensors index ``weight_map``, or None if it can't be read. + + Reads a local snapshot dir first (the common case — server-arg model paths are + resolved to local dirs at init); falls back to fetching just the index file + from the HF hub for repo-id paths. + """ + import glob + import json + import os + + if os.path.isdir(model_dir): + for index_path in glob.glob( + os.path.join(model_dir, "*.safetensors.index.json") + ): + try: + with open(index_path) as f: + return json.load(f).get("weight_map", {}) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to read %s: %s", index_path, e) + return None + + try: + import huggingface_hub + from huggingface_hub import hf_hub_download + + from sglang.srt.model_loader.weight_utils import ( + download_safetensors_index_file_from_hf, + ) + + index_file = "model.safetensors.index.json" + download_safetensors_index_file_from_hf(model_dir, index_file, None, revision) + index_path = hf_hub_download( + repo_id=model_dir, + filename=index_file, + revision=revision, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + with open(index_path) as f: + return json.load(f).get("weight_map", {}) + except Exception as e: # noqa: BLE001 + logger.debug("Could not fetch safetensors index for %s: %s", model_dir, e) + return None + + +def _nextn_moe_is_fp4_quantized_in_checkpoint( + config: PretrainedConfig, +) -> Optional[bool]: + """True/False if the checkpoint's NextN/MTP MoE experts ship fp4 scales, else None. + + The MTP/NextN decoder lives at HF prefix ``model.layers.{num_hidden_layers}`` + (the same index the loader uses to remap onto ``model.decoder``). fp4-quantized + MoE experts always carry ``weight_scale`` / ``weight_scale_2`` tensors; an + unquantized (bf16) MTP MoE (DeepSeek-R1-FP4 / V3-0324-FP4) does not. + """ + nextn_layer_id = getattr(config, "num_hidden_layers", None) + if nextn_layer_id is None: + return None + server_args = get_global_server_args() + if server_args.speculative_draft_model_path: + model_dir = server_args.speculative_draft_model_path + revision = server_args.speculative_draft_model_revision + else: + model_dir = server_args.model_path + revision = server_args.revision + if not model_dir: + return None + weight_map = _load_safetensors_index_weight_map(model_dir, revision) + if weight_map is None: + return None + expert_prefix = f"model.layers.{nextn_layer_id}.mlp.experts." + return any( + name.startswith(expert_prefix) and "weight_scale" in name + for name in weight_map + ) + + def _should_drop_nextn_modelopt_fp4_quant_config( quant_config: Optional[QuantizationConfig], + config: PretrainedConfig, ) -> bool: """Whether to drop the modelopt_fp4 quant config for the NextN/MTP module. - DeepSeek-R1-FP4 / V3-0324-FP4 ship an *unquantized* MTP module (no fp4 scales - for the MTP MoE, and their hf_quant_config does NOT list the MTP layer in - `ignore`); for those we must drop fp4 so the whole MTP loads bf16 (PR #7376). - GLM-5.2-NVFP4 instead ships fp4 routed experts on the MTP layer and runs on a - flashinfer fp4 MoE backend that consumes them — there we keep fp4 and let the - exclude-list make attn / indexer / shared_experts / gate bf16. + DeepSeek-R1-FP4 / V3-0324-FP4 ship an *unquantized* MTP MoE (no fp4 scales in + the checkpoint) and do NOT list the MTP layer in `ignore`; for those we drop + fp4 so the whole MTP loads bf16 (PR #7376). GLM-5.2-NVFP4 ships fp4 routed + experts on the MTP layer; there we keep fp4 and let the exclude-list make + attn / indexer / shared_experts / gate bf16. + + The decision is strictly per-checkpoint. The MoE runner backend is the WRONG + signal: it reads AUTO in the EAGLE draft worker, and an fp4-capable runner does + not imply an fp4 MTP (e.g. DeepSeek-V3-0324-FP4 runs flashinfer_trtllm yet ships + a bf16 MTP). So we inspect the checkpoint's safetensors index for the MTP MoE's + weight_scale tensors — the ground truth. """ if quant_config is None or quant_config.get_name() != "modelopt_fp4": return False - # The NextN/MTP draft is built by the EAGLE draft worker, where the effective - # MoE runner may be the *speculative* backend (which is AUTO when - # --speculative-moe-runner-backend is unset), not the main runner. Keep fp4 if - # EITHER the main or the speculative runner is a flashinfer fp4 backend. - from sglang.srt.layers.moe.utils import get_speculative_moe_runner_backend - - def _is_fp4_capable(b) -> bool: - return ( - b.is_flashinfer_trtllm() - or b.is_flashinfer_trtllm_routed() - or b.is_flashinfer_cutedsl() - or b.is_flashinfer_cutlass() + is_fp4 = _nextn_moe_is_fp4_quantized_in_checkpoint(config) + if is_fp4 is None: + # Couldn't read the index. Preserve the historical default (PR #7376): drop + # fp4 so the MTP loads bf16 — safe for R1/V3, and degrades a (would-be) fp4 + # MTP to bf16 rather than crashing on an unreadable checkpoint. + logger.warning( + "Could not determine NextN MoE fp4 quantization from the checkpoint " + "index; dropping modelopt_fp4 for the NextN module (bf16 MTP fallback)." ) - - main_backend = get_moe_runner_backend() - spec_backend = get_speculative_moe_runner_backend() - keep_fp4 = _is_fp4_capable(main_backend) or _is_fp4_capable(spec_backend) - logger.warning( - "[NextN-fp4] main_runner=%s spec_runner=%s keep_fp4=%s", - main_backend, - spec_backend, - keep_fp4, - ) - return not keep_fp4 + return True + return not is_fp4 class DeepseekModelNextN(nn.Module): @@ -137,7 +208,7 @@ class DeepseekModelNextN(nn.Module): else: moe_quant_config_override = None - if _should_drop_nextn_modelopt_fp4_quant_config(quant_config): + if _should_drop_nextn_modelopt_fp4_quant_config(quant_config, config): logger.warning( "Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model." )