<h3 dir="auto">System Info</h3>
<p dir="auto"><code class="notranslate">transformers-cli env</code>=</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/user/.local/bin/transformers-cli&quot;, line 5, in &lt;module&gt;
    from transformers.commands.transformers_cli import main
  File &quot;/home/user/.local/lib/python3.10/site-packages/transformers/commands/transformers_cli.py&quot;, line 24, in &lt;module&gt;
    from .pt_to_tf import PTtoTFCommand
  File &quot;/home/user/.local/lib/python3.10/site-packages/transformers/commands/pt_to_tf.py&quot;, line 21, in &lt;module&gt;
    from datasets import load_dataset
ModuleNotFoundError: No module named 'datasets'"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "/home/user/.local/bin/transformers-cli", line 5, in &lt;module&gt;
    from transformers.commands.transformers_cli import main
  File "/home/user/.local/lib/python3.10/site-packages/transformers/commands/transformers_cli.py", line 24, in &lt;module&gt;
    from .pt_to_tf import PTtoTFCommand
  File "/home/user/.local/lib/python3.10/site-packages/transformers/commands/pt_to_tf.py", line 21, in &lt;module&gt;
    from datasets import load_dataset
ModuleNotFoundError: No module named 'datasets'
</code></pre></div>
<p dir="auto"><a href="https://colab.research.google.com/drive/1YORPWx4okIHXnjW7MSAidXN29mPVNT7F?usp=sharing" rel="nofollow">https://colab.research.google.com/drive/1YORPWx4okIHXnjW7MSAidXN29mPVNT7F?usp=sharing</a></p>
<h3 dir="auto">Who can help?</h3>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/sgugger/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/sgugger">@sgugger</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patil-suraj/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patil-suraj">@patil-suraj</a></p>
<p dir="auto">Im following:<br>
<a href="https://colab.research.google.com/drive/1YORPWx4okIHXnjW7MSAidXN29mPVNT7F" rel="nofollow">https://colab.research.google.com/drive/1YORPWx4okIHXnjW7MSAidXN29mPVNT7F</a><br>
from:<br>
<a href="https://huggingface.co/blog/hf-bitsandbytes-integration" rel="nofollow">https://huggingface.co/blog/hf-bitsandbytes-integration</a></p>
<h3 dir="auto">Information</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> The official example scripts</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> My own modified scripts</li>
</ul>
<h3 dir="auto">Tasks</h3>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> An officially supported task in the <code class="notranslate">examples</code> folder (such as GLUE/SQuAD, ...)</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> My own task or dataset (give details below)</li>
</ul>
<h3 dir="auto">Reproduction</h3>
<p dir="auto">I love being able to use the int8 models! And being able to run the big LLMs on Colab. However, I just discovered that only the default search algo works. This is too bad, because for a chat application, the TopK sampling provides much more natural variation.</p>
<p dir="auto">Beam:<br>
RuntimeError: "log_softmax_lastdim_kernel_impl" not implemented for 'Half'<br>
Topk:<br>
RuntimeError: "topk_cpu" not implemented for 'Half'</p>
<h3 dir="auto">Expected behavior</h3>
<p dir="auto">being able to use TopK and TopP sampling on int8 optimized models</p>
<p dir="auto">When using a bloom model with generate, I get</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" RuntimeError: &quot;topk_cpu&quot; not implemented for 'Half'"><pre class="notranslate"><code class="notranslate"> RuntimeError: "topk_cpu" not implemented for 'Half'
</code></pre></div>
<p dir="auto">when <code class="notranslate">do_sample=True</code>.</p>
<p dir="auto">I.e.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MAX_NEW_TOKENS = 128
model_name = 'bigscience/bloom-560m'

text = &quot;&quot;&quot;
Q: On average Joe throws 25 punches per minute. A fight lasts 5 rounds of 3 minutes. 
How many punches did he throw?\n
A: Let’s think step by step.\n&quot;&quot;&quot;
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_ids = tokenizer(text, return_tensors=&quot;pt&quot;).input_ids

free_in_GB = int(torch.cuda.mem_get_info()[0]/1024**3)
max_memory = f'{free_in_GB-6}GB'

n_gpus = torch.cuda.device_count()
max_memory = {i: max_memory for i in range(n_gpus)}

model = AutoModelForCausalLM.from_pretrained(
  model_name,
  cache_dir=&quot;/home/code/fm_in_context_eval/transformers_cache&quot;,
  device_map='auto', 
  load_in_8bit=True, 
  max_memory=max_memory
)
generated_ids = model.generate(input_ids, max_length=len(input_ids[0])+1, do_sample=True)"><pre class="notranslate"><code class="notranslate">import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MAX_NEW_TOKENS = 128
model_name = 'bigscience/bloom-560m'

text = """
Q: On average Joe throws 25 punches per minute. A fight lasts 5 rounds of 3 minutes. 
How many punches did he throw?\n
A: Let’s think step by step.\n"""
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_ids = tokenizer(text, return_tensors="pt").input_ids

free_in_GB = int(torch.cuda.mem_get_info()[0]/1024**3)
max_memory = f'{free_in_GB-6}GB'

n_gpus = torch.cuda.device_count()
max_memory = {i: max_memory for i in range(n_gpus)}

model = AutoModelForCausalLM.from_pretrained(
  model_name,
  cache_dir="/home/code/fm_in_context_eval/transformers_cache",
  device_map='auto', 
  load_in_8bit=True, 
  max_memory=max_memory
)
generated_ids = model.generate(input_ids, max_length=len(input_ids[0])+1, do_sample=True)
</code></pre></div>