使用ncclReduce api的示例

以下是使用ncclReduce API的示例代码：

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

def run(rank, size):
    # 创建进程组
    dist.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=rank, world_size=size)
    
    # 创建模型和优化器
    model = nn.Linear(10, 1).to(rank)
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    # 创建输入数据
    input_data = torch.randn(10).to(rank)
    
    # 同步所有进程的模型参数
    dist.broadcast(model.state_dict(), src=0)
    
    # 前向传播
    output = model(input_data)
    
    # 计算损失
    loss = ((output - 1) ** 2).sum()
    
    # 梯度清零
    optimizer.zero_grad()
    
    # 反向传播
    loss.backward()
    
    # 使用ncclReduce API 同步梯度
    dist.reduce(model.grad, dst=0, op=dist.ReduceOp.SUM)
    
    # 更新模型参数
    optimizer.step()
    
    # 打印输出
    if rank == 0:
        print('Rank 0: loss={:.4f}, gradient={:.4f}'.format(loss.item(), model.grad.item()))

def main():
    # 设置进程数量
    size = 2
    
    # 启动多进程
    mp.spawn(run, args=(size,), nprocs=size)

if __name__ == '__main__':
    main()

在这个示例中，我们使用了torch.distributed模块中的dist.reduce函数来同步梯度。在这个函数中，我们指定了源进程（src=0）和操作（op=dist.ReduceOp.SUM）来将所有进程的梯度相加，并将结果广播给所有进程。这样，所有进程都可以使用相同的梯度更新模型参数。

在调用dist.reduce函数之前，我们先使用dist.broadcast函数将模型参数广播给所有进程，以确保所有进程使用的是相同的初始参数。

最后，我们在rank为0的进程中打印出损失和梯度的值，以验证同步是否成功。

请注意，该示例假设已经按照指定的初始化方法（init_method='tcp://localhost:23456'）启动了分布式进程组。您需要根据自己的环境设置正确的初始化方法。此外，还需要保证您的环境支持NCCL库，并且已经按照正确的方式配置了NCCL库