Initialize torch.distributed and mpu.
()
| 207 | |
| 208 | |
| 209 | def _initialize_distributed(): |
| 210 | """Initialize torch.distributed and mpu.""" |
| 211 | args = get_args() |
| 212 | |
| 213 | device_count = torch.cuda.device_count() |
| 214 | if torch.distributed.is_initialized(): |
| 215 | |
| 216 | if args.rank == 0: |
| 217 | print( |
| 218 | "torch distributed is already initialized, " |
| 219 | "skipping initialization ...", |
| 220 | flush=True, |
| 221 | ) |
| 222 | args.rank = torch.distributed.get_rank() |
| 223 | args.world_size = torch.distributed.get_world_size() |
| 224 | |
| 225 | else: |
| 226 | |
| 227 | if args.rank == 0: |
| 228 | print("> initializing torch distributed ...", flush=True) |
| 229 | # Manually set the device ids. |
| 230 | if device_count > 0: |
| 231 | device = args.rank % device_count |
| 232 | if args.local_rank is not None: |
| 233 | assert ( |
| 234 | args.local_rank == device |
| 235 | ), "expected local-rank to be the same as rank % device-count." |
| 236 | else: |
| 237 | args.local_rank = device |
| 238 | if args.force_device is not None: |
| 239 | print( |
| 240 | f" > forcefully set the device to {args.force_device}, originally {device}" |
| 241 | ) |
| 242 | device = args.force_device |
| 243 | torch.cuda.set_device(device) |
| 244 | # Call the init process |
| 245 | init_method = "tcp://" |
| 246 | master_ip = os.getenv("MASTER_ADDR", "localhost") |
| 247 | master_port = os.getenv("MASTER_PORT", "6000") |
| 248 | init_method += master_ip + ":" + master_port |
| 249 | print( |
| 250 | f" > (rank={args.rank}) initializing process group: " |
| 251 | f"world_size={args.world_size} " |
| 252 | f"backend={args.distributed_backend} " |
| 253 | f"init_method={init_method}", |
| 254 | flush=True, |
| 255 | ) |
| 256 | timeout = datetime.timedelta(minutes=args.dist_timeout) |
| 257 | torch.distributed.init_process_group( |
| 258 | backend=args.distributed_backend, |
| 259 | world_size=args.world_size, |
| 260 | rank=args.rank, |
| 261 | init_method=init_method, |
| 262 | timeout=timeout |
| 263 | ) |
| 264 | print(f" > (rank={args.rank}) process group initialized") |
| 265 | |
| 266 | # Set the tensor model-parallel, pipeline model-parallel, and |
no test coverage detected