def cleanup():
       dist.destroy_process_group()
def run_demo(demo_fn, world_size):
       mp.spawn(
           demo_fn,
           args=(world_size,),
           nprocs=world_size,
           join=True)
def demo_serialization(rank, world_size):
       setup(rank, world_size)
   CHECKPOINT = tempfile.gettempdir() + "/checkpoint.pt"

   model = SimpleModel().to(rank)
   ddp_model = DistributedDataParallel(model, device_ids=[rank])

   powersgd_hook = powerSGD.powerSGD_hook
   powersgd_state = powerSGD.PowerSGDState(process_group=None)

   optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
   ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

   state = {
       'state_dict': ddp_model.state_dict(),
       'comm_hook': hook,
       'comm_hook_state': hook_state}

   if rank == 0:
       torch.save(state, CHECKPOINT)

   dist.barrier()
   map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
