if distribute and local_rank == 0:
    acc = verification(...)
    # lr_sched is torch.optim.lr_scheduler.ReduceLROnPlateau()
    lr_sched.step(acc)
lr_broadcast = torch.tensor([0.0]).cuda()
if distribute:
    if local_rank == 0:
        lr_broadcast = torch.tensor([optimizer.state_dict()['param_groups'][0]['lr']]).cuda()
    dist.all_reduce(lr_broadcast)
    optimizer.state_dict()['param_groups'][0]['lr'] = lr_broadcast.item()
    print(local_rank, optimizer.state_dict()['param_groups'][0]['lr'])

When I run this code with 2 GPUs, I found two mistakes.
First, dist.all_reduct() failed. I try to print lr_broadcast after all_reduce, but different on GPU0 and GPU1.
Second, after I pass the lr_broadcast value to optimizer.state_dict, I print the lr in optimizer, but the value is differ from lr_broadcast.
So, I got confused.