Necessary initialization for torch.distributed for model-only mode
(model_parallel_size=1, seed=0)
| 314 | before, after - before, after)) |
| 315 | |
| 316 | def _simple_init(model_parallel_size=1, seed=0): |
| 317 | '''Necessary initialization for torch.distributed for model-only mode''' |
| 318 | args = argparse.Namespace( |
| 319 | distributed_backend='nccl' if torch.distributed.is_nccl_available() and torch.cuda.is_available() else 'gloo', |
| 320 | model_parallel_size=model_parallel_size, |
| 321 | ) |
| 322 | args.rank = int(os.getenv('RANK', '0')) |
| 323 | args.world_size = int(os.getenv("WORLD_SIZE", '1')) |
| 324 | args.local_rank = int(os.getenv("LOCAL_RANK", '0')) # torchrun |
| 325 | args.device = args.local_rank |
| 326 | if not torch.cuda.is_available(): |
| 327 | args.device = 'cpu' |
| 328 | args.deepspeed = False |
| 329 | set_random_seed(seed) |
| 330 | |
| 331 | if initialize_distributed(args): # first time init model parallel, print warning |
| 332 | print_rank0( |
| 333 | 'You are using model-only mode.\n\ |
| 334 | For torch.distributed users or loading model parallel models, set environment variables RANK, WORLD_SIZE and LOCAL_RANK.' |
| 335 | ) |
| 336 | return True |
| 337 | return False |
| 338 | |
| 339 | def get_args(args_list=None, parser=None): |
| 340 | """Parse all the args.""" |
no test coverage detected