Test: add higher atol in test_forward_with_num_logits_to_keep

[gante]

What does this PR do?

#31292 added test_forward_with_num_logits_to_keep. The test needs a higher atol at least for modern CUDA devices, it was failing locally on my end.

[Forward pass in the lm head with different shape due to num_logits_to_keep -> different order of FP operations (maybe different optimal kernels too) -> slightly different output. See this comment for more info.]

FYI @Cyrilvallez

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Cyrilvallez]

Thanks for fixing and pointing that out! I had no idea that input shape only could impact outputs of linear layers so significantly for lower precisions. Very good thing to keep in mind in the future.
Do you know why the test did not fail on the CIs when I committed? I must be missing something on how the tests are automatically run/what tests are actually run

[gante]

@Cyrilvallez

Do you know why the test did not fail on the CIs when I committed?

Our CI on push runs on a machine without GPUs :D CPU computations tend to be less sensible to these tiny fluctuations

[ArthurZucker]

thanks for fixing!