Renormalize LMNextToken.sample() probs to fix floating point errors

[gabegrand]

Due to numerical imprecision, the exponentiated next_token_logprobs occasionally do not sum to 1. When the discrepancy is too large, this causes np.random.choice to throw the following error:

    token_id = np.random.choice(len(probs), p=(probs))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "numpy/random/mtrand.pyx", line 975, in numpy.random.mtrand.RandomState.choice
ValueError: probabilities do not sum to 1

This is observed empirically when using token masking and appears to occur more frequently for smaller models (e.g., llama-3.2-1B).

This PR fixes this issue by renormalizing next_token_logprobs after exponentiation. This is intended as a stopgap fix pending future work to look into numerical stability in the token masking code.

[alex-lew]

Looks good, as a stopgap! We probably want to have some way of emitting a warning when this happens, so that we catch future bugs relating to logprob computation.

[gabegrand]

Thanks! I agree about the warning. I didn't one for now because I didn't want to introduce a ton of spam in the logs. In the future we could consider making better use of the logging / warnings modules, which allow suppression of repeated warnings through filters.