Тип bfloat16

Тип данных bfloat16 (Brain Floating Point) - это 16-битный формат с плавающей точкой, разработанный для ускорения вычислений в глубоком обучении. В отличие от стандартного float16, он использует 8 бит для экспоненты (как у float32) и 7 бит для мантиссы, что обеспечивает больший динамический диапазон и устойчивость к переполнениям при обучении нейросетей.

Основное преимущество bfloat16 - сохранение числового диапазона 32-битного формата при сокращении точности мантиссы. Это делает его идеальным выбором для современных графических процессоров, поддерживающих аппаратное ускорение операций с bfloat16, таких как NVIDIA Ampere и последующие поколения.

Создание тензора с типом bfloat16

Для создания тензора с типом bfloat16 передайте параметр dtype со значением torch.bfloat16 в функцию tensor:

import torch

t = torch.tensor([1, 2, 3, 4, 5], dtype=torch.bfloat16)
print(t)

Результат выполнения кода:

tensor([1., 2., 3., 4., 5.], dtype=torch.bfloat16)

Преобразование существующего тензора

Изменить тип данных уже существующего тензора на bfloat16 можно с помощью метода to:

import torch

t = torch.tensor([1.5, 2.7, 3.1, 4.8, 5.2])
t_bf16 = t.to(torch.bfloat16)
print(t_bf16)

Результат выполнения кода:

tensor([1.5000, 2.7000, 3.1000, 4.8000, 5.2000], dtype=torch.bfloat16)

Сравнение bfloat16 и float32

Давайте сравним точность хранения чисел в форматах bfloat16 и float32 на примере сложения чисел с разным порядком:

import torch

t1 = torch.tensor([1.0, 100.0, 10000.0], dtype=torch.bfloat16)
t2 = torch.tensor([0.1, 10.0, 1000.0], dtype=torch.bfloat16)
res_bf16 = t1 + t2

t1_f32 = torch.tensor([1.0, 100.0, 10000.0], dtype=torch.float32)
t2_f32 = torch.tensor([0.1, 10.0, 1000.0], dtype=torch.float32)
res_f32 = t1_f32 + t2_f32

print("bfloat16:", res_bf16)
print("float32:", res_f32)

Результат выполнения кода:

bfloat16: tensor([1.1000, 110.0000, 11000.0000], dtype=torch.bfloat16)
float32: tensor([1.1000, 110.0000, 11000.0000])

Использование в вычислениях

Тензоры типа bfloat16 можно использовать в математических операциях, но важно помнить, что результаты будут иметь ограниченную точность:

import torch

t = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.bfloat16)
res = t * t
print(res)

Результат выполнения кода:

tensor([
	[1., 4., 9.],
	[16., 25., 36.]
], dtype=torch.bfloat16)

Особенности работы на GPU

При использовании графического процессора тип bfloat16 позволяет ускорить вычисления и снизить потребление памяти. Проверить доступность поддержки можно с помощью функции is_bf16_supported:

import torch

if torch.cuda.is_bf16_supported():
	t = torch.tensor([1, 2, 3, 4, 5], dtype=torch.bfloat16, device='cuda')
	print("bfloat16 на GPU:", t)
else:
	print("bfloat16 не поддерживается на этом GPU")

Смотрите также

тип float32,
который является стандартным 32-битным форматом с плавающей точкой
тип float16,
который использует 16 бит с уменьшенной экспонентой
тип float64,
который предоставляет двойную точность за счёт 64 бит
функцию promote_types,
которая определяет результирующий тип при смешанных операциях