性能优化

本指南是快速入门指南中讨论的后续。我们将重点介绍在训练基本 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
torch.distributed.init_process_group(
    "nccl",
    init_method="file:///tmp/rdzv",
    world_size=1,
    rank=0,
)
world_group = torch.distributed.new_group(ranks=[0], backend="nccl")
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 解释为尺寸为 \(\text{num_nodes} \times \text{gpus_per_node}\) 的 2D 网格。行是张量并行组，列是数据并行组。

使用 Transformer Engine 启用数据并行类似于使用标准 PyTorch 模型启用数据并行：只需使用 torch.nn.parallel.DistributedDataParallel 包装模块即可。Transformer Engine 模块还原生支持张量并行和序列并行。如果用户为张量并行提供进程组，则模块将在内部分配数据并执行通信。如果启用了序列并行，则它将应用于不适合张量并行操作的操作，并将使用张量并行进程组。

多 GPU FP8 训练的一个重要考虑因素是如何在 GPU 之间同步 FP8 缩放因子。如果启用了张量并行，则必须在张量并行组上同步比例。但是，建议在数据并行组和张量并行组上都进行同步，以获得最佳收敛效果。这可以使用 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=world_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": world_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