[Bug]: apply_temperature may cause nan in probs

[courage17340]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

Recently, we met some illegal memory access problems. When we turned off flashinfer sampling, the problem disappeared, or at least it was much harder to trigger. Unfortunately, I can't provide a minimal reproduction here.

After some investigation, I believe the problem is caused by apply_temperature. Suppose we have a decode batch with both random sampling and greedy sampling. Here will set the temperature for greedy sampling to -1.0, and here will div all the logits against temperature. If the logits (for greedy sampling) contains -inf, it will become inf after the division, and hence the probs tensor will be fully nan, I guess maybe it will cause bad things when using flashinfer top_p sampling. I'm not 100% sure about this, because I use probs with nan to run against flashinfer top_p sampling, and compute-sanitizer can't find any problems. But anyway, nan in probs is dangerous, I think it should be fixed.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[youkaichao]

cc @njhill
I think the divide by -1.0 is introduced by #13587

[njhill]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

@courage17340 do you mean it impacts all rows of the logits tensor not just those whose logits contain inf?

[youkaichao]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

we do have lots of reports about IMA from sampling, and we don't know what's the root cause yet. this issue description looks reasonable.

What's the original motivation for dividing by -1.0 ? it seems to be incorrect from my understanding.

[courage17340]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

@courage17340 do you mean it impacts all rows of the logits tensor not just those whose logits contain inf?

It looks like all of the conditions are required to reproduce:

1. the batch should contain both greedy sampling (for temperature=-1) and random sampling (otherwise all_greedy is satisfied and apply_temperature is not executed)
2. the original logits for greedy sampling contains -inf, normally set by some logits modification (bad words or others?)
3. flashinfer sampling is used, which means the batch contains top_p/top_k sampling

The whole batch will then be sampled by flashinfer top_p/top_k sampling kernel, and the rows with inf in logits will get nan in probs, and other rows looks OK. I don't know what happens next because compute-sanitizer can't find errors when I just run the kernel on the dumped tensors directly.

[njhill]

But these rows are the ones that were already separately greedy sampled and the corresponding logits in this method are discarded (filtered out by the subsequent where here). Do you have a test case where nan values are actually returned to the user?

[courage17340]

But these rows are the ones that were already separately greedy sampled and the corresponding logits in this method are discarded (filtered out by the subsequent where here). Do you have a test case where nan values are actually returned to the user?

I guess that nan in probs may trigger some corner cases and some kernel (like flashinfer top_p sampling) may access invalid indices. However, I still can't verify that yet.

[zifeitong]

Thanks @courage17340, I think your theory is correct.

@njhill I think the key issue is torch.softmax() called by flashinfer will return nan for the entire row if there are any inf in the row, which cause flahsinfer crashing. The filter is too late at this moment.

I tested a hack by clamping logits to [float.MIN, float.MAX] on our workload and it fixes the crashes in #19483. The proper fix should be avoiding creating inf in the first place.

[zifeitong]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

I think it's not that easy to trigger this bug, you'll need both greedy sampling (which is never default setting) and sampling params causing "-inf" (bad words, structured output) mixed in a single batch.

[ashgold]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

@courage17340 do you mean it impacts all rows of the logits tensor not just those whose logits contain inf?

It was said that it was six months ago, so I doubt that this is the cause.
This problem also occurred at v0.8.4.
It does not occur at 0.7.2.
This problem occurs very occasionally only in operational situations, so it is very difficult to reproduce, so it does not seem to have become an issue so far.

[njhill]

Hmm, that change was made 6 months ago so it's a bit surprising the problem hasn't shown up earlier if that's the cause...

I think it's not that easy to trigger this bug, you'll need both greedy sampling (which is never default setting) and sampling params causing "-inf" (bad words, structured output) mixed in a single batch.

@courage17340 @zifeitong We should be able to create a reproducer with all of these things though? (i.e. which also uses flashifer sampling)