importosRANK=int(os.environ["RANK"])WORLD_SIZE=int(os.environ["WORLD_SIZE"])os.environ["CUDA_VISIBLE_DEVICES"]=os.environ["LOCAL_RANK"]importtorchimporttorch.distributedasdisttorch.manual_seed(0)dist.init_process_group("gloo")N=10expected=torch.sum(torch.arange(0,WORLD_SIZE,dtype=torch.float)).item()t=torch.full((1000000,),RANK,device="cuda",dtype=torch.float)tensors= [t.clone()for_inrange(N)]futs= []fortensorintensors:work=dist.all_reduce(tensor,async_op=True)futs.append(work.get_future())# create high priority stream to do the CPU copy and preempt the default streamstream=torch.cuda.Stream(priority=-1)forfut,tensorinzip(futs,tensors):withtorch.cuda.stream(stream):fut.wait()val=tensor[-1].item()assertval==expected,f"Expected{expected}, got{val}"

torchrun --nnodes 1 --nproc_per_node=gpu ~/scripts/gloo_future_stream.py

BACKEND=gloo WORLD_SIZE=4 TEMP_DIR=/tmp/foo pytest test/distributed/test_distributed_spawn.py -v -s -x  -k 'test_ddp_apply_optim_in_backward'