Odd output from falcon-mamba-7b on mps device
Author: derekelewis
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.45.0.dev0
- Platform: macOS-14.6.1-arm64-arm-64bit
- Python version: 3.12.4
- Huggingface_hub version: 0.24.5
- Safetensors version: 0.4.4
- Accelerate version: 0.33.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: No

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "tiiuae/falcon-mamba-7b"
torch_dtype = torch.bfloat16

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype).to("mps")
print(model)

tokenizer = AutoTokenizer.from_pretrained(model_id)

input_sequence = tokenizer.encode("the largest crater on the moon is", return_tensors="pt").to('mps')
print(input_sequence)

output_sequence = model.generate(input_sequence, max_new_tokens=100, do_sample=True)

print(tokenizer.decode(output_sequence[0]))
```

Output:

```
the largest crater on the moon is�haped, and the number “13” is very common in occult and Satanic symbology. It has been proposed that some Illuminati think that 666 years of rule is good enough! All the 6's adds up to a total of 21, which is half of 42 (2 x 42 is 84), which when divided by 3 (Satan's signature) adds up to 28, (7 x 4) and
```

### Expected behavior

Output is better when device = "cpu":

```
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "tiiuae/falcon-mamba-7b"
torch_dtype = torch.bfloat16

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype)
print(model)

tokenizer = AutoTokenizer.from_pretrained(model_id)

input_sequence = tokenizer.encode(
    "the largest crater on the moon is", return_tensors="pt")
print(input_sequence)

output_sequence = model.generate(input_sequence, max_new_tokens=100, do_sample=True)

print(tokenizer.decode(output_sequence[0]))
```

Output:

```
the largest crater on the moon is called Orientale; the largest on Mars is named Hellas, after the mythical Greek home of Zeus on earth).
It looks like cratering occurred, as on Earth, from a single “shot,” but on a much grander scale. If Earth were as large as Mars and got struck by the same number of asteroids, the Earth’s surface would look as battered as the moon’s. But Earth’s greater gravity and mass (the Earth is 4.5
```

Comments:
Hi @derekelewis - thanks for flagging this issue! 

What mps device e.g. M1, are you running on? 

We've previously observed a difference in behaviour of layers between mps and other devices -- #22468 -- and sometimes unexplained behaviour possibly to do with the torch version -- #32005. 

Similar [to this comment](https://github.com/huggingface/transformers/issues/32005#issuecomment-2231926114) - could you run to get the hidden states from the model with a single forward pass in both cases and compare the outputs? For there we should be able to bisect to try and find where the differences might be creeping in and how large they are 

@amyeroberts this is on a M2 Max.

Using this code to generate the output of a single forward pass:

```
for layer in range(len(hidden_states_mps[0])):
    print(layer, (hidden_states_mps[0][layer].to("cpu") - hidden_states_cpu[0][layer]).abs().max())
```

Output:

```
0 tensor(0.5914)
1 tensor(0.6874)
2 tensor(0.5764)
3 tensor(0.6183)
4 tensor(0.7371)
5 tensor(0.8249)
6 tensor(0.9075)
7 tensor(1.0664)
8 tensor(1.2437)
9 tensor(1.4555)
10 tensor(1.8229)
11 tensor(1.8912)
12 tensor(2.0372)
13 tensor(2.2490)
14 tensor(2.5427)
15 tensor(3.0060)
16 tensor(3.5862)
17 tensor(3.9439)
18 tensor(4.4555)
19 tensor(4.6528)
20 tensor(5.4690)
21 tensor(5.7601)
22 tensor(5.9815)
23 tensor(6.4383)
24 tensor(6.4707)
25 tensor(6.9566)
26 tensor(7.8468)
27 tensor(8.0318)
28 tensor(9.1572)
29 tensor(10.1214)
30 tensor(10.8080)
31 tensor(10.9183)
32 tensor(11.5801)
33 tensor(12.9529)
34 tensor(13.2965)
35 tensor(14.2415)
36 tensor(15.0923)
37 tensor(16.4985)
38 tensor(17.8206)
39 tensor(19.8181)
40 tensor(21.8513)
41 tensor(23.6099)
42 tensor(25.4560)
43 tensor(28.4208)
44 tensor(30.4484)
45 tensor(33.2785)
46 tensor(35.9397)
47 tensor(36.8553)
48 tensor(42.5090)
49 tensor(47.4634)
50 tensor(51.9332)
51 tensor(59.9800)
52 tensor(62.4336)
53 tensor(65.5664)
54 tensor(66.7109)
55 tensor(68.0918)
56 tensor(69.4824)
57 tensor(73.3965)
58 tensor(74.1289)
59 tensor(77.2475)
60 tensor(94.1821)
61 tensor(99.2111)
62 tensor(123.3049)
63 tensor(119.9967)
64 tensor(357.7265)
```

By comparison here's a forward pass w/ CUDA on a T4:

```
0 tensor(0.0001)
1 tensor(0.0005)
2 tensor(0.0010)
3 tensor(0.0010)
4 tensor(0.0015)
5 tensor(0.0016)
6 tensor(0.0016)
7 tensor(0.0016)
8 tensor(0.0023)
9 tensor(0.0023)
10 tensor(0.0043)
11 tensor(0.0038)
12 tensor(0.0048)
13 tensor(0.0048)
14 tensor(0.0048)
15 tensor(0.0067)
16 tensor(0.0106)
17 tensor(0.0087)
18 tensor(0.0067)
19 tensor(0.0096)
20 tensor(0.0096)
21 tensor(0.0057)
22 tensor(0.0057)
23 tensor(0.0057)
24 tensor(0.0074)
25 tensor(0.0092)
26 tensor(0.0092)
27 tensor(0.0094)
28 tensor(0.0094)
29 tensor(0.0211)
30 tensor(0.0250)
31 tensor(0.0250)
32 tensor(0.0250)
33 tensor(0.0172)
34 tensor(0.0133)
35 tensor(0.0153)
36 tensor(0.0190)
37 tensor(0.0209)
38 tensor(0.0248)
39 tensor(0.0268)
40 tensor(0.0316)
41 tensor(0.0355)
42 tensor(0.0433)
43 tensor(0.0394)
44 tensor(0.0551)
45 tensor(0.0490)
46 tensor(0.0660)
47 tensor(0.0709)
48 tensor(0.1168)
49 tensor(0.0865)
50 tensor(0.0865)
51 tensor(0.0985)
52 tensor(0.0957)
53 tensor(0.0957)
54 tensor(0.0957)
55 tensor(0.1102)
56 tensor(0.1102)
57 tensor(0.1248)
58 tensor(0.1560)
59 tensor(0.1873)
60 tensor(0.3123)
61 tensor(0.2285)
62 tensor(0.3066)
63 tensor(0.5599)
64 tensor(0.6773)
```

@derekelewis Thanks for sharing! 

OK, so there's already large differences happening from the first `FalconMambaBlock`. In this case, the next step is to inspect the activations at each step in `FalconMambaBlock` layer to try and isolate the offending operation

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.