295bc17576
Co-authored-by: root <root@gpu-lg-cmc-h-h200-3047.host.h.pjlab.org.cn> Co-authored-by: chengshuang <chengshuang@pjlab.org.cn> Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
from typing import Any
|
|
|
|
from sglang.srt.configs.model_config import ModelConfig
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
|
|
class DllmConfig:
|
|
def __init__(
|
|
self,
|
|
algorithm: str,
|
|
algorithm_config: dict[str, Any],
|
|
block_size: int,
|
|
mask_id: int,
|
|
max_running_requests: int,
|
|
):
|
|
self.algorithm = algorithm
|
|
self.algorithm_config = algorithm_config
|
|
self.block_size = block_size
|
|
self.mask_id = mask_id
|
|
self.max_running_requests = max_running_requests
|
|
|
|
@staticmethod
|
|
def from_server_args(
|
|
server_args: ServerArgs,
|
|
):
|
|
if server_args.dllm_algorithm is None:
|
|
return None
|
|
|
|
model_config = ModelConfig.from_server_args(
|
|
server_args,
|
|
model_path=server_args.model_path,
|
|
model_revision=server_args.revision,
|
|
)
|
|
DLLM_PARAMS = {
|
|
"LLaDA2MoeModelLM": {"block_size": 32, "mask_id": 156895},
|
|
"SDARForCausalLM": {"block_size": 4, "mask_id": 151669},
|
|
"SDARMoeForCausalLM": {"block_size": 4, "mask_id": 151669},
|
|
}
|
|
|
|
arch = model_config.hf_config.architectures[0]
|
|
if arch in DLLM_PARAMS:
|
|
params = DLLM_PARAMS[arch]
|
|
block_size = params["block_size"]
|
|
mask_id = params["mask_id"]
|
|
else:
|
|
raise RuntimeError(f"Unknown diffusion LLM: {arch}")
|
|
|
|
max_running_requests = (
|
|
1
|
|
if server_args.max_running_requests is None
|
|
else server_args.max_running_requests
|
|
)
|
|
|
|
algorithm_config = {}
|
|
if server_args.dllm_algorithm_config is not None:
|
|
try:
|
|
import yaml
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Please install PyYAML to use YAML config files. "
|
|
"`pip install pyyaml`"
|
|
)
|
|
with open(server_args.dllm_algorithm_config, "r") as f:
|
|
algorithm_config = yaml.safe_load(f)
|
|
|
|
# Parse common algorithm configurations
|
|
block_size = algorithm_config.get("block_size", block_size)
|
|
|
|
return DllmConfig(
|
|
algorithm=server_args.dllm_algorithm,
|
|
algorithm_config=algorithm_config,
|
|
block_size=block_size,
|
|
mask_id=mask_id,
|
|
max_running_requests=max_running_requests,
|
|
)
|