Possible data converting problem when using flash attention 2 with whisper
Author: changyeli
Pull request: no

### System Info

- `transformers` version: 4.35.0
- Platform: Linux-5.14.0-284.25.1.el9_2.x86_64-x86_64-with-glibc2.34
- Python version: 3.11.5
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: trainer's default

### Who can help?

@sanchit-gandhi


### Information

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

### Tasks

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

### Reproduction

I was trying to fine-tune whisper small with flash attention 2 on a private data. Followed the [post](https://huggingface.co/blog/fine-tune-whisper) here for most of the code. Here are some changes I made:

```python
model_card = "openai/whisper-small"
model_name = model_card.split("/")[-1]
config = configparser.ConfigParser()
config.read("config.ini")
tran_df = pd.read_csv("../total_df.csv")
processor = AutoProcessor.from_pretrained(
    model_card)
tokenizer = WhisperTokenizer.from_pretrained(
    model_card)
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_card)
temo_dt = load_dataset(
    "audiofolder", data_dir=config['DATA']['dataset'],
    split="train[:1%]")
temo_dt = temo_dt.train_test_split(test_size=0.3)
temo_dt = temo_dt.cast_column("audio", Audio(sampling_rate=16000))

model = WhisperForConditionalGeneration.from_pretrained(
    model_card, use_flash_attention_2=True,
    torch_dtype=torch.float16)
model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(
    language="english", task="transcribe")
model.config.suppress_tokens = []
data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
# training process
training_args = Seq2SeqTrainingArguments(
    output_dir=f"../{model_name}", 
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    learning_rate=1e-5,
    warmup_steps=500,
    max_steps=6000,
    # speed up
    gradient_checkpointing=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=16,
    predict_with_generate=True,
    generation_max_length=225,
    save_steps=1000,
    eval_steps=1000,
    logging_steps=25,
    report_to="none",
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    auto_find_batch_size=True,
    torch_compile=True,
)
trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=temo_dt["train"],
    eval_dataset=temo_dt["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics_wer,
    tokenizer=processor.feature_extractor,
)
trainer.train()
```

It gave me this error: `RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same.`

So I tried to convert the `temo_dt` to half tensor using the following code:

```python
format = {'type': 'torch', 'format_kwargs' :{'dtype': torch.float16}}
temo_dt.set_format(**format)
```

But it returned this error: `RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding).`

Very interestingly, I can fine-tune the whisper small model perfectly without flash attention 2 using the code above. Is there anything I missed?

### Expected behavior

Fine-tuning whisper should go as expected with `use_flash_attention_2=True`.

Comments:
cc @younesbelkada 

hi @changyeli 
Can you share the full traceback of the two errors you are getting?

Can you also share the content of the dataset? note only the `input_features` needs to be casted in float16 and not the entire entries on your dataset. Perhaps you can use `dataset.map(xxx, batched=True)` to only cast `input_features` in float16 : https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py#L1684

@younesbelkada Sure, here is the full traceback using `torch_dtype=torch.float16`:

```shell
Traceback (most recent call last):
  File "/home/suppl/scripts/fine_tune_whisper.py", line 161, in <module>
    trainer.train()
  File "/home//anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home//anaconda3/envs/whisper/lib/python3.11/site-packages/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home//anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1860, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2725, in training_step
    loss = self.compute_loss(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2748, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 185, in forward
    outputs = self.parallel_apply(replicas, inputs, module_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 200, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 110, in parallel_apply
    output.reraise()
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in _worker
    output = module(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1683, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1543, in forward
    encoder_outputs = self.encoder(
                      ^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1119, in forward
    inputs_embeds = nn.functional.gelu(self.conv1(input_features))
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 310, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 306, in _conv_forward
    return F.conv1d(input, weight, bias, self.stride,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same
```

I also have a somehow related question regarding the gradient checkpoint. I got the following error when I set `gradient_checkpointing=True` on a small subset.

```
Segmentation fault (core dumped)
```

It looks like something related to memory. I used the whisper small model and ~500 samples for fine-tuning and got this error. Interestingly, It was perfectly fine with inference - I can inference using whisper large on a much bigger dataset. How should I proceed?

The input features are indeed in `float32` dtype. Tried to cast using the following code, but didn't work:

```python
format = {'type': 'torch', 'format_kwargs' :{'dtype': torch.float16}}
temo_dt['train'].set_format(columns=['input_features'], **format)
```

Could you try removing:
```diff
model = WhisperForConditionalGeneration.from_pretrained(
    model_card, use_flash_attention_2=True,
-   torch_dtype=torch.float16)
```

And then setting:
```diff
training_args = Seq2SeqTrainingArguments(
    output_dir=f"../{model_name}", 
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    learning_rate=1e-5,
    warmup_steps=500,
    max_steps=6000,
    # speed up
    gradient_checkpointing=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=16,
    predict_with_generate=True,
    generation_max_length=225,
    save_steps=1000,
    eval_steps=1000,
    logging_steps=25,
    report_to="none",
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    auto_find_batch_size=True,
    torch_compile=True,
+   fp16=True
)
```
=> we should cast the `input_features` to the correct `dtype` if we let the Trainer handle it.

Hi @sanchit-gandhi, thanks for the suggestion! But I still got a data converting related error with your approach - here is full trace of error message:

```
ializing it on CPU with `model.to('cuda')`.
  0%|                                                                                               | 0/870 [00:00<?, ?it/s]The input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in torch.float32.
Traceback (most recent call last):
  File "/home/coraal-suppl/scripts/fine_tune_whisper.py", line 164, in <module>
    trainer.train()
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1860, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2725, in training_step
    loss = self.compute_loss(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2748, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 185, in forward
    outputs = self.parallel_apply(replicas, inputs, module_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 200, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 110, in parallel_apply
    output.reraise()
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in _worker
    output = module(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1683, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1543, in forward
    encoder_outputs = self.encoder(
                      ^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1159, in forward
    layer_outputs = encoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 722, in forward
    hidden_states, attn_weights, _ = self.self_attn(
                                     ^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 569, in forward
    attn_output = self._flash_attention_forward(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 629, in _flash_attention_forward
    attn_output = flash_attn_func(
                  ^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/flash_attn/flash_attn_interface.py", line 708, in flash_attn_func
    return FlashAttnFunc.apply(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/flash_attn/flash_attn_interface.py", line 437, in forward
    out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
                                                                ^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/flash_attn/flash_attn_interface.py", line 49, in _flash_attn_forward
    out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = flash_attn_cuda.fwd(
                                                                ^^^^^^^^^^^^^^^^^^^^
RuntimeError: FlashAttention only support fp16 and bf16 data type

  0%|          | 0/870 [00:09<?, ?it/s]
```

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.

Hi @sanchit-gandhi I keep getting similar errors after upgrading to the most recent version (as of 01/13/24). Loaded the model using 

```python
model = WhisperForConditionalGeneration.from_pretrained(
        model_card,
        attn_implementation="flash_attention_2",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        use_safetensors=True)
```
As the `torch_dtype` here is required, or it will raise the following error even if I set `fp16=True` in the `Seq2SeqTrainingArguments`: ValueError: Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes. You passed torch.float32, this might lead to unexpected behaviour.

I got this error when trying to fine-tune whisper model: RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same. This is the same error I got when I open this issue. I also converted the the input data using 

```python
dt.set_format(
        columns=['input_features', 'labels'],
        **format
    )
```
And call the model via:

```python
model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_card,
        attn_implementation="flash_attention_2",
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True)
```

But I got a different error message:

 trainer.train()
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1537, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 1854, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2735, in training_step
    loss = self.compute_loss(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/trainer.py", line 2758, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 185, in forward
    outputs = self.parallel_apply(replicas, inputs, module_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py", line 200, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 110, in parallel_apply
    output.reraise()
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in _worker
    output = module(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1818, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1694, in forward
    decoder_outputs = self.decoder(
                      ^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/transformers/models/whisper/modeling_whisper.py", line 1442, in forward
    inputs_embeds = self.embed_tokens(input_ids)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/modules/sparse.py", line 162, in forward
    return F.embedding(
           ^^^^^^^^^^^^
  File "/home/anaconda3/envs/whisper/lib/python3.11/site-packages/torch/nn/functional.py", line 2233, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)

  0%|          | 0/22 [00:08<?, ?it/s]

It's interesting as I use FA2 in decoder-only language models and it works fine, but no luck on whisper yet. Any suggestion on how to proceed?

Did several experiments - turned out I forgot to drop unused columns for preprocessing but the problem still exists.

For fine-tuning, if I remove `torch_dtype=torch.float16` when loading the pre-trained model and enable `fp16=True` in the training arguments. I got the following error:

ValueError: Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes. You passed torch.float32, this might lead to unexpected behaviour.

@sanchit-gandhi any update on this? How do we use flash attention 2 here? 

My guess is that this has to do with the mixed precision training not working fine. In `accelerate`, by passing mixed_precision as "no", this error wasn't coming up but how do we do the same in `Trainer`? 

I am encountering the same kind of issue. Has any solution been found yet? I only get the error on the evaluation step of the trainer, the training itself seems to work fine. 

`RuntimeError: Input type (float) and bias type (c10::BFloat16) should be the same`

cc @sanchit-gandhi @ylacombe 

Hey @changyeli @cvl01 and  @RohitMidha23, is it still an issue?

Hey @ylacombe Somehow I bypassed this issue with #32370 and #32366 

It works with `attn_implementation="sdpa"` when retrieving model's pre-trained weights.

@changyeli that sort of goes against the point as you aren't using `flash` attention anymore - which is inherently designed to be faster... 

@RohitMidha23 I think sdpa supports flash attention. This is what I found in the [documentation](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention):

> PyTorch’s [torch.nn.functional.scaled_dot_product_attention](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html) (SDPA) can also call FlashAttention and memory-efficient attention kernels under the hood. SDPA support is currently being added natively in Transformers and is used by default for torch>=2.1.1 when an implementation is available. You may also set attn_implementation="sdpa" in from_pretrained() to explicitly request SDPA to be used.

But I agree regarding the speed - it does not speed up _that_ much using SDPA.

> Hey @changyeli @cvl01 and @RohitMidha23, is it still an issue?嘿，还有问题吗？

Yes I still meet same question on transformers 3.37.1 (latest from now)


```python
from transformers import AutoTokenizer,Trainer,TrainingArguments,AutoModelForCausalLM
from argparse import ArgumentParser
from template import modelType2Template
import datasets
from functools import partial
from data_utils.load_func import dname2load
from data_utils.dataset_func import dname2func
from torch.utils.data import DataLoader
from config import *
from collator import DataCollatorForContinuousRegionMask
import os
def parse_args():
    parser=ArgumentParser()
    parser.add_argument("--model_name", type=str, default="qwen2.5-0.5b")
    parser.add_argument("--dataset", type=str, default="fineweb_1b")
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--max_length", type=int, default=512)
    parser.add_argument("--template",default='qwen2.5')
    parser.add_argument("--gradient_accumulation_steps",default=32)
    parser.add_argument("--num_train_epochs",default=3)
    return parser.parse_args()



args=parse_args()


tokenizer=AutoTokenizer.from_pretrained(model_dir.get(args.model_name, None))
tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
if tokenizer.mask_token==None:
    special_tokens_dict = {"mask_token": "[MASK]"}
    tokenizer.add_special_tokens(special_tokens_dict)
template = modelType2Template[args.template](tokenizer)
# ignore my dataset process, I check this by adding collator and dataset to torch dataset and print out dtype of input_ids, return torch.int64
dataset_name_list = args.dataset.split(",")
dataset_list = []
for dname in dataset_name_list:
    train_dataset = dname2load[dname](dataset_dir.get(dname, None))
    train_dataset = train_dataset.map(
        partial(
            dname2func[dname],
            template=template,
            mode=0,
            test=False,
            pt=True,
        ),
        batched=True,
        num_proc=os.cpu_count(),
        remove_columns=train_dataset.features.keys(),
        desc="tokenize",
        load_from_cache_file=False,
        
    )
    dataset_list.append(train_dataset)
train_dataset = datasets.concatenate_datasets(dataset_list)
collator=DataCollatorForContinuousRegionMask(tokenizer)

training_args=TrainingArguments(
    seed=42,
    data_seed=42,
    overwrite_output_dir=True,
    output_dir='models',
    logging_steps=1,
    remove_unused_columns=True,
    gradient_accumulation_steps=args.gradient_accumulation_steps,
    save_strategy="epoch",
    dataloader_pin_memory=True,
    dataloader_prefetch_factor=3,
    dataloader_num_workers=16,
    num_train_epochs=args.num_train_epochs,
    per_device_train_batch_size=1,
    bf16=True,
    warmup_ratio=0.05,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    include_num_input_tokens_seen =True,
)

model = AutoModelForCausalLM.from_pretrained(model_dir.get(args.model_name, None))
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=collator,
    train_dataset =train_dataset,
)

trainer.train()
```


It’s puzzling that errors only started to appear after dozens of steps of training.
```bash
17:40:08.507   0%|          | 0/32307 [00:00<?, ?it/s]
17:40:08.510   0%|          | 1/32307 [00:15<142:05:43, 15.83s/it]
17:40:08.510                                                      
17:40:08.511 
17:40:18.421   0%|          | 1/32307 [00:15<142:05:43, 15.83s/it]
17:40:18.423   0%|          | 2/32307 [00:25<110:50:13, 12.35s/it]
17:40:18.423                                                      
17:40:18.423 
17:40:29.065   0%|          | 2/32307 [00:25<110:50:13, 12.35s/it]
17:40:29.070   0%|          | 3/32307 [00:36<103:50:21, 11.57s/it]
17:40:29.070                                                      
17:40:29.070 
17:40:38.748   0%|          | 3/32307 [00:36<103:50:21, 11.57s/it]
17:40:38.751   0%|          | 4/32307 [00:46<97:08:38, 10.83s/it] 
17:40:38.751                                                     
17:40:38.751 
17:40:49.651   0%|          | 4/32307 [00:46<97:08:38, 10.83s/it]
17:40:49.655   0%|          | 5/32307 [00:56<97:23:26, 10.85s/it]
17:40:49.655                                                     
17:40:49.655 
17:40:59.642   0%|          | 5/32307 [00:56<97:23:26, 10.85s/it]
17:40:59.644   0%|          | 6/32307 [01:06<94:45:09, 10.56s/it]
17:40:59.644                                                     
17:40:59.644 
17:41:09.615   0%|          | 6/32307 [01:06<94:45:09, 10.56s/it]
17:41:09.618   0%|          | 7/32307 [01:16<93:01:46, 10.37s/it]
17:41:09.618                                                     
17:41:09.618 
17:41:19.581   0%|          | 7/32307 [01:16<93:01:46, 10.37s/it]
17:41:19.584   0%|          | 8/32307 [01:26<91:52:27, 10.24s/it]
17:41:19.584                                                     
17:41:19.584 
17:41:29.323   0%|          | 8/32307 [01:26<91:52:27, 10.24s/it]
17:41:29.326   0%|          | 9/32307 [01:36<90:28:34, 10.08s/it]
17:41:29.326                                                     
17:41:29.326 
17:41:39.665   0%|          | 9/32307 [01:36<90:28:34, 10.08s/it]
17:41:39.668   0%|          | 10/32307 [01:46<91:11:02, 10.16s/it]
17:41:39.668                                                      
17:41:39.668 
17:41:50.166   0%|          | 10/32307 [01:46<91:11:02, 10.16s/it]
17:41:50.169   0%|          | 11/32307 [01:57<92:06:26, 10.27s/it]
17:41:50.169                                                      
17:41:50.169 
17:41:59.787   0%|          | 11/32307 [01:57<92:06:26, 10.27s/it]
17:41:59.790   0%|          | 12/32307 [02:07<90:20:31, 10.07s/it]
17:41:59.790                                                      
17:41:59.790 
17:42:09.774   0%|          | 12/32307 [02:07<90:20:31, 10.07s/it]
17:42:09.777   0%|          | 13/32307 [02:17<90:06:43, 10.05s/it]
17:42:09.777                                                      
17:42:09.777 
17:42:19.594   0%|          | 13/32307 [02:17<90:06:43, 10.05s/it]
17:42:19.596   0%|          | 14/32307 [02:26<89:29:50,  9.98s/it]
17:42:19.596                                                      
17:42:19.596 
17:42:29.328   0%|          | 14/32307 [02:26<89:29:50,  9.98s/it]
17:42:29.331   0%|          | 15/32307 [02:36<88:50:22,  9.90s/it]
17:42:29.331                                                      
17:42:29.331 
17:42:39.040   0%|          | 15/32307 [02:36<88:50:22,  9.90s/it]
17:42:39.043   0%|          | 16/32307 [02:46<88:19:02,  9.85s/it]
17:42:39.043                                                      
17:42:39.043 
17:42:48.525   0%|          | 16/32307 [02:46<88:19:02,  9.85s/it]
17:42:48.527   0%|          | 17/32307 [02:55<87:20:20,  9.74s/it]
17:42:48.527                                                      
17:42:48.527 
17:42:58.152   0%|          | 17/32307 [02:55<87:20:20,  9.74s/it]
17:42:58.155   0%|          | 18/32307 [03:05<87:02:22,  9.70s/it]
17:42:58.155                                                      
17:42:58.155 
17:43:07.884   0%|          | 18/32307 [03:05<87:02:22,  9.70s/it]
17:43:07.887   0%|          | 19/32307 [03:15<87:06:48,  9.71s/it]
17:43:07.887                                                      
17:43:07.887 
17:43:18.860   0%|          | 19/32307 [03:15<87:06:48,  9.71s/it]
17:43:18.862   0%|          | 20/32307 [03:26<90:30:34, 10.09s/it]
17:43:18.862                                                      
17:43:18.862 
17:43:29.165   0%|          | 20/32307 [03:26<90:30:34, 10.09s/it]
17:43:29.167   0%|          | 21/32307 [03:36<91:04:51, 10.16s/it]
17:43:29.167                                                      
17:43:29.167 
17:43:39.218   0%|          | 21/32307 [03:36<91:04:51, 10.16s/it]
17:43:39.219   0%|          | 22/32307 [03:46<90:48:03, 10.12s/it]
17:43:39.219                                                      
17:43:39.219 
17:43:49.079   0%|          | 22/32307 [03:46<90:48:03, 10.12s/it]
17:43:49.080   0%|          | 23/32307 [03:56<90:05:16, 10.05s/it]
17:43:49.080                                                      
17:43:49.080 
17:43:59.179   0%|          | 23/32307 [03:56<90:05:16, 10.05s/it]
17:43:59.181   0%|          | 24/32307 [04:06<90:13:59, 10.06s/it]
17:43:59.181                                                      
17:43:59.181 
17:44:09.691   0%|          | 24/32307 [04:06<90:13:59, 10.06s/it]
17:44:09.692   0%|          | 25/32307 [04:17<91:26:18, 10.20s/it]
17:44:09.692                                                      
17:44:09.692 
17:44:19.358   0%|          | 25/32307 [04:17<91:26:18, 10.20s/it]
17:44:19.360   0%|          | 26/32307 [04:26<90:00:41, 10.04s/it]
17:44:19.360                                                      
17:44:19.360 
17:44:29.024   0%|          | 26/32307 [04:26<90:00:41, 10.04s/it]
17:44:29.026   0%|          | 27/32307 [04:36<89:00:25,  9.93s/it]
17:44:29.026                                                      
17:44:29.026 
17:44:39.154   0%|          | 27/32307 [04:36<89:00:25,  9.93s/it]
17:44:39.156   0%|          | 28/32307 [04:46<89:33:09,  9.99s/it]
17:44:39.156                                                      
17:44:39.156 
17:44:49.396   0%|          | 28/32307 [04:46<89:33:09,  9.99s/it]
17:44:49.398   0%|          | 29/32307 [04:56<90:14:02, 10.06s/it]
17:44:49.398                                                      
17:44:49.398 
17:44:59.303   0%|          | 29/32307 [04:56<90:14:02, 10.06s/it]
17:44:59.305   0%|          | 30/32307 [05:06<89:48:32, 10.02s/it]
17:44:59.305                                                      
17:44:59.305 
17:45:09.057   0%|          | 30/32307 [05:06<89:48:32, 10.02s/it]
17:45:09.059   0%|          | 31/32307 [05:16<89:05:59,  9.94s/it]
17:45:09.059                                                      
17:45:09.059 
17:45:19.239   0%|          | 31/32307 [05:16<89:05:59,  9.94s/it]
17:45:19.241   0%|          | 32/32307 [05:26<89:45:11, 10.01s/it]
17:45:19.241                                                      
17:45:19.241 
17:45:29.519   0%|          | 32/32307 [05:26<89:45:11, 10.01s/it]
17:45:29.521   0%|          | 33/32307 [05:36<90:28:17, 10.09s/it]
17:45:29.521                                                      
17:45:29.521 
17:45:39.569   0%|          | 33/32307 [05:36<90:28:17, 10.09s/it]
17:45:39.570   0%|          | 34/32307 [05:46<90:21:27, 10.08s/it]
17:45:39.570                                                      
17:45:39.571 
17:45:49.453   0%|          | 34/32307 [05:46<90:21:27, 10.08s/it]
17:45:49.455   0%|          | 35/32307 [05:56<89:49:51, 10.02s/it]
17:45:49.455                                                      
17:45:49.455 
17:45:59.181   0%|          | 35/32307 [05:56<89:49:51, 10.02s/it]
17:45:59.184   0%|          | 36/32307 [06:06<89:02:27,  9.93s/it]
17:45:59.184                                                      
17:45:59.184 
17:46:09.525   0%|          | 36/32307 [06:06<89:02:27,  9.93s/it]
17:46:09.527   0%|          | 37/32307 [06:16<90:08:32, 10.06s/it]
17:46:09.527                                                      
17:46:09.527 
17:46:19.639   0%|          | 37/32307 [06:16<90:08:32, 10.06s/it]
17:46:19.641   0%|          | 38/32307 [06:26<90:17:44, 10.07s/it]
17:46:19.641                                                      
17:46:19.641 
17:46:29.725   0%|          | 38/32307 [06:26<90:17:44, 10.07s/it]
17:46:29.726   0%|          | 39/32307 [06:37<90:19:28, 10.08s/it]
17:46:29.726                                                      
17:46:29.726 
17:46:39.962   0%|          | 39/32307 [06:37<90:19:28, 10.08s/it]
17:46:39.964   0%|          | 40/32307 [06:47<90:45:07, 10.13s/it]
17:46:39.964                                                      
17:46:39.964 
17:46:49.979   0%|          | 40/32307 [06:47<90:45:07, 10.13s/it]
17:46:49.981   0%|          | 41/32307 [06:57<90:27:28, 10.09s/it]
17:46:49.981                                                      
17:46:49.981 
17:47:00.064   0%|          | 41/32307 [06:57<90:27:28, 10.09s/it]
17:47:00.067   0%|          | 42/32307 [07:07<90:26:11, 10.09s/it]
17:47:00.067                                                      
17:47:00.067 
17:47:10.359   0%|          | 42/32307 [07:07<90:26:11, 10.09s/it]
17:47:10.361   0%|          | 43/32307 [07:17<90:58:54, 10.15s/it]
17:47:10.361                                                      
17:47:10.361 
17:47:20.176   0%|          | 43/32307 [07:17<90:58:54, 10.15s/it]
17:47:20.179   0%|          | 44/32307 [07:27<90:04:47, 10.05s/it]
17:47:20.179                                                      
17:47:20.179 
17:47:30.124   0%|          | 44/32307 [07:27<90:04:47, 10.05s/it]
17:47:30.126   0%|          | 45/32307 [07:37<89:47:53, 10.02s/it]
17:47:30.126                                                      
17:47:30.126 
17:47:40.010   0%|          | 45/32307 [07:37<89:47:53, 10.02s/it]
17:47:40.013   0%|          | 46/32307 [07:47<89:26:10,  9.98s/it]
17:47:40.013                                                      
17:47:40.013 
17:47:48.038   0%|          | 46/32307 [07:47<89:26:10,  9.98s/it][rank1]: Traceback (most recent call last):
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/INS/ruanjunhao04/ICP/middle_training/data.py", line 87, in <module>
17:47:48.038 [rank1]:     trainer.train()
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/trainer.py", line 2164, in train
17:47:48.038 [rank1]:     return inner_training_loop(
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/trainer.py", line 2522, in _inner_training_loop
17:47:48.038 [rank1]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
17:47:48.038 [rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/trainer.py", line 3653, in training_step
17:47:48.038 [rank1]:     loss = self.compute_loss(model, inputs)
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/trainer.py", line 3709, in compute_loss
17:47:48.038 [rank1]:     outputs = model(**inputs)
17:47:48.038 [rank1]:               ^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
17:47:48.038 [rank1]:     return self._call_impl(*args, **kwargs)
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
17:47:48.038 [rank1]:     return forward_call(*args, **kwargs)
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
17:47:48.038 [rank1]:     ret_val = func(*args, **kwargs)
17:47:48.038 [rank1]:               ^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 1899, in forward
17:47:48.038 [rank1]:     loss = self.module(*inputs, **kwargs)
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
17:47:48.038 [rank1]:     return self._call_impl(*args, **kwargs)
17:47:48.038 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
17:47:48.038 [rank1]:     return inner()
17:47:48.038 [rank1]:            ^^^^^^^
17:47:48.038 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
17:47:48.038 [rank1]:     result = forward_call(*args, **kwargs)
17:47:48.039 [rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 1165, in forward
17:47:48.039 [rank1]:     outputs = self.model(
17:47:48.039 [rank1]:               ^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
17:47:48.039 [rank1]:     return self._call_impl(*args, **kwargs)
17:47:48.039 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
17:47:48.039 [rank1]:     return inner()
17:47:48.039 [rank1]:            ^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
17:47:48.039 [rank1]:     result = forward_call(*args, **kwargs)
17:47:48.039 [rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 854, in forward
17:47:48.039 [rank1]:     inputs_embeds = self.embed_tokens(input_ids)
17:47:48.039 [rank1]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
17:47:48.039 [rank1]:     return self._call_impl(*args, **kwargs)
17:47:48.039 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
17:47:48.039 [rank1]:     return inner()
17:47:48.039 [rank1]:            ^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
17:47:48.039 [rank1]:     result = forward_call(*args, **kwargs)
17:47:48.039 [rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 190, in forward
17:47:48.039 [rank1]:     return F.embedding(
17:47:48.039 [rank1]:            ^^^^^^^^^^^^
17:47:48.039 [rank1]:   File "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-aipnlp/ruanjunhao04/env/rjh/lib/python3.12/site-packages/torch/nn/functional.py", line 2551, in embedding
17:47:48.039 [rank1]:     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
17:47:48.039 [rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
17:47:48.039 [rank1]: RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got CUDABFloat16Type instead (while checking arguments for embedding)

```