LlamaFactory

DDP 通信钩子 ​

分布式数据并行 (DDP) 通信钩子提供了一个通用接口，通过覆盖 DistributedDataParallel 中的默认全归约（allreduce）来控制梯度在工作节点之间的通信。提供了一些内置的通信钩子，用户可以轻松应用这些钩子来优化通信。

  • FP16 压缩钩子：通过将梯度转换为半精度浮点格式（torch.float16）来压缩梯度，减少通信开销。
  • BF16 压缩钩子：类似于 FP16，但使用脑浮点格式（torch.bfloat16），在某些硬件上可能更高效。
  • PowerSGD 钩子：一种高级梯度压缩算法，提供高压缩率，可以加速带宽受限的分布式训练。

在本教程中，你将看到如何快速设置 DDP 通信钩子，并使用 Accelerate 库提供的工具进行训练，这可能只需要添加一行代码！这将展示如何使用 DDP 通信钩子来优化分布式训练中的梯度通信。

FP16 压缩钩子 ​

BF16 压缩 Hook ​

在深度学习中，使用 BF16（Brain Floating Point 16）数据类型可以显著减少内存使用和加速计算。BF16 是一种 16 位浮点数格式，专为机器学习设计，具有与 32 位浮点数（FP32）相似的动态范围，但精度较低。

使用 BF16 压缩 Hook ​

为了在训练过程中使用 BF16，可以使用 BF16CompressionHook。这个 Hook 会在前向传播和反向传播中自动将张量转换为 BF16 格式，从而减少内存使用和加速计算。

示例代码 ​
python
from torch.distributed.algorithms.ddp_comm_hooks import BF16CompressionHook

# 假设你已经初始化了一个 DistributedDataParallel 模型
model = torch.nn.parallel.DistributedDataParallel(model)

# 注册 BF16 压缩 Hook
BF16CompressionHook.register(model)
注意事项 ​
  • 兼容性：确保你的硬件和软件环境支持 BF16。大多数现代 GPU 和 CPU 都支持 BF16。
  • 精度损失：虽然 BF16 可以显著加速计算，但可能会导致一些精度损失。在某些任务中，这种损失是可以接受的，但在其他任务中可能需要权衡。
  • 调试：在使用 BF16 时，建议进行充分的调试和验证，以确保模型的性能和准确性。
性能提升 ​

使用 BF16 压缩 Hook 可以显著减少内存使用和加速计算，特别是在大规模分布式训练中。以下是一些性能提升的示例：

  • 内存使用：BF16 只占用 16 位，而 FP32 占用 32 位，因此可以减少一半的内存使用。
  • 计算速度：现代 GPU 和 CPU 对 BF16 的支持使得计算速度显著提升。
结论 ​

BF16CompressionHook 是一个强大的工具，可以在深度学习训练中显著减少内存使用和加速计算。通过合理使用，可以在保持模型性能的同时，提高训练效率。

PowerSGD 钩子 ​

DDP 通信钩子工具 ​

有两个额外的工具用于支持通信钩子的可选功能。

comm_wrapper ​

comm_wrapper 是一个选项，用于将通信钩子包装以添加额外的功能。例如，它可以用于将 FP16 压缩与其他通信策略结合使用。目前支持的包装器有 no、fp16 和 bf16。

python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_wrapper=DDPCommunicationHookType.FP16
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

comm_state_option ​

comm_state_option 允许你传递某些通信钩子所需的额外状态信息。这在处理像 PowerSGD 这样的有状态钩子时特别有用，这些钩子需要在训练步骤之间维护超参数和内部状态。以下是一个使用 comm_state_option 和 PowerSGD 钩子的示例。

python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_state_option={"matrix_approximation_rank": 2}
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

对于更高级的用法和额外的钩子，请参阅 PyTorch DDP 通信钩子文档。

实用工具

Copyright © 2024-2025 llamafactory.cn