Repository: pytorch
Issue: Runtime error when running huggingface pretrained model with torch-xla
Labels: triaged,module: xla

### 🐛 Describe the bug

Encountered `RuntimeError: isDifferentiableType(variable.scalar_type())` during the forward pass of the training script of a hugging face wav2vec2 conformer model. This error happens with a CPU device. 

Full trace of the error is 
```
Traceback (most recent call last):
  File "run_wav2vec2_pretraining_bug.py", line 190, in <module>
    model(**batch)
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 2124, in forward
    logits = self.compute_contrastive_logits(
  File "/home/ubuntu/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1997, in compute_contrastive_logits
    logits = torch.cosine_similarity(predicted_features.float(), target_features.float(), dim=-1)
RuntimeError: isDifferentiableType(variable.scalar_type()) INTERNAL ASSERT FAILED at "../torch/csrc/autograd/functions/utils.h":75, please report a bug to PyTorch. 
```

# Reproduction example
Please find the reproducing script in the following section. This script is modified from https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py, to include only the forward part of the training loop. 

To run the script, invoke with the following command in a venv satisfying the requirement:
```
PJRT_DEVICE=CPU python run_wav2vec2_pretraining_bug.py
```
Or similarly
```
PJRT_DEVICE=CPU python  run_wav2vec2_pretraining_no_trainer.py --dataset_name="librispeech_asr" --dataset_config_names clean clean --dataset_split_names validation test --model_name_or_path="facebook/wav2vec2-conformer-rope-large-960h-ft" --output_dir="./wav2vec2-pretrained-demo" --max_train_steps="1" --num_warmup_steps="1" --gradient_accumulation_steps="1" --learning_rate="0.005" --weight_decay="0.01" --max_duration_in_seconds="20.0" --min_duration_in_seconds="2.0" --logging_steps="1" --saving_steps="1" --per_device_train_batch_size="1" --per_device_eval_batch_size="1" --adam_beta1="0.9" --adam_beta2="0.98" --adam_epsilon="1e-06"  --mask_time_prob="0.65" --mask_time_length="10"
```

```
#  run_wav2vec2_pretraining_bug.py
# # Copyright 2021 The HuggingFace Inc. team. All rights reserved.
# #
# # Licensed under the Apache License, Version 2.0 (the "License");
# # you may not use this file except in compliance with the License.
# # You may obtain a copy of the License at
# #
# #     http://www.apache.org/licenses/LICENSE-2.0
# #
# # Unless required by applicable law or agreed to in writing, software
# # distributed under the License is distributed on an "AS IS" BASIS,
# # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# # See the License for the specific language governing permissions and

# """Pre-Training a 🤗 Wav2Vec2 model on unlabeled audio data"""

from dataclasses import dataclass
from typing import Dict, List, Optional, Union
import torch
from datasets import DatasetDict, concatenate_datasets, load_dataset
from torch.utils.data.dataloader import DataLoader
import torch
import datasets

from transformers import (
    Wav2Vec2Config,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForPreTraining,
)
from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices, _sample_negative_indices

@dataclass
class DataCollatorForWav2Vec2Pretraining:
    """
    Data collator that will dynamically pad the inputs received and prepare masked indices
    for self-supervised pretraining.

    Args:
        model (:class:`~transformers.Wav2Vec2ForPreTraining`):
            The Wav2Vec2 model used for pretraining. The data collator needs to have access
            to config and ``_get_feat_extract_output_lengths`` function for correct padding.
        feature_extractor (:class:`~transformers.Wav2Vec2FeatureExtractor`):
            The processor used for proccessing the data.
        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:
            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
              sequence if provided).
            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
              maximum acceptable input length for the model if that argument is not provided.
            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
              different lengths).
        max_length (:obj:`int`, `optional`):
            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).
        pad_to_multiple_of (:obj:`int`, `optional`):
            If set will pad the sequence to a multiple of the provided value.
            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
        mask_time_prob (:obj:`float`, `optional`, defaults to :obj:`0.65`):
            Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked for the contrastive task.
            Note that overlap between masked sequences may decrease the actual percentage of masked vectors.
            The default value is taken from the original wav2vec 2.0 article (https://arxiv.org/abs/2006.11477),
            and results in about 49 percent of each sequence being masked on average.
        mask_time_length (:obj:`int`, `optional`, defaults to :obj:`10`):
            Length of each vector mask span to mask along the time axis in the contrastive task. The default value
            originates from the original wav2vec 2.0 article and corresponds to the ``M`` variable mentioned there.
    """

    model: Wav2Vec2ForPreTraining
    feature_extractor: Wav2Vec2FeatureExtractor
    padding: Union[bool, str] = "longest"
    pad_to_multiple_of: Optional[int] = None
    mask_time_prob: Optional[float] = 0.65
    mask_time_length: Optional[int] = 10

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # reformat list to dict and set to pytorch format
        batch = self.feature_extractor.pad(
            features,
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors="pt",
        )

        device = batch["input_values"].device
        batch_size = batch["input_values"].shape[0]

        mask_indices_seq_length = self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1])
        # make sure masked sequence length is a Python scalar
        mask_indices_seq_length = int(mask_indices_seq_length)

        # make sure that no loss is computed on padded inputs
        if batch.get("attention_mask") is not None:
            # compute real output lengths according to convolution formula
            batch["sub_attention_mask"] = self.model._get_feature_vector_attention_mask(
                mask_indices_seq_length, batch["attention_mask"]
            )

        features_shape = (batch_size, mask_indices_seq_length)

        # sample randomly masked indices
        mask_time_indices = _compute_mask_indices(
            features_shape,
            self.mask_time_prob,
            self.mask_time_length,
            attention_mask=batch.get("sub_attention_mask"),
        )

        # sample negative indices
        sampled_negative_indices = _sample_negative_indices(
            features_shape,
            self.model.config.num_negatives,
            mask_time_indices=mask_time_indices,
        )
        batch["mask_time_indices"] = torch.tensor(mask_time_indices, dtype=torch.long, device=device)
        batch["sampled_negative_indices"] = torch.tensor(sampled_negative_indices, dtype=torch.long, device=device)

        return batch

model_id = "facebook/wav2vec2-conformer-rope-large-960h-ft"

# retreive model, feature extractor, and instantiate data collator
config = Wav2Vec2Config.from_pretrained(model_id)
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_id)
model = Wav2Vec2ForPreTraining(config)
data_collator = DataCollatorForWav2Vec2Pretraining(
    model=model,
    feature_extractor=feature_extractor,
    mask_time_prob=config.mask_time_prob,
    mask_time_length=config.mask_time_length,
)

raw_datasets = DatasetDict()
dataset_config_name = "clean"
train_split_name = "test"
dataset_split = load_dataset(
    "librispeech_asr",
    dataset_config_name,
    split=train_split_name,
    trust_remote_code=True,
)
raw_datasets["train"] = dataset_split

# make sure that dataset decodes audio with correct sampling rate
raw_datasets = raw_datasets.cast_column(
    "audio", datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
)

max_length = int(5.0 * feature_extractor.sampling_rate)
def prepare_dataset(batch):
    sample = batch["audio"]
    inputs = feature_extractor(
        sample["array"], sampling_rate=sample["sampling_rate"], max_length=max_length, truncation=True
    )
    batch["input_values"] = inputs.input_values[0]
    return batch

# load audio files into numpy arrays
vectorized_datasets = raw_datasets.map(
    prepare_dataset,
    remove_columns=raw_datasets["train"].column_names,
    cache_file_names=None,
)

train_dataloader = DataLoader(
    vectorized_datasets["train"],
    collate_fn=data_collator,
    batch_size=1,
)

dev = torch.device("xla")
model = model.to(dev)

for step, batch in enumerate(train_dataloader):
    # compute num of losses
    num_losses = batch["mask_time_indices"].sum()
    sub_attention_mask = batch.pop("sub_attention_mask", None)
    sub_attention_mask = (
        sub_attention_mask if sub_attention_mask is not None else torch.ones_like(batch["mask_time_indices"])
    )
    sub_attention_mask = sub_attention_mask.to(dev)
    num_losses = num_losses.to(dev)
    batch["input_values"] = batch["input_values"].to(dev)
    batch["attention_mask"] = batch["attention_mask"].to(dev)
    batch["mask_time_indices"] = batch["mask_time_indices"].to(dev)
    batch["sampled_negative_indices"] = batch["sampled_negative_indices"].to(dev)

    percent_masked = num_losses / sub_attention_mask.sum()
    model(**batch)
```


### Versions

The result include:
```
[pip3] numpy==1.24.4
[pip3] pytorchhighlevelperfmodel==1.0.0.1806+beb242c34
[pip3] torch==2.1.2
[pip3] torch-xla==2.1.3
[pip3] torchvision==0.16.2
[pip3] triton==2.1.0
```
Additional packages include: 
```
librosa                       0.10.2.post1
transformers                  4.41.2
datasets                      2.20.0
```

cc @bdhirsh