Why the inputs of baddbmm are fp32 the outputs of baddbmm are fp16?

Joyce_Chen: <p>I’m trying to run starcoder with QLORA and find that the loss has been 0 since the very first step.<br>
By debugging, I found that the inputs of baddbmm (attn_weights, query, key) are fp32 and the outputs are fp16.</p>
<blockquote>
<p><code>attn_weights = torch.baddbmm(attn_weights, query, key, beta=beta, alpha=scale_factor).view(attn_shape)</code></p>
</blockquote>
<p>Why the baddbmm has different datatypes for inputs and outputs?</p>

ptrblck: <p>What are the <code>dtype</code>s of all inputs? Also, are you executing the code in an <code>autocast</code> context?</p>

Joyce_Chen: <p>Thanks for your response.<br>
Regarding your questions,<br>
1、no autocast<br>
2、dtype: (attn_weights=float32, query=float32,key=float32,beta=int,alpha=float), return attn_weights=float16</p>

ptrblck: <p>In this case no casting will be performed as seen here:</p>
<pre><code class="lang-python">device = "cuda"
attn_weights = torch.randn(1, 10, 10, device=device)
query = torch.randn(1, 10, 10, device=device)
key = torch.randn(1, 10, 10, device=device)
beta = 1
scale_factor = 1.

attn_weights = torch.baddbmm(attn_weights, query, key, beta=beta, alpha=scale_factor)
print(attn_weights.dtype)
# torch.float32
</code></pre>

Joyce_Chen: <p>The issue appears in the Hugging Face Transformers code. The code file is located as follows:</p>
<blockquote>
<p>xxx\Anaconda3\envs\belle\Lib\site-packages\transformers\models\gpt_bigcode\modeling_gpt_bigcode.py</p>
</blockquote>
<p>The <code>_attn</code> function within this file is causing the problem.</p>
<blockquote>
<p>def _attn(self, query, key, value, attention_mask=None, head_mask=None):</p>
<pre><code>    dtype = query.dtype
    softmax_dtype = torch.float32 if self.attention_softmax_in_fp32 else dtype
    upcast = dtype != softmax_dtype

    unscale = self.layer_idx + 1 if self.scale_attention_softmax_in_fp32 and upcast else 1
    scale_factor = unscale**-1
    if self.scale_attn_weights:
        scale_factor /= self.head_dim**0.5

    query_shape = query.shape
    batch_size = query_shape[0]
    key_length = key.size(-1)
    if self.multi_query:
        # (batch_size, query_length, num_heads, head_dim) x (batch_size, head_dim, key_length)
        # -&gt; (batch_size, query_length, num_heads, key_length)
        query_length = query_shape[1]
        attn_shape = (batch_size, query_length, self.num_heads, key_length)
        attn_view = (batch_size, query_length * self.num_heads, key_length)
        # No copy needed for MQA 2, or when layer_past is provided.
        query = query.reshape(batch_size, query_length * self.num_heads, self.head_dim)
    else:
        # (batch_size, num_heads, query_length, head_dim) x (batch_size, num_heads, head_dim, key_length)
        # -&gt; (batch_size, num_heads, query_length, key_length)
        query_length = query_shape[2]
        attn_shape = (batch_size, self.num_heads, query_length, key_length)
        attn_view = (batch_size * self.num_heads, query_length, key_length)
        # Always copies
        query = query.reshape(batch_size * self.num_heads, query_length, self.head_dim)
        # No copy when layer_past is provided.
        key = key.reshape(batch_size * self.num_heads, self.head_dim, key_length)

    attn_weights = torch.empty(attn_view, device=query.device, dtype=query.dtype)
    if query.device.type == "cpu":
        # This is needed because of a bug in pytorch https://github.com/pytorch/pytorch/issues/80588.
        # The bug was fixed in https://github.com/pytorch/pytorch/pull/96086,
        # but the fix has not been released as of pytorch version 2.0.0.
        attn_weights.zero_()
        beta = 1
    else:
        beta = 0
</code></pre>
</blockquote>
<pre><code>     **##!!error here!!**
</code></pre>
<blockquote>
<pre><code>    attn_weights = torch.baddbmm(attn_weights, query, key, beta=beta, alpha=scale_factor).view(attn_shape)
    ......
</code></pre>
</blockquote>

Joyce_Chen: <p>Additional Information of testing environment…</p>
<blockquote>
<p>Collecting environment information…<br>
PyTorch version: 2.0.0+cu118<br>
Is debug build: False<br>
CUDA used to build PyTorch: 11.8<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Microsoft Windows Server 2019 Datacenter<br>
GCC version: (tdm64-1) 5.1.0<br>
Clang version: Could not collect<br>
CMake version: Could not collect<br>
Libc version: N/A</p>
<p>Python version: 3.9.16 (main, Mar  8 2023, 10:39:24) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)<br>
Python platform: Windows-10-10.0.17763-SP0<br>
Is CUDA available: True<br>
CUDA runtime version: 11.6.55<br>
CUDA_MODULE_LOADING set to: LAZY<br>
GPU models and configuration:<br>
GPU 0: Tesla V100-PCIE-32GB<br>
GPU 1: Tesla V100-PCIE-32GB<br>
GPU 2: Tesla V100-PCIE-32GB<br>
GPU 3: Tesla V100-PCIE-32GB</p>
<p>Nvidia driver version: 511.65<br>
cuDNN version: Could not collect<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A<br>
Is XNNPACK available: True</p>
<p>CPU:<br>
Architecture=9<br>
CurrentClockSpeed=3093<br>
DeviceID=CPU0<br>
Family=179<br>
L2CacheSize=20480<br>
L2CacheSpeed=<br>
Manufacturer=GenuineIntel<br>
MaxClockSpeed=3093<br>
Name=Intel(R) Xeon(R) Gold 6242R CPU @ 3.10GHz<br>
ProcessorType=3<br>
Revision=21767</p>
<p>Architecture=9<br>
CurrentClockSpeed=3093<br>
DeviceID=CPU1<br>
Family=179<br>
L2CacheSize=20480<br>
L2CacheSpeed=<br>
Manufacturer=GenuineIntel<br>
MaxClockSpeed=3093<br>
Name=Intel(R) Xeon(R) Gold 6242R CPU @ 3.10GHz<br>
ProcessorType=3<br>
Revision=21767</p>
<p>Versions of relevant libraries:<br>
[pip3] mypy-extensions==1.0.0<br>
[pip3] numpy==1.24.3<br>
[pip3] torch==2.0.0+cu118<br>
[pip3] torchaudio==2.0.2+cu118<br>
[pip3] torchvision==0.15.2+cu118<br>
[conda] numpy                     1.24.3                   pypi_0    pypi<br>
[conda] torch                     2.0.1+cu118              pypi_0    pypi<br>
[conda] torchaudio                2.0.2+cu118              pypi_0    pypi<br>
[conda] torchvision               0.15.2+cu118             pypi_0    pypi</p>
</blockquote>

Joyce_Chen: <p>I’ve done exactly the same testing on my lab and got the same result. However, when it runs in the code mentioned above, the output dtype becomes bf16.</p>

ptrblck: <p>I guess HF might use <code>autocast</code> for you behind your back without you realizing it.<br>
PyTorch won’t cast activations down without the user asking for it. Your code is unfortunately not executable, so if you get stuck, narrow it down further and post a minimal and executable code snippet to reproduce the issue.</p>

Joyce_Chen: <p>Thanks for your support! I’ll try to narrow it down.</p>

ptrblck: <p>In case you have access to the <code>forward</code> methods you could use <code>torch.is_autocast_enabled()</code> to check if you are inside an <code>autocast</code> context to further narrow it down.</p>

Joyce_Chen: <p>Thanks for the advice！ I inserted <code>torch.is_autocast_enabled()</code> in the code and found it is actually set to TRUE.</p>
<p>Thanks for your support! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>OK, great! Then I would indeed believe a higher level API enables it for you.</p>