Potential bug in xLSTM (native) training
Author: stefan-it
Pull request: no
Labels: name: bug

### System Info

Ha, probably it needs another bug report, but `transformers env` is showing me:

```
Traceback (most recent call last):
  File "/home/ubuntu/venvs/dev/bin/transformers", line 5, in <module>
    from transformers.cli.transformers import main
  File "/home/ubuntu/transformers/src/transformers/cli/transformers.py", line 23, in <module>
    from transformers.cli.serve import Serve
  File "/home/ubuntu/transformers/src/transformers/cli/serve.py", line 351, in <module>
    class Serve:
  File "/home/ubuntu/transformers/src/transformers/cli/serve.py", line 658, in Serve
    ) -> ChatCompletionChunk:
         ^^^^^^^^^^^^^^^^^^^
NameError: name 'ChatCompletionChunk' is not defined
```

However, the most relevant stuff:

* Transformers latest commit, 370fc65ee594b82aafcb971ce5a0d60ca0851954
* PyTorch 2.9
* CUDA 12.8
* H100 (it needs around 50GB of GPU RAM)

### Who can help?

@kpoeppel (for adding the original xLSTM implementation into Transformers)
@Cyrilvallez

### Information

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

### Tasks

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

### Reproduction

# xLSTM Training with **native** Transformers implementation

Install latest Transformers version.

Notice: for this section of my bug report it is very important, that the `xlstm` package is **not** installed! 

```bash
cd transformers/examples/pytorch/language-modeling
```

then execute:

```
python3 run_clm.py \
--dataset_name stefan-it/xlstm-transformers-bug-data \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 32 \
--do_train \
--output_dir ./xlstm-native \
--config_name stefan-it/xlstm-transformers-bug-native \
--tokenizer_name stefan-it/xlstm-transformers-bug-native \
--block_size 8192   --optim adamw_torch   --learning_rate 3e-03 \
--save_strategy steps   --save_steps 10000 \
--logging_strategy steps   --logging_steps 500   --bf16 \
--dataloader_drop_last yes  \
--preprocessing_num_workers 8 \
--max_steps 10700
```

The config in my demo repo contains:

* `"chunkwise_kernel": "chunkwise--native_autograd"`
* `"sequence_kernel": "native_sequence__native"`

You can see the full config in [my repo](https://huggingface.co/stefan-it/xlstm-transformers-bug-native)

Unfortunately, the training won't start because of:

```bash
  File "/home/ubuntu/transformers/src/transformers/utils/generic.py", line 757, in wrapper
    output = func(self, *args, **kwargs)                                          
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                            
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 1467, in forward
    hidden_states, rnn_state = xlstm_block(                                                                                                                         
                               ^^^^^^^^^^^^             
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl                                           
    return self._call_impl(*args, **kwargs)                                                                                                                         
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                         
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)                                                                                                                            
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                          
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 1192, in forward                                                           
    x_mlstm, state = self.mlstm_layer(x_mlstm, state)                                                                                                               
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                               
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)                                       
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                       
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl                                                   
    return forward_call(*args, **kwargs)                                                                                                                            
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                            
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 1139, in forward
    h, state = self.mlstm_backend(
               ^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/venvs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 815, in forward
    return self._train_fn(
           ^^^^^^^^^^^^^^^
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 358, in mlstm_chunkwise_native_autograd
    matH_out, vecN_out, vecM_out = mlstm_chunkwise_parallel_fw_H(
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/transformers/src/transformers/models/xlstm/modeling_xlstm.py", line 179, in mlstm_chunkwise_parallel_fw_H
    matQ = matQ.view(batch_size, nh, nc, chunk_size, dqk)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: shape '[4, 4, 128, 64, 8192]' is invalid for input of size 8388608
```

# xLSTM Training with `xlstm` package

In the next step, the `xlstm` package must be installed:

```bash
pip3 install xlstm
```

We use the same command for pretraining:

```bash
python3 run_clm.py \
--dataset_name stefan-it/xlstm-transformers-bug-data \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 32 \
--do_train \
--output_dir ./xlstm-integration \
--config_name stefan-it/xlstm-transformers-bug-native \
--tokenizer_name stefan-it/xlstm-transformers-bug-native \
--block_size 8192   --optim adamw_torch   --learning_rate 3e-03 \
--save_strategy steps   --save_steps 10000 \
--logging_strategy steps   --logging_steps 500   --bf16 \
--dataloader_drop_last yes  \
--preprocessing_num_workers 8 \
--max_steps 10700
```

It is now working! However, it is very slow.

# xLSTM Training with `xlstm` + Triton

To verify, if the xLSTM implementation with `xlstm` package and Triton is working, I have set-up another config repo.

Installation steps:

```bash
pip3 install xlstm
pipe install triton==3.2.0
```

(It only works with an older Triton version).

Then the pretraining command is:

```bash
python3 run_clm.py \
--dataset_name stefan-it/xlstm-transformers-bug-data \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 32 \
--do_train \
--output_dir ./xlstm-triton \
--config_name stefan-it/xlstm-transformers-bug-triton \
--tokenizer_name stefan-it/xlstm-transformers-bug-triton \
--block_size 8192   --optim adamw_torch   --learning_rate 3e-03 \
--save_strategy steps   --save_steps 10000 \
--logging_strategy steps   --logging_steps 500   --bf16 \
--dataloader_drop_last yes  \
--preprocessing_num_workers 8 \
--max_steps 10700
```

It has the following configuration modifications:

* `"chunkwise_kernel": "chunkwise--triton_xl_chunk"`
* `"sequence_kernel": "native_sequence__triton"`

Training is working, this time much more faster, because of the Triton kernels.

### Expected behavior

I would expect, that xLSTM training is working without the `xlstm` package, because the Transformers implementation has an extra check for that:

https://github.com/huggingface/transformers/blob/370fc65ee594b82aafcb971ce5a0d60ca0851954/src/transformers/models/xlstm/modeling_xlstm.py#L30-L36

Comments:
The `env` issue feels like some of your versions might be out of sync - try uninstalling and reinstalling transformers and seeing if other packages get upgraded!

As for the `xlstm` issue, yes, that shape error should probably not be thrown. However, this might be an issue with `run_clm`, since that script probably assumes that you're training a transformer-like model.

Can you try making a short script to reproduce the issue that just loads the model with `AutoModelForCausalLM` or something?

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.