ValueError: Cannot find backend for cpu in flash_attn/ops/triton/rotary.py

username: Pritish_Mishra
cooked: <h3><a name="issue-description-1" class="anchor" href="#issue-description-1"></a>Issue Description</h3>
<h4><a name="problem-2" class="anchor" href="#problem-2"></a>Problem</h4>
<p>Executing the provided code results in a <code>ValueError: Cannot find backend for cpu</code> originating from <code>flash_attn/ops/triton/rotary.py</code>.</p>
<h4><a name="context-3" class="anchor" href="#context-3"></a>Context</h4>
<p>The code employs the <code>flash_attn</code> library in conjunction with a pretrained model for text generation. The specific error occurs during the application of rotary transformations.</p>
<h4><a name="steps-to-reproduce-4" class="anchor" href="#steps-to-reproduce-4"></a>Steps to Reproduce</h4>
<ol>
<li>
<p>Import necessary libraries:</p>
<pre><code class="lang-python">import transformers
import torch
import time
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
</code></pre>
</li>
<li>
<p>Define base model name and new model path:</p>
<pre><code class="lang-python">base_model_name = 'NousResearch/Yarn-Llama-2-7b-128k'
new_model_path = '/opt/llama_models/testing'
</code></pre>
</li>
<li>
<p>Load the base model:</p>
<pre><code class="lang-python">base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    trust_remote_code=True
)
</code></pre>
</li>
<li>
<p>Initialize a <code>PeftModel</code> using the base model and a custom path:</p>
<pre><code class="lang-python">model = PeftModel.from_pretrained(base_model, new_model_path)
</code></pre>
</li>
<li>
<p>Merge and unload the model:</p>
<pre><code class="lang-python">model = model.merge_and_unload()
</code></pre>
</li>
<li>
<p>Initialize a tokenizer:</p>
<pre><code class="lang-python">tokenizer = AutoTokenizer.from_pretrained(base_model_name)
tokenizer.padding_side = "right"
tokenizer.add_eos_token = True
</code></pre>
</li>
<li>
<p>Set up a text generation pipeline:</p>
<pre><code class="lang-python">pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto"
)
</code></pre>
</li>
<li>
<p>Provide a prompt and specify generation parameters:</p>
<pre><code class="lang-python">prompt = "something here".strip()
max_new_tokens = 1
</code></pre>
</li>
<li>
<p>Generate text:</p>
<pre><code class="lang-python">start_time = time.time()
sequences = pipeline(
    prompt,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    max_new_tokens=max_new_tokens
)
elapsed_time = time.time() - start_time
</code></pre>
</li>
</ol>
<h4><a name="expected-behavior-5" class="anchor" href="#expected-behavior-5"></a>Expected Behavior</h4>
<p>The code should generate text using the specified prompt and parameters without encountering any errors.</p>
<h4><a name="actual-behavior-6" class="anchor" href="#actual-behavior-6"></a>Actual Behavior</h4>
<p>The code raises a <code>ValueError: Cannot find backend for cpu</code> during the rotary transformation step.</p>
<p>Full Error:</p>
<pre><code class="lang-auto">...
/usr/local/lib/python3.8/dist-packages/flash_attn/layers/rotary.py in forward(ctx, x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen)
     46         max_seqlen: Optional[int] = None,
     47     ):
---&gt; 48         out = apply_rotary(
     49             x,
     50             cos,

/usr/local/lib/python3.8/dist-packages/flash_attn/ops/triton/rotary.py in apply_rotary(x, cos, sin, seqlen_offsets, cu_seqlens, max_seqlen, interleaved, inplace, conjugate)
    211     # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
    212     with torch.cuda.device(x.device.index):
--&gt; 213         rotary_kernel[grid](
    214             output,  # data ptrs
    215             x,

&lt;string&gt; in rotary_kernel(OUT, X, COS, SIN, CU_SEQLENS, SEQLEN_OFFSETS, seqlen, nheads, rotary_dim, seqlen_ro, CACHE_KEY_SEQLEN, stride_out_batch, stride_out_seqlen, stride_out_nheads, stride_out_headdim, stride_x_batch, stride_x_seqlen, stride_x_nheads, stride_x_headdim, BLOCK_K, IS_SEQLEN_OFFSETS_TENSOR, IS_VARLEN, INTERLEAVED, CONJUGATE, BLOCK_M, grid, num_warps, num_stages, extern_libs, stream, warmup, device, device_type)

ValueError: Cannot find backend for cpu
</code></pre>
<h4><a name="environment-7" class="anchor" href="#environment-7"></a>Environment</h4>
<ul>
<li>Python Version: 3.8.10</li>
<li>PyTorch Version: 2.0.1+cu118</li>
<li>Transformers Version: 4.33.0.dev0</li>
<li>flash-attn Version: 2.2.2</li>
<li>Triton Version: 2.1.0</li>
</ul>
<h4><a name="additional-information-8" class="anchor" href="#additional-information-8"></a>Additional Information</h4>
<ul>
<li>Before this error I was encountering <a href="https://github.com/Dao-AILab/flash-attention/issues/523" rel="noopener nofollow ugc">this</a> error but it got solved after upgrading Triton to 2.1.0 but right after that I am stuck with this error.</li>
</ul>

username: ptrblck
cooked: <p>You are explicitly using the GPU via:</p>
<pre><code class="lang-python">with torch.cuda.device
</code></pre>
<p>so did you check if <code>flash_attn</code> supports CPU-only workloads?</p>