性能优化¶

本指南是快速入门指南中讨论的后续内容。我们将重点介绍在训练基本 GPT 编码器层时实现最佳性能的技术。为方便起见，我们使用 quickstart_utils.py 中定义的一些辅助函数。

[1]:

import torch
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling
import quickstart_utils as utils

# Layer configuration
hidden_size = 4096
sequence_length = 2048
batch_size = 4
ffn_hidden_size = 16384
num_attention_heads = 32
dtype = torch.float16

# Synthetic data
x = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)
dy = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)

[2]:

# Construct layer
basic_transformer = te.TransformerLayer(
    hidden_size,
    ffn_hidden_size,
    num_attention_heads,
)
basic_transformer.to(dtype=dtype).cuda()

fp8_format = Format.HYBRID
fp8_recipe = DelayedScaling(
    fp8_format=fp8_format,
    amax_history_len=16,
    amax_compute_algo="max",
)
# Training step
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = basic_transformer(x, attention_mask=None)
y.backward(dy)

# Measure step time
utils.speedometer(
    basic_transformer,
    x,
    dy,
    forward_kwargs = { "attention_mask": None },
    fp8_autocast_kwargs = { "enabled": True, "fp8_recipe": fp8_recipe },
)

Mean time: 27.82952880859375 ms

多 GPU 训练¶

摘要

我们使用数据并行、张量并行和序列并行来并行化 Transformer 层。

可以使用多种并行策略来启用 Transformer 模型的多 GPU 训练，这些策略通常基于不同的方法来分配它们的 \(\text{sequence_length} \times \text{batch_size} \times \text{hidden_size}\) 激活张量。最常见的方法是数据并行，它沿着 \(\text{batch_size}\) 维度进行分配。通过在每个 GPU 上存储模型的重复副本，可以独立完成训练步骤的前向和后向传递，然后进行梯度同步。一种更高级的策略是张量并行，这是一种模型并行类型，它沿着 \(\text{hidden_size}\) 维度进行分配。这使我们能够扩展超过数据并行的限制（通常 \(\text{hidden_size} > \text{batch_size}\)），并减少每个 GPU 的内存使用量（因为模型参数也被分配），但它也会产生在每个步骤中在 GPU 之间传递激活张量的开销。有关更详细的说明，请参阅 Megatron-LM 论文。最后，序列并行沿着 \(\text{sequence_length}\) 维度进行分配。这可以在启用张量并行时使用，以便并行化在张量并行区域之外运行的操作（例如，层归一化）。有关更多详细信息，请参阅这篇论文。

为了展示这一点，让我们首先使用一个简单的进程组初始化 NCCL

[3]:

# Configure parallel groups
import os
import torch
world_group = torch.distributed.init_process_group(
    "nccl",
    init_method="file:///tmp/rdzv",
    world_size=1,
    rank=0,
)
data_parallel_group = torch.distributed.new_group(ranks=[0], backend="nccl")
tensor_parallel_group = torch.distributed.new_group(ranks=[0], backend="nccl")

我们只使用一个 GPU 进行初始化，以使本示例简单。有关使用多个 GPU 运行的指导，请查阅文档 torch.distributed。请注意，我们要求每个分布式进程精确对应于一个 GPU，因此我们将它们互换使用。在实践中，有多种因素会影响最佳并行布局：系统硬件、网络拓扑、其他并行方案（如流水线并行）的使用。一个粗略的经验法则是将 GPU 解释为一个 2D 网格，其维度为 \(\text{num_nodes} \times \text{gpus_per_node}\)。行是张量并行组，列是数据并行组。

使用 Transformer Engine 启用数据并行类似于使用标准 PyTorch 模型启用数据并行：只需使用 torch.nn.parallel.DistributedDataParallel 包装模块。FP8 训练需要对缩放因子进行额外的同步，因此数据并行进程组也必须传递给 fp8_autocast 上下文管理器。Transformer Engine 模块还原生支持张量并行和序列并行。如果用户为张量并行提供进程组，则模块将在内部分配数据并执行通信。如果启用了序列并行，它将应用于不适合张量并行的操作，并将使用张量并行进程组。在这种情况下，张量并行组还必须传递给 fp8_autocast 上下文管理器中的 fp8_group 参数，可以直接传递，也可以作为更大的分布式组的子集传递。

[4]:

# Construct layer
parallel_transformer = te.TransformerLayer(
    hidden_size,
    ffn_hidden_size,
    num_attention_heads,
    set_parallel_mode=True,
    tp_group=tensor_parallel_group,
    sequence_parallel=True,
)
parallel_transformer.to(dtype=dtype).cuda()
parallel_transformer = torch.nn.parallel.DistributedDataParallel(
    parallel_transformer,
    process_group=data_parallel_group,
)

# Training step
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=data_parallel_group):
    y = parallel_transformer(x, attention_mask=None)
y.backward(dy)

# Measure step time
utils.speedometer(
    parallel_transformer,
    x,
    dy,
    forward_kwargs = { "attention_mask": None },
    fp8_autocast_kwargs = {
        "enabled": True,
        "fp8_recipe": fp8_recipe,
        "fp8_group": data_parallel_group,
    },
)

Mean time: 29.09606689453125 ms

梯度累积融合¶

摘要

我们利用 Tensor Core 的能力将输出直接累积到 FP32 中。

PyTorch 的 autograd 功能假设模型参数及其对应的梯度具有相同的数据类型。但是，虽然像 FP8 这样的低精度数据类型足以评估神经网络的前向和后向传递，但优化步骤通常需要完整的 FP32 精度，以避免显著的学习退化。此外，Hopper GPU 上的 Tensor Core 可以选择将矩阵乘积直接累积到 FP32 中，从而获得更好的数值精度，并避免需要单独的类型转换内核。因此，Transformer Engine 提供了一个选项，可以直接为权重张量生成 FP32 梯度。FP32 梯度不会输出到参数的 grad 张量，而是输出到在后向传递之前必须初始化的 main_grad 张量。

[5]:

# Construct layer
wgrad_transformer = te.TransformerLayer(
    hidden_size,
    ffn_hidden_size,
    num_attention_heads,
    fuse_wgrad_accumulation=True,
    fuse_qkv_params=True, # Required for fuse_wgrad_accumulation
)
wgrad_transformer.to(dtype=dtype).cuda()
for param in wgrad_transformer.parameters():
    param.grad = None
    param.main_grad = torch.zeros_like(param, dtype=torch.float32)

# Training step
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = wgrad_transformer(x, attention_mask=None)
y.backward(dy)
for param in wgrad_transformer.parameters():
    if param.grad is not None:
        param.main_grad.copy_(param.grad)
        param.grad = None

# Measure step time
utils.speedometer(
    wgrad_transformer,
    x,
    dy,
    forward_kwargs = { "attention_mask": None },
    fp8_autocast_kwargs = { "enabled": True, "fp8_recipe": fp8_recipe },
)

Mean time: 27.510029296875 ms

FP8 权重缓存¶

摘要

当使用多个梯度累积步骤进行训练时，我们避免了冗余的 FP8 类型转换。

由于权重通常以 FP32 格式进行训练，因此在执行 FP8 计算之前需要进行类型转换。默认情况下，fp8_autocast 上下文管理器会在内部处理此问题，将遇到的非 FP8 张量转换为 FP8。但是，在某些情况下，我们可以对此进行改进。特别是，如果我们的训练迭代被分成多个梯度累积步骤，则每个小批量将遇到相同的权重张量。因此，我们只需要在第一个梯度累积步骤中将权重转换为 FP8，并且可以缓存生成的 FP8 权重以用于剩余的梯度累积步骤。

警告！

使用和不使用 FP8 权重缓存优化时的精确数值输出可能不是逐位相同的。这是因为虽然权重在梯度累积周期中保持冻结，但 FP8 权重的缩放因子和 amax 值会随着每次迭代结束时的更新而发生变化。 amax 张量的这些变化会合并到 amax 历史记录中，该历史记录未被冻结。

[6]:

# Construct layer
weight_caching_transformer = te.TransformerLayer(
    hidden_size,
    ffn_hidden_size,
    num_attention_heads,
)
weight_caching_transformer.to(dtype=dtype).cuda()

# Cast weights in first gradient accumulation step
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = weight_caching_transformer(x, attention_mask=None, is_first_microbatch=True)
y.backward(dy)

# Reuse FP8 weights in subsequent gradient accumulation steps
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = weight_caching_transformer(x, attention_mask=None, is_first_microbatch=False)
y.backward(dy)

# Measure step time
utils.speedometer(
    weight_caching_transformer,
    x,
    dy,
    forward_kwargs = { "attention_mask": None, "is_first_microbatch": False },
    fp8_autocast_kwargs = { "enabled": True, "fp8_recipe": fp8_recipe },
)

Mean time: 27.262666015625 ms