Repository: pytorch
Issue: xpu: LayerNorm kernel produces different results on XPU than CPU/CUDA
Labels: triaged,module: xpu

With:
* PyTorch: https://github.com/pytorch/pytorch/commit/f2d388eddd25b5ab5b909a8b4cbf20ff34c6f3f2
* https://github.com/intel/torch-xpu-ops/commit/15f6d654685a98b6bbf6b76ee094258b16dd21ca
* https://github.com/huggingface/transformers/commit/1f6b423f0ce0c6e7afd300b59e5fd1a816e4896c
* https://github.com/huggingface/accelerate/commit/e11d3ceff3a49378796cdff5b466586d877d5c60

Running Huggingface `hiera` model with non-default `initializer_range=1e-10` (default is `0.02` and it works) on Pytorch XPU backend (I used Intel PVC platform) fails - output tensor contains `Nan` values. This behavior was first found running the following Huggingface Transformer tests for `hiera` model:
```
$ cat spec.py
import torch
DEVICE_NAME = 'xpu'
MANUAL_SEED_FN = torch.xpu.manual_seed
EMPTY_CACHE_FN = torch.xpu.empty_cache
DEVICE_COUNT_FN = torch.xpu.device_count

$ TRANSFORMERS_TEST_DEVICE_SPEC=spec.py python3 -m pytest --pspec tests/models/hiera/ -k torch_fx
...
FAILED tests/models/hiera/test_modeling_hiera.py::
    Here we also overwrite some of the tests of test_modeling_common.py, as Hiera does not use input_ids, inputs_embeds,
    attention_mask and seq_length.
    ::test_torch_fx - AssertionError: False is not true : traced 0th output doesn't match model 0th output for <class...
FAILED tests/models/hiera/test_modeling_hiera.py::
    Here we also overwrite some of the tests of test_modeling_common.py, as Hiera does not use input_ids, inputs_embeds,
    attention_mask and seq_length.
    ::test_torch_fx_output_loss - AssertionError: False is not true : traced 0th output doesn't match model 0th output for <class...
================================ 2 failed, 160 deselected in 2.32s =================================
```

Instead of running Huggingface Transformers, it's possible to reproduce the issue on the following script which is my attempt to extract essential behavior from HF test:

```
import copy
import random
import torch
from transformers import HieraConfig, HieraModel, PretrainedConfig

config = HieraConfig(
    embed_dim=8,
    image_size=[64, 64],
    patch_stride=[4, 4],
    patch_size=[7, 7],
    patch_padding=[3, 3],
    masked_unit_size=[8, 8],
    mlp_ratio=1.0,
    num_channels=3,
    depths=[1, 1, 1, 1],
    num_heads=[1, 1, 1, 1],
    embed_dim_multiplier=2.0,
    hidden_act='gelu',
    decoder_hidden_size=2,
    decoder_depth=1,
    decoder_num_heads=1,
    initializer_range=0.02)

setattr(config, "initializer_range", 1e-10)  # NOTE THIS LINE !!!

model = HieraModel(config).to("xpu")
model.eval()

def floats_tensor(shape, scale=1.0, rng=None, name=None):
    """Creates a random float32 tensor"""
    if rng is None:
        rng = random.Random()

    total_dims = 1
    for dim in shape:
        total_dims *= dim

    values = []
    for _ in range(total_dims):
        values.append(rng.random() * scale)

    return torch.tensor(data=values, dtype=torch.float, device="xpu").view(shape).contiguous()

inputs = {"pixel_values": floats_tensor([13, 3, 64, 64])}

model.config.use_cache = False

with torch.no_grad():
    outputs = model(**inputs)
print(outputs)
```

Issue appears when `initializer_range=1e-10`. That's what HF Transformers are doing in the tests (see [tests/test_modeling_common.py#L149](https://github.com/huggingface/transformers/blob/6c3f168b36882f0beebaa9121eafa1928ba29633/tests/test_modeling_common.py#L149)). Output on XPU will be:
```
HieraModelOutput(last_hidden_state=tensor([[[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]],
...
         nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]],
       device='xpu:0'), bool_masked_pos=None, ids_restore=None, hidden_states=None, attentions=None, reshaped_hidden_states=None)
```
And running with default `initializer_range=0,02` on XPU you will get reasonable results (named HF hiera tests will pass with this value, but other tests form hiera will start to fail):
```
HieraModelOutput(last_hidden_state=tensor([[[ 0.1861, -0.1193,  0.2276,  ...,  0.1105,  0.1467, -0.0326],
         [ 0.2299, -0.0409,  0.1859,  ...,  0.1770,  0.0710,  0.0243],
         [ 0.2048, -0.1447,  0.1758,  ...,  0.1698,  0.0585, -0.0254],
         [ 0.1917, -0.0898,  0.2141,  ...,  0.1629,  0.0251,  0.0306]],
...
          5.1592e-01,  2.1765e-01,  2.0610e+00,  1.5068e+00, -1.6449e+00,
          1.5775e-01,  7.1410e-01, -4.9712e-01, -8.2570e-01]], device='xpu:0'), bool_masked_pos=None, ids_restore=None, hidden_states=None, attentions=None, reshaped_hidden_states=None)
```

**Note that CUDA does handle `initializer_range=1e-10` in a different way** and does not return `Nan` tensors. However, what CUDA returns is suspiciously "good":
```
HieraModelOutput(last_hidden_state=tensor([[[3.0000e-10, 3.0000e-10, 3.0000e-10,  ..., 3.0000e-10,
          3.0000e-10, 3.0000e-10],
         [3.0000e-10, 3.0000e-10, 3.0000e-10,  ..., 3.0000e-10,
          3.0000e-10, 3.0000e-10],
...
         1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10,
         1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10,
         1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10,
         1.0000e-10, 1.0000e-10, 1.0000e-10, 1.0000e-10]], device='cuda:0'), bool_masked_pos=None, ids_restore=None, hidden_states=None, attentions=None, reshaped_hidden_states=None)
```

And for complete picture on CUDA, that's what it returns on `initializer_range=0,02`:
```
HieraModelOutput(last_hidden_state=tensor([[[-0.0146,  0.1575,  0.0880,  ...,  0.1405,  0.0805, -0.1740],
         [ 0.0672,  0.1291,  0.0963,  ...,  0.1333,  0.1244, -0.0825],
         [-0.0050,  0.1595,  0.0767,  ...,  0.1392,  0.1436, -0.1975],
         [ 0.0101,  0.1367,  0.0705,  ...,  0.1299,  0.1178, -0.1880]],
...
         -2.5323e-01,  9.3338e-01, -5.0501e-01, -4.9314e-01,  7.3834e-01,
          9.5067e-02,  1.7623e-01,  1.8487e+00, -1.5834e-01,  9.8299e-01,
          1.8513e+00,  1.1275e+00,  1.8372e+00,  1.0047e+00,  1.9499e-01,
         -1.7089e+00,  4.8383e-01,  1.8169e-01, -2.4395e+00]], device='cuda:0'), bool_masked_pos=None, ids_restore=None, hidden_states=None, attentions=None, reshaped_hidden_states=None)
```

CC: @gujinghui @EikanWang @fengyuan14 @guangyey @jgong5 @xytintel @faaany

cc @gujinghui @EikanWang @fengyuan14 @guangyey