Replace inefficient torch.sqrt taking scalar input with numpy.sqrt
Author: FindHao
Pull request: yes

# What does this PR do?

This PR fixes a performance issue for the model reformer. In model reformer running on GPUs, there are two calls for `torch.sqrt` and `torch.rsqrt` of a scalar `self.attention_head_size`. From the profiling results, we found that it calls `aten::to`, `aten::copy`, and `cudaStreamSynchronize`  etc. It is because `torch.sqrt` and `torch.rsqrt` only accept tensors as arguments, so it has to change the scalar to a tensor with only one element. By replacing them with `numpy.sqrt`, those extra operations are removed. Here is an [example script](https://gist.github.com/FindHao/abbaa7a8a0b74173d9eb2233684bb6a9) to test the impacted function. 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker and @younesbelkada

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts and @NielsRogge
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: HF Trainer: @stas00, Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger, @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: @sgugger
- TensorFlow: @Rocketknight1

 -->


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Thanks for working on this. Can you add the improvements that you have in terms of performances? (this will be useful if we look back at this code to know the motivation).

> Thanks for working on this. Can you add the improvements that you have in terms of performances? (this will be useful if we look back at this code to know the motivation).

@ArthurZucker @sgugger Thanks for your approval! Since this function call doesn't take a big portion of the whole model, the speedup at the model level is trivial. But at the function level, function `_len_and_dim_norm` obtains 27X speedup and removes CPU-GPU synchronizations.

Thanks for the number!