| 50 | static bool all_ranks_local_p = false; |
| 51 | |
| 52 | void initialize(int size, int rank, torch::Tensor& kvs_data) |
| 53 | { |
| 54 | if (is_initialized) return; |
| 55 | |
| 56 | // Check whether all ranks is on the same physical machine. |
| 57 | // If true, we will use an SHM based low latency allreduce |
| 58 | |
| 59 | auto ls_string = std::getenv("LOCAL_SIZE"); |
| 60 | int ls = 0; |
| 61 | if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } |
| 62 | |
| 63 | if (size >= 1 && size == ls) { all_ranks_local_p = true; } |
| 64 | |
| 65 | world_size = size; |
| 66 | world_rank = rank; |
| 67 | is_initialized = 1; |
| 68 | |
| 69 | ccl::kvs::address_type main_addr; |
| 70 | |
| 71 | if (rank != 0) { |
| 72 | memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); |
| 73 | kvs = ccl::create_kvs(main_addr); |
| 74 | } |
| 75 | |
| 76 | _ccl_comms.emplace_back(ccl::create_communicator(size, rank, kvs)); |
| 77 | |
| 78 | auto addr_string = std::getenv("MASTER_ADDR"); |
| 79 | if (addr_string == NULL) { addr_string = ""; } |
| 80 | auto port_string = std::getenv("MASTER_PORT"); |
| 81 | if (port_string == NULL) { port_string = ""; } |
| 82 | |
| 83 | if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } |
| 84 | } |
| 85 | |
| 86 | /* |
| 87 | rank == 0: create main kvs and return its address |
nothing calls this directly
no test coverage detected