Merge branch 'develop' into suport_spoling

Author: lizexu123

docs/parameters.md

@@ -48,6 +48,7 @@ When using FastDeploy to deploy models (including offline inference and service
 | ```dynamic_load_weight``` | `int` | Whether to enable dynamic weight loading, default: 0 |
 | ```enable_expert_parallel``` | `bool` | Whether to enable expert parallel |
 | ```enable_logprob``` | `bool` | Whether to enable return log probabilities of the output tokens or not. If true, returns the log probabilities of each output token returned in the content of message.If logrpob is not used, this parameter can be omitted when starting |
+| ```logprobs_mode``` | `str` | Indicates the content returned in the logprobs. Supported mode: `raw_logprobs`, `processed_logprobs`, `raw_logits`, `processed_logits`. Raw means the values before applying logit processors, like bad words. Processed means the values after applying such processors. |
 | ```served_model_name```| `str`| The model name used in the API. If not specified, the model name will be the same as the --model argument |
 | ```revision``` | `str` | 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. |
 | ```chat_template``` | `str` | Specify the template used for model concatenation, It supports both string input and file path input. The default value is None. If not specified, the model's default template will be used. |

docs/zh/parameters.md

@@ -46,6 +46,7 @@
 | ```dynamic_load_weight```          | `int`       | 是否动态加载权重，默认0 |
 | ```enable_expert_parallel```       | `bool`      | 是否启用专家并行 |
 | ```enable_logprob```       | `bool`      | 是否启用输出token返回logprob。如果未使用 logrpob，则在启动时可以省略此参数。 |
+| ```logprobs_mode```       | `str`      | 指定logprobs中返回的内容。支持的模式：`raw_logprobs`、`processed_logprobs'、`raw_logits`,`processed_logits'。processed表示logits应用温度、惩罚、禁止词处理后计算的logprobs。|
 | ```served_model_name```       | `str`      | API 中使用的模型名称，如果未指定，模型名称将与--model参数相同 |
 | ```revision```       | `str`      | 自动下载模型时，用于指定模型的Git版本，分支名或tag |
 | ```chat_template```       | `str`      | 指定模型拼接使用的模板，支持字符串与文件路径，默认为None，如未指定，则使用模型默认模板 |

fastdeploy/config.py

@@ -183,6 +183,7 @@ def __init__(
         self.max_model_len = 0
         self.dtype = "bfloat16"
         self.enable_logprob = False
+        self.logprobs_mode = "raw_logprobs"
         self.enable_redundant_experts = False
         self.redundant_experts_num = 0
         self.seed = 0

fastdeploy/engine/args_utils.py

@@ -367,6 +367,15 @@ class EngineArgs:
     Must be explicitly enabled via the `--enable-logprob` startup parameter to output logprob values.
     """
 
+    logprobs_mode: str = "raw_logprobs"
+    """
+    Indicates the content returned in the logprobs.
+    Supported mode:
+    1) raw_logprobs, 2) processed_logprobs, 3) raw_logits, 4) processed_logits.
+    Raw means the values before applying logit processors, like bad words.
+    Processed means the values after applying such processors.
+    """
+
     seed: int = 0
     """
     Random seed to use for initialization. If not set, defaults to 0.
@@ -412,6 +421,8 @@ def __post_init__(self):
         if self.enable_logprob:
             if not current_platform.is_cuda():
                 raise NotImplementedError("Only CUDA platform supports logprob.")
+            if self.speculative_config is not None and self.logprobs_mode.startswith("processed"):
+                raise NotImplementedError("processed_logprobs not support in speculative.")
         if self.speculative_config is not None:
             envs.ENABLE_V1_KVCACHE_SCHEDULER = 0
         if self.splitwise_role != "mixed" and self.cache_transfer_protocol != "rdma":
@@ -610,6 +621,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.enable_logprob,
             help="Enable output of token-level log probabilities.",
         )
+        model_group.add_argument(
+            "--logprobs-mode",
+            type=str,
+            choices=["raw_logprobs", "processed_logprobs", "processed_logits"],
+            default=EngineArgs.logprobs_mode,
+            help="Indicates the content returned in the logprobs.",
+        )
         model_group.add_argument(
             "--seed",
             type=int,

fastdeploy/engine/async_llm.py

@@ -830,6 +830,7 @@ def _start_worker_service(self):
             f" --runner {self.cfg.model_config.runner}"
             f" --convert {self.cfg.model_config.convert}"
             f" --override-pooler-config {self.cfg.model_config.override_pooler_config}"
+            f" --logprobs_mode {self.cfg.model_config.logprobs_mode}"
         )
 
         worker_append_flag = {

fastdeploy/engine/engine.py

@@ -532,6 +532,7 @@ def _start_worker_service(self):
             f" --runner {self.cfg.model_config.runner}"
             f" --convert {self.cfg.model_config.convert}"
             f" --override-pooler-config {self.cfg.model_config.override_pooler_config}"
+            f" --logprobs_mode {self.cfg.model_config.logprobs_mode}"
         )
 
         worker_append_flag = {

fastdeploy/model_executor/layers/sample/sampler.py

@@ -199,7 +199,7 @@ class Sampler(nn.Layer):
     Sampler for normal generation.
     """
 
-    def __init__(self, fd_config: FDConfig = None):
+    def __init__(self, fd_config: FDConfig = None, logprobs_mode: str = "raw_logprobs"):
         """ """
         super().__init__()
         if (
@@ -217,6 +217,7 @@ def __init__(self, fd_config: FDConfig = None):
             raise NotImplementedError
 
         self.processor = SamplerProcessor()
+        self.logprobs_mode = fd_config.model_config.logprobs_mode if fd_config is not None else logprobs_mode
         # Can only be created when fd_config.early_stopper_config.enable_early_stop = True
         if (
             fd_config is not None
@@ -335,7 +336,10 @@ def forward_cuda(
 
         num_logprobs = sampling_metadata.max_num_logprobs
         if num_logprobs is not None:
-            raw_logprobs = self.compute_logprobs(logits, sampling_metadata)
+            if self.logprobs_mode == "raw_logprobs":
+                raw_logprobs = self.compute_logprobs(logits, sampling_metadata)
+            elif self.logprobs_mode == "raw_logits":
+                raw_logprobs = logits.clone()
 
         logits = apply_penalty_multi_scores(
             sampling_metadata.pre_token_ids,
@@ -352,6 +356,12 @@ def forward_cuda(
             sampling_metadata.eos_token_ids,
         )
 
+        if num_logprobs is not None:
+            if self.logprobs_mode == "processed_logprobs":
+                raw_logprobs = self.compute_logprobs(logits, sampling_metadata)
+            elif self.logprobs_mode == "processed_logits":
+                raw_logprobs = logits.clone()
+
         probs = F.softmax(logits)
 
         probs = min_p_sampling(probs, sampling_metadata.min_p, sampling_metadata.min_p_list)
@@ -437,6 +447,7 @@ def __init__(self, fd_config: FDConfig):
             self.forward = self.forward_cuda
         else:
             raise NotImplementedError
+        self.logprobs_mode = fd_config.model_config.logprobs_mode
         self.speculative_verify_window = fd_config.speculative_config.verify_window
         self.speculative_max_candidate_len = fd_config.speculative_config.max_candidate_len
         self.speculative_benchmark_mode = fd_config.speculative_config.benchmark_mode
@@ -644,7 +655,10 @@ def forward_cuda(
                 share_inputs["seq_lens_encoder"],
                 share_inputs["accept_num"],
             )
-            raw_logprobs = self.compute_logprobs(target_logtis, sampling_metadata)
+            if self.logprobs_mode == "raw_logprobs":
+                raw_logprobs = self.compute_logprobs(target_logtis, sampling_metadata)
+            elif self.logprobs_mode == "raw_logits":
+                raw_logprobs = target_logtis.clone()
 
         logprobs_tensors = None
         token_ids = share_inputs["accept_tokens"]
@@ -677,6 +691,7 @@ def __init__(self, fd_config: FDConfig):
             self.forward = self.forward_cuda
         else:
             raise NotImplementedError
+        self.logprobs_mode = fd_config.model_config.logprobs_mode
 
     def pre_process(self, skip_idx_list: List[int] = []):
         """pre process before running"""
@@ -808,7 +823,12 @@ def forward_cuda(
         real_bsz = share_inputs["seq_lens_this_time"].shape[0]
         if num_logprobs is not None and share_inputs["substep"] == 0:
             real_token_num = share_inputs["batch_token_num"][:real_bsz].sum()
-            raw_logprobs = self.compute_logprobs(share_inputs["draft_logits"][:real_token_num, :], sampling_metadata)
+            if self.logprobs_mode == "raw_logprobs":
+                raw_logprobs = self.compute_logprobs(
+                    share_inputs["draft_logits"][:real_token_num, :], sampling_metadata
+                )
+            elif self.logprobs_mode == "raw_logits":
+                raw_logprobs = share_inputs["draft_logits"][:real_token_num, :].clone()
 
         logits = apply_speculative_penalty_multi_scores(
             sampling_metadata.pre_token_ids,

fastdeploy/worker/worker_process.py

@@ -614,6 +614,12 @@ def parse_args():
         action="store_true",
         help="Enable output of token-level log probabilities.",
     )
+    parser.add_argument(
+        "--logprobs_mode",
+        type=str,
+        default="raw_logprobs",
+        help="Indicates the content returned in the logprobs.",
+    )
     parser.add_argument(
         "--reasoning_parser",
         type=str,

tests/layers/test_sampler.py

@@ -15,13 +15,14 @@
 """
 
 import paddle
+import paddle.nn.functional as F
 
 from fastdeploy.model_executor.layers.sample.meta_data import SamplingMetadata
 from fastdeploy.model_executor.layers.sample.sampler import Sampler
 
 
 def _create_fake_logits(batch_size: int, vocab_size: int) -> paddle.Tensor:
-    fake_logits = paddle.full(shape=[batch_size, vocab_size], fill_value=1e-2, dtype="float32")
+    fake_logits = paddle.rand(shape=[batch_size, vocab_size], dtype="float32")
     return fake_logits
 
 
@@ -41,6 +42,7 @@ def _create_default_sampling_metadata(
     batch_size: int,
     min_seq_len: int,
     max_seq_len: int,
+    max_num_logprobs: int = None,
 ) -> SamplingMetadata:
 
     fake_sampling_metadata = SamplingMetadata(
@@ -59,6 +61,8 @@ def _create_default_sampling_metadata(
         min_p=paddle.randn([batch_size]),
         seed=paddle.to_tensor([[2025]]),
     )
+    if max_num_logprobs is not None:
+        fake_sampling_metadata.max_num_logprobs = max_num_logprobs
     return fake_sampling_metadata
 
 
@@ -75,5 +79,77 @@ def test_sampler():
     print(next_tokens)
 
 
+def get_baseline_logprobs(logits, sampling_metadata, logprobs_mode, token_ids):
+    if logprobs_mode == "raw_logprobs":
+        logprobs = F.log_softmax(logits, axis=-1)
+    elif logprobs_mode == "raw_logits":
+        logprobs = logits.clone()
+    elif logprobs_mode == "processed_logprobs":
+        from fastdeploy.model_executor.layers.sample.ops import (
+            apply_penalty_multi_scores,
+        )
+
+        logits = apply_penalty_multi_scores(
+            sampling_metadata.pre_token_ids,
+            sampling_metadata.prompt_ids,
+            sampling_metadata.prompt_lens,
+            logits,
+            sampling_metadata.repetition_penalties,
+            sampling_metadata.frequency_penalties,
+            sampling_metadata.presence_penalties,
+            sampling_metadata.temperature,
+            sampling_metadata.bad_words_token_ids,
+            sampling_metadata.step_idx,
+            sampling_metadata.min_dec_lens,
+            sampling_metadata.eos_token_ids,
+        )
+        logprobs = F.log_softmax(logits, axis=-1)
+    else:
+        from fastdeploy.model_executor.layers.sample.ops import (
+            apply_penalty_multi_scores,
+        )
+
+        logits = apply_penalty_multi_scores(
+            sampling_metadata.pre_token_ids,
+            sampling_metadata.prompt_ids,
+            sampling_metadata.prompt_lens,
+            logits,
+            sampling_metadata.repetition_penalties,
+            sampling_metadata.frequency_penalties,
+            sampling_metadata.presence_penalties,
+            sampling_metadata.temperature,
+            sampling_metadata.bad_words_token_ids,
+            sampling_metadata.step_idx,
+            sampling_metadata.min_dec_lens,
+            sampling_metadata.eos_token_ids,
+        )
+        logprobs = logits
+    token_logprobs = paddle.take_along_axis(logprobs, token_ids, axis=-1)
+    return token_logprobs
+
+
+def test_sampler_logprobs():
+    batch_size = 32
+    vocab_size = 1024
+    min_seq_len = 1
+    max_seq_len = 1024
+    logprobs_mode_list = ["raw_logprobs", "raw_logits", "processed_logprobs", "processed_logits"]
+    logits = _create_fake_logits(batch_size, vocab_size)
+    sampling_metadata = _create_default_sampling_metadata(batch_size, min_seq_len, max_seq_len, max_num_logprobs=0)
+    for logprobs_mode in logprobs_mode_list:
+        sampler = Sampler(logprobs_mode=logprobs_mode)
+        sampler_output = sampler(logits.clone(), sampling_metadata)
+        baseline_logprobs = get_baseline_logprobs(
+            logits.clone(), sampling_metadata, logprobs_mode=logprobs_mode, token_ids=sampler_output.sampled_token_ids
+        )
+        logprobs = sampler_output.logprobs_tensors.logprobs
+        print(f"baseline_logprobs = {baseline_logprobs}")
+        print(f"logprobs = {logprobs}")
+        equal = paddle.allclose(baseline_logprobs, logprobs, atol=1e-03, rtol=1e-03).item()
+        print(f"logprobs_mode: {logprobs_mode} equal={equal}")
+        assert equal
+
+
 if __name__ == "__main__":
     test_sampler()
+    test_sampler_logprobs()