def _test_collect_states(rank, world_size): os.environ["MASTER_ADDR"] = "localhost" torch.cuda.set_device(f"cuda:{rank}") # initialize the process group torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size) state = {"something": torch.tensor([rank])} collected_state = _collect_states_on_rank_zero(state) assert collected_state == {1: {"something": torch.tensor([1])}, 0: {"something": torch.tensor([0])}}
def fn(state: Dict): if key in state: return _collect_states_on_rank_zero(state) return {k: apply_to_collection(v, Dict, fn) for k, v in state.items()}