[Modeling] Fix encoder CPU offloading for whisper

[kylesayrs]

Purpose

• Fix CPU offloading for WhisperEncoder
  • This supports downstream applications which offload the layers of even small models like whisper. See [Performance] Sequential onloading vllm-project/llm-compressor#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

device_map={
    "model.encoder": "cpu",
    "model.decoder": 0,
    "proj_out": 0,
},

test_whisper_offload.py

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))

Potential Reviewers

• @SunMarc @ArthurZucker @zucchini-nlp

[zucchini-nlp]

Thanks, I like that we removed direct access to weight.data. Can you also un-skip offload tests in whisper and make sure they are green?

For ex:

transformers/tests/models/whisper/test_modeling_whisper.py

Lines 3357 to 3368 in 21cb353

	@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

[kylesayrs]

@zucchini-nlp @SunMarc Tests unskipped and passing!

[zucchini-nlp]

run-slow: whisper

[github-actions]

This comment contains run-slow, running the specified jobs:

models: ['models/whisper']
quantizations: [] ...

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[kylesayrs]

@zucchini-nlp Does this test failure indicate something to fix, or is this test noisy?

First differing element 3:
" Fol[1422 chars]ugitives bug out bindle of news that is my segment. Meanwhile."
" Fol[1422 chars]ugitives bug out bindle of news that is my segment. Meanwhile!"

[vasqu]

@kylesayrs Can you rebase/merge? The failing tests are expected, no worries :D

cc @gante @eustlb for viz

[kylesayrs]

@vasqu Merged, thank to hear it :)

[vasqu]

Thanks @kylesayrs 🤗