使用 ZeroRedundancyOptimizer 优化器状态分片#
创建日期：2021 年 2 月 26 日 | 最后更新：2021 年 10 月 20 日 | 最后验证：未验证
在本实践中，您将学习
ZeroRedundancyOptimizer 的高级思想。
如何在分布式训练中使用 ZeroRedundancyOptimizer 及其影响。
要求#
PyTorch 1.8+
什么是 ZeroRedundancyOptimizer？#
ZeroRedundancyOptimizer 的思想源自 DeepSpeed/ZeRO 项目 和 Marian，它将优化器状态分片到分布式数据并行进程中，以减少每个进程的内存占用。在 开始使用分布式数据并行 教程中，我们展示了如何使用 DistributedDataParallel (DDP) 来训练模型。在该教程中，每个进程都维护一个专用的优化器副本。由于 DDP 在反向传播时已经同步了梯度，所有优化器副本在每次迭代中都会操作相同的参数和梯度值，这就是 DDP 使模型副本保持相同状态的方式。通常，优化器还维护本地状态。例如，Adam 优化器使用每个参数的 exp_avg 和 exp_avg_sq 状态。因此，Adam 优化器的内存消耗至少是模型大小的两倍。基于此观察，我们可以通过分片优化器状态到 DDP 进程来减少优化器的内存占用。更具体地说，每个 DDP 进程中的优化器实例不再为所有参数创建每个参数状态，而是仅为所有模型参数的一个分片保留优化器状态。优化器 step() 函数仅更新其分片中的参数，然后将其更新后的参数广播到所有其他对等 DDP 进程，从而使所有模型副本保持相同的状态。
如何使用 ZeroRedundancyOptimizer？#
以下代码演示了如何使用 ZeroRedundancyOptimizer。大部分代码与 分布式数据并行说明 中提出的简单 DDP 示例相似。主要区别在于 example 函数中的 if-else 子句，该子句包装了优化器构造，在 ZeroRedundancyOptimizer 和 Adam 优化器之间进行切换。
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.distributed.optim import ZeroRedundancyOptimizer
from torch.nn.parallel import DistributedDataParallel as DDP
def print_peak_memory(prefix, device):
if device == 0:
print(f"{prefix}: {torch.cuda.max_memory_allocated(device) // 1e6}MB ")
def example(rank, world_size, use_zero):
torch.manual_seed(0)
torch.cuda.manual_seed(0)
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
# create default process group
dist.init_process_group("gloo", rank=rank, world_size=world_size)
# create local model
model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
print_peak_memory("Max memory allocated after creating local model", rank)
# construct DDP model
ddp_model = DDP(model, device_ids=[rank])
print_peak_memory("Max memory allocated after creating DDP", rank)
# define loss function and optimizer
loss_fn = nn.MSELoss()
if use_zero:
optimizer = ZeroRedundancyOptimizer(
ddp_model.parameters(),
optimizer_class=torch.optim.Adam,
lr=0.01
)
else:
optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.01)
# forward pass
outputs = ddp_model(torch.randn(20, 2000).to(rank))
labels = torch.randn(20, 2000).to(rank)
# backward pass
loss_fn(outputs, labels).backward()
# update parameters
print_peak_memory("Max memory allocated before optimizer step()", rank)
optimizer.step()
print_peak_memory("Max memory allocated after optimizer step()", rank)
print(f"params sum is: {sum(model.parameters()).sum()}")
def main():
world_size = 2
print("=== Using ZeroRedundancyOptimizer ===")
mp.spawn(example,
args=(world_size, True),
nprocs=world_size,
join=True)
print("=== Not Using ZeroRedundancyOptimizer ===")
mp.spawn(example,
args=(world_size, False),
nprocs=world_size,
join=True)
if __name__=="__main__":
main()
输出如下所示。当启用 ZeroRedundancyOptimizer 和 Adam 时，优化器 step() 的峰值内存消耗是普通 Adam 内存消耗的一半。这与我们的预期一致，因为我们将 Adam 优化器状态分片到了两个进程中。输出还表明，使用 ZeroRedundancyOptimizer 时，一个迭代后模型参数的值仍然相同（参数的总和与未使用 ZeroRedundancyOptimizer 时的总和相同）。
=== Using ZeroRedundancyOptimizer ===
Max memory allocated after creating local model: 335.0MB
Max memory allocated after creating DDP: 656.0MB
Max memory allocated before optimizer step(): 992.0MB
Max memory allocated after optimizer step(): 1361.0MB
params sum is: -3453.6123046875
params sum is: -3453.6123046875
=== Not Using ZeroRedundancyOptimizer ===
Max memory allocated after creating local model: 335.0MB
Max memory allocated after creating DDP: 656.0MB
Max memory allocated before optimizer step(): 992.0MB
Max memory allocated after optimizer step(): 1697.0MB
params sum is: -3453.6123046875
params sum is: -3453.6123046875
