From 11d3369a07ef79fd68355d301bbbefcc9b103b26 Mon Sep 17 00:00:00 2001 From: ispobock Date: Sun, 21 Jul 2024 00:32:40 +0800 Subject: [PATCH 1/2] fallback when sampling failed --- .../srt/managers/controller/infer_batch.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/controller/infer_batch.py b/python/sglang/srt/managers/controller/infer_batch.py index 0d50276cdf76..c0025a6eff4d 100644 --- a/python/sglang/srt/managers/controller/infer_batch.py +++ b/python/sglang/srt/managers/controller/infer_batch.py @@ -668,18 +668,16 @@ def sample(self, logits: torch.Tensor): max_top_k_round, batch_size = 32, probs.shape[0] uniform_samples = torch.rand((max_top_k_round, batch_size), device=probs.device) - batch_next_token_ids, _ = top_k_top_p_sampling_from_probs( + batch_next_token_ids, success = top_k_top_p_sampling_from_probs( probs, uniform_samples, self.top_ks, self.top_ps ) - # FIXME: this is a temporary fix for the illegal token ids - illegal_mask = torch.logical_or( - batch_next_token_ids < 0, batch_next_token_ids >= probs.shape[-1] - ) - if torch.any(illegal_mask): - warnings.warn("Illegal sampled token ids") - probs = probs.masked_fill(torch.isnan(probs), 0.0) - batch_next_token_ids = torch.argmax(probs, dim=-1) + if torch.any(~success): + warnings.warn("Sampling failed, fallback to top_k=1 strategy") + argmax_ids = torch.argmax(probs, dim=-1) + batch_next_token_ids = torch.where( + success, batch_next_token_ids, argmax_ids + ) if has_regex: batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy() From 87bdf7934f2b3f628656093a936329a2f603b52c Mon Sep 17 00:00:00 2001 From: ispobock Date: Sun, 21 Jul 2024 01:37:28 +0800 Subject: [PATCH 2/2] fill nan --- python/sglang/srt/managers/controller/infer_batch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/managers/controller/infer_batch.py b/python/sglang/srt/managers/controller/infer_batch.py index c0025a6eff4d..eda68ec462ae 100644 --- a/python/sglang/srt/managers/controller/infer_batch.py +++ b/python/sglang/srt/managers/controller/infer_batch.py @@ -674,6 +674,7 @@ def sample(self, logits: torch.Tensor): if torch.any(~success): warnings.warn("Sampling failed, fallback to top_k=1 strategy") + probs = probs.masked_fill(torch.isnan(probs), 0.0) argmax_ids = torch.argmax(probs, dim=-1) batch_next_token_ids = torch.where( success, batch_next_token_ids, argmax_ids