How to use float16 for all tensor operations?

mixed-precision
1

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?

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)

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.

I can provide the sample row and column files but the website does not allow me to upload NumPy files.

PyTorch Version: 2.1.1
NumPy Version: 1.26.2

2

No, pure float16 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.

1 Like
3

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

4

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.

In PyTorch, it’s torch.bfloat16.

5

Thanks for sharing this information. IIRC, bfloat16 uses float32 for accumulation.