Repository: huggingface/transformers
Pull request: [Modeling] Fix encoder CPU offloading for whisper

## Purpose ##
* Fix CPU offloading for WhisperEncoder
  * This supports downstream applications which offload the layers of even small models like whisper. See https://github.com/vllm-project/llm-compressor/pull/1263

Without this change, attempting to CPU offload the encoder layer raises a device error
```
RuntimeError: Tensor on device meta is not on the expected device cuda:0!
```

## Changes ##
* Instead of getting the `embed_positions.weight` attribute directly, leverage the hf hooks attached to the `embed_positions` module to onload the weight properly.
  * This induces a small, once per request runtime cost as `F.embedding` must be called with an identity matrix, rather than grabbing the weight value directly

## Testing ##
Use the following test script to verify that generation works with the device map
```python3
device_map={
    "model.encoder": "cpu",
    "model.decoder": 0,
    "proj_out": 0,
},
```
<details><summary>test_whisper_offload.py</summary>

```python3
import torch
from datasets import load_dataset
from transformers import WhisperForConditionalGeneration, WhisperProcessor


def load_sample(processor):
    ds = load_dataset(
        "MLCommons/peoples_speech",
        "test", split="test[:1]",
        trust_remote_code=True,
    )

    sample = next(iter(ds))
    sample = processor(
        audio=sample["audio"]["array"],
        sampling_rate=sample["audio"]["sampling_rate"],
        text=(" " + sample["text"].capitalize()),
        add_special_tokens=True,
        return_tensors="pt",
    )

    sample["input_features"] = sample["input_features"].to(dtype=torch.bfloat16)
    sample["decoder_input_ids"] = torch.tensor([processor.tokenizer.prefix_tokens])
    del sample["labels"]

    return sample


if __name__ == "__main__":
    model_id = "openai/whisper-large-v3"
    model = WhisperForConditionalGeneration.from_pretrained(
        model_id,
        device_map={
            "model.encoder": "cpu",
            "model.decoder": 0,
            "proj_out": 0,
        },
        torch_dtype=torch.bfloat16
    )
    processor = WhisperProcessor.from_pretrained(model_id)

    assert model.model.encoder.embed_positions.weight.device == torch.device("meta")
    sample = load_sample(processor)
    output = model.generate(**sample, language="en")
    print(processor.batch_decode(output, skip_special_tokens=True))
```
</details>

## Potential Reviewers ##
* @SunMarc @ArthurZucker @zucchini-nlp

Changed files:
filename: src/transformers/models/whisper/modeling_whisper.py
patch: @@ -687,9 +687,9 @@ def forward(
         inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
 
         inputs_embeds = inputs_embeds.permute(0, 2, 1)
-        embed_pos = self.embed_positions.weight
+        all_positions = torch.arange(self.embed_positions.num_embeddings, device=inputs_embeds.device)
 
-        hidden_states = inputs_embeds + embed_pos
+        hidden_states = inputs_embeds + self.embed_positions(all_positions)
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
 
         encoder_states = () if output_hidden_states else None

filename: tests/models/whisper/test_modeling_whisper.py
patch: @@ -3356,22 +3356,6 @@ def test_forward_pass_weighted_layer_sum(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_forward(*config_and_inputs, use_weighted_layer_sum=True)
 
-    @unittest.skip(reason="Some undefined behavior encountered with tiny versions of this model. Skip for now.")
-    def test_cpu_offload(self):
-        pass
-
-    @unittest.skip(reason="Some undefined behavior encountered with tiny versions of this model. Skip for now.")
-    def test_disk_offload_bin(self):
-        pass
-
-    @unittest.skip(reason="Some undefined behavior encountered with tiny versions of this model. Skip for now.")
-    def test_disk_offload_safetensors(self):
-        pass
-
-    @unittest.skip(reason="Some undefined behavior encountered with tiny versions of this model. Skip for now.")
-    def test_model_parallelism(self):
-        pass
-
     @unittest.skip(reason="Not applicable for an encoder-only acoustic model")
     def test_inputs_embeds(self):
         # input embeds is meaningless for an encoder-only acoustic model