cohere_asr: fix bug for model_parallel_beam_search test case

[kaixuanliu]

This PR fixes failed test case: tests/models/cohere_asr/test_modeling_cohere_asr.py::CohereAsrModelTest::test_model_parallel_beam_search, and add some adjustment to make the test cases pass for Intel XPU device. @ydshieh pls help review, thx!

[ydshieh]

Could you share the (full) error log you have for this test, please 🙏 ? Thanks.

I don't see this model's tests using device_map = "auto", so it's a bit strange we have the device issue.

[kaixuanliu]

When I use 4 cards: export CUDA_VISIBLE_DEVICES=0,4,5,6, and for this test case, it will throw error:

encoder_hidden_states = self.proj(encoder_hidden_states)
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))

        if position_ids is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
            position_ids = position_ids.unsqueeze(0)

        # Fixed sinusoidal position embedding added to token embeddings, then layernorm
        pos_emb = self.pos_emb(position_ids.squeeze(0))
>       inputs_embeds = self.embedding_layernorm(inputs_embeds + pos_emb)
                                                 ^^^^^^^^^^^^^^^^^^^^^^^
E       RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:3!

src/transformers/models/cohere_asr/modeling_cohere_asr.py:388: RuntimeError

[ydshieh]

ok, my bad, this test is actually using "auto"

new_model = model_class.from_pretrained(tmp_dir, device_map="auto")

[ydshieh]

run-slow: cohere_asr

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: cohere_asr

[github-actions]

Workflow Run ⚙️

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

models: ["models/cohere_asr"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	e7ab438d	workflow commit (merge commit)
PR	3e0186bd	branch commit (from PR)
main	47d77657	base commit (on main)

Model CI Report

❌ 1 new failed tests from this PR 😭

• cohere_asr:
  tests/models/cohere_asr/test_modeling_cohere_asr.py::CohereAsrIntegrationTest::test_shortform_english_no_punctuation (❌ ⟹ ❌)

[ydshieh]

run-slow: cohere_asr

[github-actions]

Workflow Run ⚙️

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

models: ["models/cohere_asr"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	fab7171c	workflow commit (merge commit)
PR	b0a56481	branch commit (from PR)
main	4fd862fc	base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !