第8章：混合精度训练
DGL 兼容 PyTorch 自动混合精度 (AMP) 包，用于混合精度训练，从而节省训练时间和 GPU/CPU 内存消耗。此功能需要 DGL 0.9+ 和 1.1+ 以支持 CPU 的 bloat16。
使用半精度进行消息传递
DGL 允许在 float16 (fp16) / bfloat16 (bf16) 特征上进行消息传递，适用于用户定义函数（UDFs）和内置函数（例如，dgl.function.sum, dgl.function.copy_u）。
注意
在使用之前，请通过torch.cuda.is_bf16_supported()检查bfloat16支持。
通常需要CUDA >= 11.0和GPU计算能力 >= 8.0。
以下示例展示了如何在半精度特征上使用DGL的消息传递API：
>>> import torch
>>> import dgl
>>> import dgl.function as fn
>>> dev = torch.device('cuda')
>>> g = dgl.rand_graph(30, 100).to(dev) # Create a graph on GPU w/ 30 nodes and 100 edges.
>>> g.ndata['h'] = torch.rand(30, 16).to(dev).half() # Create fp16 node features.
>>> g.edata['w'] = torch.rand(100, 1).to(dev).half() # Create fp16 edge features.
>>> # Use DGL's built-in functions for message passing on fp16 features.
>>> g.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'x'))
>>> g.ndata['x'].dtype
torch.float16
>>> g.apply_edges(fn.u_dot_v('h', 'x', 'hx'))
>>> g.edata['hx'].dtype
torch.float16
>>> # Use UDFs for message passing on fp16 features.
>>> def message(edges):
... return {'m': edges.src['h'] * edges.data['w']}
...
>>> def reduce(nodes):
... return {'y': torch.sum(nodes.mailbox['m'], 1)}
...
>>> def dot(edges):
... return {'hy': (edges.src['h'] * edges.dst['y']).sum(-1, keepdims=True)}
...
>>> g.update_all(message, reduce)
>>> g.ndata['y'].dtype
torch.float16
>>> g.apply_edges(dot)
>>> g.edata['hy'].dtype
torch.float16
端到端混合精度训练
DGL 依赖 PyTorch 的 AMP 包进行混合精度训练，用户体验与 PyTorch 的 完全相同。
通过使用torch.amp.autocast()包装前向传递，PyTorch会自动为每个操作和张量选择适当的数据类型。半精度张量内存效率高，大多数半精度张量上的操作更快，因为它们利用了GPU的张量核心和CPU的特殊指令集。
import torch.nn.functional as F
from torch.amp import autocast
def forward(device_type, g, feat, label, mask, model, amp_dtype):
amp_enabled = amp_dtype in (torch.float16, torch.bfloat16)
with autocast(device_type, enabled=amp_enabled, dtype=amp_dtype):
logit = model(g, feat)
loss = F.cross_entropy(logit[mask], label[mask])
return loss
在float16格式中的小梯度存在下溢问题（归零）。
PyTorch提供了一个GradScaler模块来解决这个问题。它将损失乘以一个因子，并在缩放后的损失上调用反向传播，以防止下溢问题。然后，在优化器更新参数之前，它会将计算出的梯度进行反缩放。缩放因子是自动确定的。
请注意，bfloat16不需要GradScaler。
from torch.cuda.amp import GradScaler
scaler = GradScaler()
def backward(scaler, loss, optimizer):
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
以下示例在Reddit数据集（包含1.14亿条边）上训练一个3层GAT。 请注意在激活AMP与否时代码中的差异。
import torch
import torch.nn as nn
import dgl
from dgl.data import RedditDataset
from dgl.nn import GATConv
from dgl.transforms import AddSelfLoop
amp_dtype = torch.bfloat16 # or torch.float16
class GAT(nn.Module):
def __init__(self,
in_feats,
n_hidden,
n_classes,
heads):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(GATConv(in_feats, n_hidden, heads[0], activation=F.elu))
self.layers.append(GATConv(n_hidden * heads[0], n_hidden, heads[1], activation=F.elu))
self.layers.append(GATConv(n_hidden * heads[1], n_classes, heads[2], activation=F.elu))
def forward(self, g, h):
for l, layer in enumerate(self.layers):
h = layer(g, h)
if l != len(self.layers) - 1:
h = h.flatten(1)
else:
h = h.mean(1)
return h
# Data loading
transform = AddSelfLoop()
data = RedditDataset(transform)
device_type = 'cuda' # or 'cpu'
dev = torch.device(device_type)
g = data[0]
g = g.int().to(dev)
train_mask = g.ndata['train_mask']
feat = g.ndata['feat']
label = g.ndata['label']
in_feats = feat.shape[1]
n_hidden = 256
n_classes = data.num_classes
heads = [1, 1, 1]
model = GAT(in_feats, n_hidden, n_classes, heads)
model = model.to(dev)
model.train()
# Create optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
for epoch in range(100):
optimizer.zero_grad()
loss = forward(device_type, g, feat, label, train_mask, model, amp_dtype)
if amp_dtype == torch.float16:
# Backprop w/ gradient scaling
backward(scaler, loss, optimizer)
else:
loss.backward()
optimizer.step()
print('Epoch {} | Loss {}'.format(epoch, loss.item()))
在NVIDIA V100（16GB）机器上，训练此模型时不使用fp16会消耗15.2GB的GPU内存；启用fp16后，训练消耗12.8G GPU内存，两种设置下的损失收敛到相似的值。如果我们将头数更改为[2, 2, 2]，不使用fp16训练会触发GPU内存不足（OOM）问题，而使用fp16训练则消耗15.7G GPU内存。
BFloat16 CPU 示例
DGL 支持在 CPU 上以 bfloat16 数据类型运行训练。 这种数据类型不需要任何 CPU 特性，并且可以提高内存受限模型的性能。 从具有 AMX 指令集的英特尔至强第四代处理器开始，bfloat16 应该能够显著提高训练和推理性能，而无需进行大量的代码更改。 以下是一个简单的 GCN bfloat16 训练示例：
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
from dgl.data import CiteseerGraphDataset
from dgl.nn import GraphConv
from dgl.transforms import AddSelfLoop
class GCN(nn.Module):
def __init__(self, in_size, hid_size, out_size):
super().__init__()
self.layers = nn.ModuleList()
# two-layer GCN
self.layers.append(
GraphConv(in_size, hid_size, activation=F.relu)
)
self.layers.append(GraphConv(hid_size, out_size))
self.dropout = nn.Dropout(0.5)
def forward(self, g, features):
h = features
for i, layer in enumerate(self.layers):
if i != 0:
h = self.dropout(h)
h = layer(g, h)
return h
# Data loading
transform = AddSelfLoop()
data = CiteseerGraphDataset(transform=transform)
g = data[0]
g = g.int()
train_mask = g.ndata['train_mask']
feat = g.ndata['feat']
label = g.ndata['label']
in_size = feat.shape[1]
hid_size = 16
out_size = data.num_classes
model = GCN(in_size, hid_size, out_size)
# Convert model and graph to bfloat16
g = dgl.to_bfloat16(g)
feat = feat.to(dtype=torch.bfloat16)
model = model.to(dtype=torch.bfloat16)
model.train()
# Create optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
loss_fcn = nn.CrossEntropyLoss()
for epoch in range(100):
logits = model(g, feat)
loss = loss_fcn(logits[train_mask], label[train_mask])
loss.backward()
optimizer.step()
print('Epoch {} | Loss {}'.format(epoch, loss.item()))
与常规训练的唯一区别在于训练/推理前的模型和图转换。
DGL 仍在改进其半精度支持，计算内核的性能远未达到最佳，请继续关注我们未来的更新。