分布式并行训练 - FSDP

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def fsdp_main(rank, world_size, args):
    # 初始化分布式训练环境
    setup(rank, world_size)

    # 定义数据预处理变换
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将PIL图像转换为Tensor
        transforms.Normalize((0.1307,), (0.3081,))  # 标准化（MNIST数据集的均值和标准差）
    ])

    # 加载MNIST训练集和测试集
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                        transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                        transform=transform)

    # 为分布式训练创建采样器
    # 训练集采样器，会打乱数据
    sampler1 = DistributedSampler(dataset1, rank=rank, num_replicas=world_size, shuffle=True)
    # 测试集采样器，不打乱数据
    sampler2 = DistributedSampler(dataset2, rank=rank, num_replicas=world_size)

    # 配置训练和测试的数据加载器参数
    train_kwargs = {'batch_size': args.batch_size, 'sampler': sampler1}
    test_kwargs = {'batch_size': args.test_batch_size, 'sampler': sampler2}
    # CUDA相关的数据加载配置
    cuda_kwargs = {'num_workers': 2,  # 数据加载的线程数
                    'pin_memory': True,  # 将数据固定在内存中加速传输到GPU
                    'shuffle': False}  # 采样器已经处理了shuffle，所以这里设为False
    # 合并参数
    train_kwargs.update(cuda_kwargs)
    test_kwargs.update(cuda_kwargs)

    # 创建训练和测试数据加载器
    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
    
    # 定义FSDP的自动包装策略 - 当模块参数大于100时进行分片
    my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=100
    )
    
    # 设置当前使用的GPU设备
    torch.cuda.set_device(rank)

    # 创建 CUDA 事件用于计时
    init_start_event = torch.cuda.Event(enable_timing=True)
    init_end_event = torch.cuda.Event(enable_timing=True)

    # 初始化模型并移动到当前GPU
    model = Net().to(rank)
    # 使用FSDP包装模型
    model = FSDP(model)

    # 定义优化器（Adadelta）和学习率调度器
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    
    # 开始计时
    init_start_event.record()
    
    # 训练循环
    for epoch in range(1, args.epochs + 1):
        train(args, model, rank, world_size, train_loader, optimizer, epoch, sampler=sampler1)
        test(model, rank, world_size, test_loader)
        scheduler.step()  # 更新学习率

    # 结束计时
    init_end_event.record()

    # 主进程（rank 0）打印训练时间和模型信息
    if rank == 0:
        init_end_event.synchronize()
        print(f"CUDA event elapsed time: {init_start_event.elapsed_time(init_end_event) / 1000}sec")
        print(f"{model}")

    # 如果需要保存模型
    if args.save_model:
        # 使用屏障确保所有进程完成训练
        dist.barrier()
        # 获取模型状态字典
        states = model.state_dict()
        # 只有主进程保存模型
        if rank == 0:
            torch.save(states, "mnist_cnn.pt")

    # 清理分布式训练环境
    cleanup()

if __name__ == '__main__':
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()

    torch.manual_seed(args.seed)

    WORLD_SIZE = torch.cuda.device_count()
    # 使用多进程启动分布式训练
    mp.spawn(fsdp_main,        # 主训练函数
            args=(WORLD_SIZE, args),  # 传递给主训练函数的参数
            nprocs=WORLD_SIZE, # 启动的进程数（等于GPU数量）
            join=True)        # 等待所有进程完成

   FullyShardedDataParallel(
     (_fsdp_wrapped_module): FlattenParamsWrapper(			# FSDP 内部的一个辅助模块，用于扁平化参数
       (_fpw_module): Net(
       (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
       (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
       (dropout1): Dropout(p=0.25, inplace=False)
       (dropout2): Dropout(p=0.5, inplace=False)
       (fc1): Linear(in_features=9216, out_features=128, bias=True)
       (fc2): Linear(in_features=128, out_features=10, bias=True)
       )
   )
)

my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=20000
    )
torch.cuda.set_device(rank)
model = Net().to(rank)

model = FSDP(model,
    auto_wrap_policy=my_auto_wrap_policy)

  FullyShardedDataParallel(
(_fsdp_wrapped_module): FlattenParamsWrapper(
  (_fpw_module): Net(
    (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
    (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
    (dropout1): Dropout(p=0.25, inplace=False)
    (dropout2): Dropout(p=0.5, inplace=False)
    (fc1): FullyShardedDataParallel(			# fc1 嵌套封装 FSDP
      (_fsdp_wrapped_module): FlattenParamsWrapper(
        (_fpw_module): Linear(in_features=9216, out_features=128, bias=True)
      )
    )
    (fc2): Linear(in_features=128, out_features=10, bias=True)
  )
)