introduce arg_groups/ with nemotron_h hook (#24328)
This commit is contained in:
@@ -0,0 +1,51 @@
|
|||||||
|
import logging
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from sglang.srt.utils.common import is_sm100_supported
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sglang.srt.server_args import ServerArgs
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def apply_nemotron_h_defaults(server_args: "ServerArgs", model_arch: str) -> None:
|
||||||
|
"""Apply NemotronH model-specific server arg defaults and constraints."""
|
||||||
|
model_config = server_args.get_model_config()
|
||||||
|
if model_config.quantization in [
|
||||||
|
"modelopt",
|
||||||
|
"modelopt_fp8",
|
||||||
|
"modelopt_fp4",
|
||||||
|
"modelopt_mixed",
|
||||||
|
]:
|
||||||
|
assert model_config.hf_config.mlp_hidden_act == "relu2"
|
||||||
|
if model_config.quantization == "modelopt":
|
||||||
|
quant_algo = model_config.hf_config.quantization_config["quant_algo"]
|
||||||
|
if quant_algo == "MIXED_PRECISION":
|
||||||
|
server_args.quantization = "modelopt_mixed"
|
||||||
|
else:
|
||||||
|
server_args.quantization = (
|
||||||
|
"modelopt_fp4" if quant_algo == "NVFP4" else "modelopt_fp8"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
server_args.quantization = model_config.quantization
|
||||||
|
if server_args.moe_runner_backend == "auto":
|
||||||
|
if is_sm100_supported() and server_args.moe_a2a_backend == "none":
|
||||||
|
server_args.moe_runner_backend = "flashinfer_trtllm"
|
||||||
|
logger.info(
|
||||||
|
"Use flashinfer_trtllm as MoE runner backend on sm100 for "
|
||||||
|
f"{model_arch}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
server_args.moe_runner_backend = "flashinfer_cutlass"
|
||||||
|
|
||||||
|
server_args._handle_mamba_radix_cache(
|
||||||
|
model_arch=model_arch,
|
||||||
|
support_mamba_cache=True,
|
||||||
|
support_mamba_cache_extra_buffer=False,
|
||||||
|
sm100_default_attention_backend="flashinfer",
|
||||||
|
)
|
||||||
|
assert server_args.attention_backend != "triton", (
|
||||||
|
"NemotronHForCausalLM does not support triton attention backend,"
|
||||||
|
"as the first layer might not be an attention layer"
|
||||||
|
)
|
||||||
@@ -2143,46 +2143,11 @@ class ServerArgs:
|
|||||||
support_mamba_cache=False,
|
support_mamba_cache=False,
|
||||||
)
|
)
|
||||||
elif model_arch in ["NemotronHForCausalLM"]:
|
elif model_arch in ["NemotronHForCausalLM"]:
|
||||||
model_config = self.get_model_config()
|
from sglang.srt.arg_groups.nemotron_h_hook import (
|
||||||
if model_config.quantization in [
|
apply_nemotron_h_defaults,
|
||||||
"modelopt",
|
|
||||||
"modelopt_fp8",
|
|
||||||
"modelopt_fp4",
|
|
||||||
"modelopt_mixed",
|
|
||||||
]:
|
|
||||||
assert model_config.hf_config.mlp_hidden_act == "relu2"
|
|
||||||
if model_config.quantization == "modelopt":
|
|
||||||
quant_algo = model_config.hf_config.quantization_config[
|
|
||||||
"quant_algo"
|
|
||||||
]
|
|
||||||
if quant_algo == "MIXED_PRECISION":
|
|
||||||
self.quantization = "modelopt_mixed"
|
|
||||||
else:
|
|
||||||
self.quantization = (
|
|
||||||
"modelopt_fp4" if quant_algo == "NVFP4" else "modelopt_fp8"
|
|
||||||
)
|
)
|
||||||
else:
|
|
||||||
self.quantization = model_config.quantization
|
|
||||||
if self.moe_runner_backend == "auto":
|
|
||||||
if is_sm100_supported() and self.moe_a2a_backend == "none":
|
|
||||||
self.moe_runner_backend = "flashinfer_trtllm"
|
|
||||||
logger.info(
|
|
||||||
"Use flashinfer_trtllm as MoE runner backend on sm100 for "
|
|
||||||
f"{model_arch}"
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
self.moe_runner_backend = "flashinfer_cutlass"
|
|
||||||
|
|
||||||
self._handle_mamba_radix_cache(
|
apply_nemotron_h_defaults(self, model_arch)
|
||||||
model_arch=model_arch,
|
|
||||||
support_mamba_cache=True,
|
|
||||||
support_mamba_cache_extra_buffer=False,
|
|
||||||
sm100_default_attention_backend="flashinfer",
|
|
||||||
)
|
|
||||||
assert self.attention_backend != "triton", (
|
|
||||||
"NemotronHForCausalLM does not support triton attention backend,"
|
|
||||||
"as the first layer might not be an attention layer"
|
|
||||||
)
|
|
||||||
elif model_arch in [
|
elif model_arch in [
|
||||||
"Qwen3MoeForCausalLM",
|
"Qwen3MoeForCausalLM",
|
||||||
"Qwen3VLMoeForConditionalGeneration",
|
"Qwen3VLMoeForConditionalGeneration",
|
||||||
|
|||||||
Reference in New Issue
Block a user