DeepInfer/deepinfer/config.py

# SPDX-License-Identifier: Apache-2.0
# Derived from vLLM; modified by DeepSeek.

import torch

from deepinfer import envs
from deepinfer.engine.engine_type import EngineType
from deepinfer.model_config_loader import load_model_config


def _get_runtime_max_model_len(max_model_len: int) -> int:
    """Return the exclusive upper bound of every runtime position."""
    nextn = envs.DEEPINFER_USE_NEXTN_DECODE
    scheduler_overhead = nextn + (nextn + 1 if envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING else 0)
    draft_overhead = nextn
    return max_model_len + scheduler_overhead + draft_overhead


class ModelConfig:
    """Configuration for the model.

    Args:
        model: Name or path of the huggingface model to use.
        tokenizer: Name or path of the huggingface tokenizer to use.
        trust_remote_code: Trust remote code (e.g., from HuggingFace) when
            downloading the model and tokenizer.
        load_format: Weight initialization mode:
            "auto" loads the HF safetensors index and shards from the model
                directory.
            "dummy" initializes synthetic weights for profiling, using the
                same HF model configuration.
        seed: Random seed for reproducibility.
        revision: The specific model version to use. It can be a branch name,
            a tag name, or a commit id. If unspecified, will use the default
            version.
        tokenizer_revision: The specific tokenizer version to use. It can be a
            branch name, a tag name, or a commit id. If unspecified, will use
            the default version.
        max_model_len: Maximum length of a sequence (including prompt and
            output). If None, will be derived from the model.
    """

    def __init__(
        self,
        model: str,
        tokenizer: str,
        trust_remote_code: bool,
        load_format: str,
        seed: int,
        max_num_seqs: int,
        revision: str | None = None,
        tokenizer_revision: str | None = None,
        max_model_len: int | None = None,
        max_num_batched_tokens: int | None = None,
        ignore_len: int = -1,
        served_model_name: str = "deepinfer-model",
        engine_type: EngineType = EngineType.NORMAL,
    ) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.trust_remote_code = trust_remote_code
        self.load_format = load_format
        self._verify_load_format()
        self.seed = seed
        self.revision = revision
        self.tokenizer_revision = tokenizer_revision

        self.max_num_seqs = max_num_seqs
        self.max_model_len = max_model_len if max_model_len is not None else 4096
        self.runtime_max_model_len = _get_runtime_max_model_len(self.max_model_len)
        self.max_num_batched_tokens = _get_max_num_batched_tokens(self.max_model_len, max_num_batched_tokens, max_num_seqs)

        self.ignore_len = ignore_len
        self.served_model_name = served_model_name
        self.dtype = torch.bfloat16
        self.use_compressed_tokenizer = False
        self.engram_tokenizer = self.tokenizer
        self.engine_type = engine_type

        self._load_model_cfg()

    _GPT_RUNTIME_ATTRS = (
        "max_num_seqs",
        "max_num_batched_tokens",
        "runtime_max_model_len",
        "engine_type",
    )

    def _load_model_cfg(self) -> None:
        self.model_cfg = load_model_config(self.model)
        self.architecture = self.model_cfg.architecture

        for attr in self._GPT_RUNTIME_ATTRS:
            setattr(self.model_cfg.gpt, attr, getattr(self, attr))

    def __getstate__(self) -> dict:
        # model_cfg 不可 pickle，改为只传递必要的运行时标量，
        # 由 worker 进程在 __setstate__ 中重新 load model_cfg 后回填。
        state = self.__dict__.copy()
        model_cfg = state.pop("model_cfg", None)
        overrides: dict = {}
        if model_cfg is not None:
            overrides = {k: model_cfg.gpt[k] for k in self._GPT_RUNTIME_ATTRS}
        state["_gpt_runtime_overrides"] = overrides
        return state

    def __setstate__(self, newstate: dict) -> None:
        overrides = newstate.pop("_gpt_runtime_overrides", {})
        self.__dict__.update(newstate)
        self.model_cfg = load_model_config(self.model)
        self.architecture = self.model_cfg.architecture
        for key, value in overrides.items():
            setattr(self.model_cfg.gpt, key, value)

    def _verify_load_format(self) -> None:
        load_format = self.load_format.lower()
        supported_load_format = ["auto", "dummy"]
        assert load_format in supported_load_format
        self.load_format = load_format

    def get_vocab_size(self) -> int:
        return self.model_cfg.gpt.vocab_size


class CacheConfig:
    def __init__(
        self,
        block_size: int,
        gpu_memory_utilization: float,
        lookahead_tokens: int = 0,
    ) -> None:
        self.block_size = block_size
        self.gpu_memory_utilization = gpu_memory_utilization
        self._verify_args()

        # Will be set after profiling.
        self.num_gpu_blocks = None
        # Will be set after init cache.
        self.num_state_seqs = None
        self.bytes_of_state_per_seqs = None
        self.bytes_per_block = None

        self.lookahead_tokens = lookahead_tokens

    def _verify_args(self) -> None:
        if self.gpu_memory_utilization > 1.0:
            raise ValueError(f"GPU memory utilization must be less than 1.0. Got {self.gpu_memory_utilization}.")


class ParallelConfig:
    def __init__(
        self,
        data_parallel_size: int,
        context_parallel_size: int,
        expert_parallel_size: int,
        engram_parallel_size: int,
        engine_rank: int,
        engine_group_rank: int,
        engine_type: EngineType = EngineType.NORMAL,
    ) -> None:
        if data_parallel_size == 0:
            data_parallel_size = max(1, expert_parallel_size // context_parallel_size)
        self.data_parallel_size = data_parallel_size
        self.context_parallel_size = context_parallel_size
        self.expert_parallel_size = expert_parallel_size
        self.engine_rank = engine_rank
        self.data_parallel_rank = engine_rank % data_parallel_size
        self.engine_group_rank = engine_group_rank
        self.engine_type = engine_type

        self.model_parallel_size = self.context_parallel_size
        self.engram_parallel_size = engram_parallel_size or (self.model_parallel_size * data_parallel_size)
        self.world_size = max(self.model_parallel_size * data_parallel_size, self.engram_parallel_size)
        self._verify_args()

    def _verify_args(self) -> None:
        assert self.expert_parallel_size > 0
        assert (self.model_parallel_size * self.data_parallel_size) % self.expert_parallel_size == 0
        assert self.engram_parallel_size > 0
        assert self.world_size % (self.model_parallel_size * self.data_parallel_size) == 0
        assert self.world_size % self.engram_parallel_size == 0
        assert 0 <= self.engine_rank < self.world_size // self.model_parallel_size


class SchedulerConfig:
    def __init__(
        self,
        max_num_batched_tokens: int | None,
        max_num_seqs: int,
        max_model_len: int,
        stop_preempted_request: bool = False,
        chunked_prefill: bool = False,
        prefill_attempt_interval: int = envs.DEEPINFER_PREFILL_ATTEMPT_INTERVAL,
        prefill_any_attempt_interval: int = envs.DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL,
    ) -> None:
        self.max_num_batched_tokens = _get_max_num_batched_tokens(max_model_len, max_num_batched_tokens, max_num_seqs)
        self.max_num_seqs = max_num_seqs
        self.max_model_len = max_model_len
        self.runtime_max_model_len = _get_runtime_max_model_len(max_model_len)
        self.stop_preempted_request = stop_preempted_request
        self.chunked_prefill = chunked_prefill
        # NORMAL 引擎每隔多少个 decode step 尝试一次 prefill（ALL 投票点：全员有活才做）
        self.prefill_attempt_interval = prefill_attempt_interval
        # NORMAL 引擎 ANY 保底投票点间隔：任一 rank 有活即集体 prefill，防饿死
        self.prefill_any_attempt_interval = prefill_any_attempt_interval
        self._verify_args()

    def _verify_args(self) -> None:
        if (not self.chunked_prefill) and self.max_num_batched_tokens < self.max_model_len:
            raise ValueError(
                f"max_num_batched_tokens ({self.max_num_batched_tokens}) is smaller than max_model_len ({self.max_model_len}). This effectively limits the maximum sequence length to max_num_batched_tokens and makes DeepInfer reject longer sequences. Please increase max_num_batched_tokens or decrease max_model_len."
            )
        if self.prefill_attempt_interval < 1:
            raise ValueError(f"prefill_attempt_interval ({self.prefill_attempt_interval}) should be >= 1")
        if self.prefill_any_attempt_interval < 1:
            raise ValueError(f"prefill_any_attempt_interval ({self.prefill_any_attempt_interval}) should be >= 1")


def _get_max_num_batched_tokens(
    max_model_len: int,
    max_num_batched_tokens: int | None,
    max_num_seqs: int,
) -> int:
    max_num_decoding_tokens = max_num_seqs * (1 + envs.DEEPINFER_USE_NEXTN_DECODE)
    if max_num_batched_tokens is not None:
        assert max_num_batched_tokens >= max_num_decoding_tokens
        return max_num_batched_tokens
    return max(max_model_len, 4096, max_num_decoding_tokens)