[Bugfix] Fix entropy calculation bugs (#5941)

* fix entropy bugs

Author: Deleter-D

fastdeploy/model_executor/entropy_utils.py

@@ -19,6 +19,19 @@
 from fastdeploy.utils import data_processor_logger
 
 
+def get_entropy(logits):
+    # Check for -inf values in logits
+    if paddle.any(paddle.isinf(logits) & (logits < 0)):
+        data_processor_logger.debug("Detected -inf values in logits, clipping to minimum value")
+        logits = paddle.clip(logits, min=1e-9)
+
+    a0 = logits - paddle.max(logits, axis=-1, keepdim=True)
+    ea0 = paddle.exp(a0)
+    z0 = paddle.sum(ea0, axis=-1, keepdim=True)
+    p0 = ea0 / z0
+    return paddle.sum(p0 * (paddle.log(z0) - a0), axis=-1)
+
+
 def calculate_logits_entropy(logits, share_inputs, temperature):
     real_bsz = share_inputs["seq_lens_this_time"].shape[0]
     real_seq_lens = paddle.where(
@@ -27,13 +40,6 @@ def calculate_logits_entropy(logits, share_inputs, temperature):
         share_inputs["seq_lens_this_time"].squeeze(1),
     )
 
-    def get_entropy(logits):
-        a0 = logits - paddle.max(logits, axis=-1, keepdim=True)
-        ea0 = paddle.exp(a0)
-        z0 = paddle.sum(ea0, axis=-1, keepdim=True)
-        p0 = ea0 / z0
-        return paddle.sum(p0 * (paddle.log(z0) - a0), axis=-1)
-
     batch_indices = paddle.arange(real_bsz, dtype="int32")
     batch_id_per_token = paddle.repeat_interleave(batch_indices, real_seq_lens)
     for i in range(logits.shape[0]):
@@ -77,13 +83,6 @@ def speculate_calculate_logits_entropy(logits, share_inputs, temperature):
     for i in range(total_accepted_num):
         accepted_logits[i] = logits[accepted_idx[i]]
 
-    def get_entropy(logits):
-        a0 = logits - paddle.max(logits, axis=-1, keepdim=True)
-        ea0 = paddle.exp(a0)
-        z0 = paddle.sum(ea0, axis=-1, keepdim=True)
-        p0 = ea0 / z0
-        return paddle.sum(p0 * (paddle.log(z0) - a0), axis=-1)
-
     batch_indices = paddle.arange(share_inputs["accept_num"].shape[0], dtype="int32")
     batch_id_per_token = paddle.repeat_interleave(batch_indices, share_inputs["accept_num"])
     for i in range(accepted_logits.shape[0]):

fastdeploy/worker/gpu_model_runner.py

@@ -1875,15 +1875,14 @@ def _dummy_sampler_run(
                     group=self.parallel_config.tp_group,
                 )
         else:
-            self.sampler(
+            sampler_output = self.sampler(
                 logits,
                 self.sampling_metadata,
                 self.model_config.max_model_len,
                 self.share_inputs,
                 accept_all_drafts,
                 reject_all_drafts,
             )
-            sampler_output = None
             if self.parallel_config.tensor_parallel_size > 1:
                 paddle.distributed.broadcast(
                     self.share_inputs["accept_tokens"],

tests/model_executor/test_entropy_utils.py

@@ -107,6 +107,34 @@ def test_entropy_list_clear(self):
 
         self.assertAlmostEqual(share_inputs["entropy_list"][2][0], 0.0003187173861078918, places=6)
 
+    def test_negative_inf_clip(self):
+        share_inputs = {
+            "seq_lens_this_time": paddle.to_tensor([[1], [0], [15]], dtype="int32"),
+            "seq_lens_encoder": paddle.to_tensor([[0], [0], [15]], dtype="int32"),
+            "seq_lens_decoder": paddle.to_tensor([[30], [0], [15]], dtype="int32"),
+            "entropy_list": [[], [], []],
+            "stop_flags": paddle.to_tensor([[False], [True], [False]], dtype="bool"),
+            "req_ids": ["req_1", "req_2", "req_3"],
+        }
+
+        logits = paddle.to_tensor(
+            [
+                [10.0, 1.0, -float("inf")],
+                [1.0, 1.0, -float("inf")],
+            ],
+            dtype="float32",
+        )
+        temperature = paddle.ones([3], dtype="float32")
+
+        calculate_logits_entropy(logits, share_inputs, temperature)
+
+        self.assertEqual(len(share_inputs["entropy_list"][0]), 1)
+        self.assertEqual(len(share_inputs["entropy_list"][1]), 0)
+        self.assertEqual(len(share_inputs["entropy_list"][2]), 1)
+
+        self.assertAlmostEqual(share_inputs["entropy_list"][0][0], 0.0017332095885649323, places=6)
+        self.assertAlmostEqual(share_inputs["entropy_list"][2][0], 1.017357349395752, places=6)
+
 
 class TestSpeculateCalculateLogitsEntropy(unittest.TestCase):
 
@@ -207,6 +235,34 @@ def test_entropy_list_clear(self):
 
         self.assertAlmostEqual(share_inputs["entropy_list"][1][0], 0.0024676250759512186, places=6)
 
+    def test_negative_inf_clip(self):
+        share_inputs = {
+            "seq_lens_this_time": paddle.to_tensor([[1], [0], [15]], dtype="int32"),
+            "seq_lens_encoder": paddle.to_tensor([[0], [0], [15]], dtype="int32"),
+            "seq_lens_decoder": paddle.to_tensor([[30], [0], [15]], dtype="int32"),
+            "entropy_list": [[], [], []],
+            "stop_flags": paddle.to_tensor([[False], [True], [False]], dtype="bool"),
+            "req_ids": ["req_1", "req_2", "req_3"],
+        }
+
+        logits = paddle.to_tensor(
+            [
+                [10.0, 1.0, -float("inf")],
+                [1.0, 1.0, -float("inf")],
+            ],
+            dtype="float32",
+        )
+        temperature = paddle.ones([3], dtype="float32")
+
+        calculate_logits_entropy(logits, share_inputs, temperature)
+
+        self.assertEqual(len(share_inputs["entropy_list"][0]), 1)
+        self.assertEqual(len(share_inputs["entropy_list"][1]), 0)
+        self.assertEqual(len(share_inputs["entropy_list"][2]), 1)
+
+        self.assertAlmostEqual(share_inputs["entropy_list"][0][0], 0.0017332095885649323, places=6)
+        self.assertAlmostEqual(share_inputs["entropy_list"][2][0], 1.017357349395752, places=6)
+
 
 if __name__ == "__main__":
     unittest.main()