Parameters do not properly sync across processes when combining DDP and AMP with opt level O2

With my current setup, it appears that model parameters are not synchronized across processes when running with DDP and AMP set to opt level O2. However, opt level O1 seems to behave as expected -- i.e., all processes appear to be working with the same version of the model.

It's possible that this is intended behavior or that I am doing something wrong, but I did not expect it to work this way and it causes a variety of problems, from subtle to glaring. If there's some incorrect usage of AMP or DDP here, please let me know.

I've written a script that reproduces this behavior in my setup and that of a co-worker with a slightly different dev environment -- i.e., no explicit matching of CUDA, apex, or pytorch versions. The idea is to make a minimal model, then run a forward pass through the model version managed by process 0 to get a baseline output, then run forward passes through the model versions managed by all processes, and then compare those outputs with the baseline output from process 0. All of the forward passes use the same input.

The script can be run like:

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 amp_bug_script.py --opt_level O1

When I run this script with --opt_level O1, it behaves as expected and the diff between outputs from all processes is 0. When I run it with --opt_level O2, the diff between process 0 and the baseline is 0, since it provides the baseline, but all other processes have non-trivial diff from the baseline.

Here's the script:

import os
import time
import pickle
import argparse

import torch

from apex import amp


def test_opt_level(opt_level, local_rank):
    # initialize distributed/multiprocessing stuff
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '8787'
    torch.cuda.set_device(local_rank)
    torch.distributed.init_process_group(backend='nccl',
                                         init_method='env://')

    # get a basic model and optimizer
    n_input, n_output = 128, 128
    model = torch.nn.Sequential(
        torch.nn.Linear(n_input, 512),
        torch.nn.ReLU(),
        torch.nn.Linear(512, n_output)
    )
    model.to('cuda')
    optimizer = torch.optim.Adam(model.parameters())

    # wrap model and optimizer with AMP stuff
    model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
    # wrap model with DDP stuff
    model = torch.nn.parallel.DistributedDataParallel(model,
                                device_ids=[local_rank],
                                output_device=local_rank,
                                find_unused_parameters=True)
    
    # run a baseline pass through model for comparison across processes
    if local_rank == 0:
        # run an input through model to get a baseline output
        shared_input = torch.randn((16, n_input))
        with torch.no_grad():
            shared_output = model(shared_input.to('cuda'))
        probe_dict = {'input': shared_input.cpu(),
                      'output': shared_output.cpu()}
        # save shared input and output to pkl
        pkl_file = open('probe_dict.pkl', 'wb')
        pickle.dump(probe_dict, pkl_file)
        pkl_file.close()
        print('--- Measuring diffs against process {0:d} --- '.format(local_rank))
    torch.distributed.barrier()  # sync processes before moving on

    # load probe dict shared across processes
    pkl_file = open('probe_dict.pkl', 'rb')
    probe_dict = pickle.load(pkl_file)
    pkl_file.close()
    # compute output for this process on shared input
    with torch.no_grad():
        local_output = model(probe_dict['input'].to('cuda'))
    # compare local process output with shared baseline
    out_diff = torch.abs(local_output.cpu() - probe_dict['output']).mean()
    torch.distributed.barrier()  # sync processes before moving on
    # print output diffs
    time.sleep(0.1 + float(local_rank))
    print('  process {0:d}, out_diff={1:.4f}'.format(local_rank, out_diff))


if __name__=='__main__':
    # parser with minimal args for current purposes
    parser = argparse.ArgumentParser(description='arg parser')
    parser.add_argument("--opt_level", default='O1', type=str)
    parser.add_argument("--local_rank", default=0, type=int)
    args = parser.parse_args()
    # run test for a given opt level
    test_opt_level(args.opt_level, args.local_rank)

Start with the attached amp_bug_script.py and run it with opt_level O1 and O2 using the shown torch.distributed.launch command. Read the amp.initialize and DistributedDataParallel setup in the script, then compare each process's out_diff against process 0. Done means determining whether O2's differing outputs are intended or identifying the cause of the synchronization failure.