第 8 章：混合精度训练
DGL 兼容 PyTorch 自动混合精度 (AMP) 包，用于混合精度训练，从而节省训练时间和 GPU/CPU 内存消耗。此功能需要 DGL 0.9+，以及 CPU 的 bloat16 需要 1.1+。
半精度消息传递
DGL 支持在 float16 (fp16) / bfloat16 (bf16) 特征上进行消息传递，对 UDF (用户定义函数) 和内置函数（例如 dgl.function.sum, dgl.function.copy_u）均适用。
注意
在使用 bfloat16 之前，请通过 torch.cuda.is_bf16_supported() 检查其支持情况。通常它需要 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 Tensorcores 和 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 仍在改进其半精度支持，计算核的性能远未达到最佳状态，请继续关注我们未来的更新。