Daily Dev Tips

How to Fix "triu_tril_cuda_template" not implemented for 'BFloat16'

Qualiteg Product Development Team

Qualiteg Product Development Team

— 1 min read

This post explains how to deal with the error "triu_tril_cuda_template" not implemented for 'BFloat16', which can occur when you specify torch_dtype=torch.bfloat16 while loading a model.

Below is the log from when this occurred with Llama 3.

  File "/home/mlu/.virtualenvs/ChatStream/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mlu/.virtualenvs/ChatStream/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1208, in forward
    outputs = self.model(
  File "/home/mlu/.virtualenvs/ChatStream/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mlu/.virtualenvs/ChatStream/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 992, in forward
    causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens)
  File "/home/mlu/.virtualenvs/ChatStream/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1095, in _update_causal_mask
    causal_mask = torch.triu(causal_mask, diagonal=1)
RuntimeError: "triu_tril_cuda_template" not implemented for 'BFloat16'

This problem occurs when your PyTorch version is 2.0.1 or earlier.

Check your torch version with pip list.

pip list

As shown below, in 2.0.1 triu_tril_cuda_template is quite literally not implemented for this dtype, which is why the error is raised.

Fix

Updating PyTorch to the latest version resolves the problem.

  • CUDA 12.x
pip install --upgrade torch torchvision torchaudio
  • CUDA 11.8

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Pinning a specific version also works.

pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0

Read more