How to use float16 for all tensor operations?

MahdiNazemi: <p>When calculating the dot product of two half-precision vectors, it appears that PyTorch uses float32 for accumulation, and finally converts the output back to float16. Is it possible to carry out all operations in float16?</p>
<pre data-code-wrap="py"><code class="lang-py">import numpy as np
import torch

row = np.load("row.npy")  # (1, 4096)
col = np.load("col.npy")  # (4096, 1)
DIM = 4096

# Calculate the output using the dot product function
np.dot(row, col)
# array([[-0.01642]], dtype=float16)

# Accumulate the result in a float16 variable
result = np.float16(0)
for i in range(DIM):
    result += row[0, i] * col[i, 0]
print(result)
# -0.01443

# Accumulate the result in a float32 variable
result = np.float32(0)
for i in range(DIM):
    result += row[0, i] * col[i, 0]
print(result)
# -0.016465545

row = torch.from_numpy(row).cuda()
col = torch.from_numpy(col).cuda()

# Calculate the output using the dot product function
torch.dot(row[0], col[:, 0])
# tensor(-0.0164, device='cuda:0', dtype=torch.float16)

# Accumulate the result in a float16 variable
result = torch.tensor([0]).half().cuda()
for i in range(DIM):
    result += row[0, i] * col[i, 0]
print(result)
# tensor([-0.0144], device='cuda:0', dtype=torch.float16)
</code></pre>
<p>So it looks like PyTorch’s dot product calculates results that are very close to that of NumPy’s, but they both use float32 for accumulation.</p>
<p>I can provide the sample row and column files but the website does not allow me to upload NumPy files.</p>
<p>PyTorch Version: 2.1.1<br>
NumPy Version: 1.26.2</p>

ptrblck: <p>No, pure <code>float16</code> accumulations are not provided as they were causing convergence issues in the past if I’m not mistaken. If you need them for some reason you might want to write a custom extension.</p>

MahdiNazemi: <p>Thanks very much. I’ll try to find the existing implementation and see if I can modify it to use float16 for all operations.</p>

J_Johnson: <p>Are you familiar with bfloat16? It was developed specifically for machine learning. It has the range of float32 but lower precision, which, it turns out, is not as important for individual parameters in a model.</p>
<aside class="onebox wikipedia" data-onebox-src="https://en.m.wikipedia.org/wiki/Bfloat16_floating-point_format">
  <header class="source">

      <a href="https://en.m.wikipedia.org/wiki/Bfloat16_floating-point_format" target="_blank" rel="noopener nofollow ugc">en.m.wikipedia.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://en.m.wikipedia.org/wiki/Bfloat16_floating-point_format" target="_blank" rel="noopener nofollow ugc">bfloat16 floating-point format</a></h3>

<p>The bfloat16 (brain floating point) floating-point format is a computer number format occupying 16 bits in computer memory; it represents a wide dynamic range of numeric values by using a floating radix point. This format is a shortened (16-bit) version of the 32-bit IEEE 754 single-precision floating-point format (binary32) with the intent of accelerating machine learning and near-sensor computing. It preserves the approximate dynamic range of 32-bit floating-point numbers by retaining  The bflo...</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>In PyTorch, it’s <code>torch.bfloat16</code>.</p>

MahdiNazemi: <p>Thanks for sharing this information. IIRC, bfloat16 uses float32 for accumulation.</p>