Initialize torch.distributed.
(args)
| 506 | |
| 507 | |
| 508 | def initialize_distributed(args): |
| 509 | """Initialize torch.distributed.""" |
| 510 | if torch.distributed.is_initialized(): |
| 511 | if mpu.model_parallel_is_initialized(): |
| 512 | if args.model_parallel_size != mpu.get_model_parallel_world_size(): |
| 513 | raise ValueError('model_parallel_size is inconsistent with prior configuration.' |
| 514 | 'We currently do not support changing model_parallel_size.') |
| 515 | return False |
| 516 | else: |
| 517 | if args.model_parallel_size > 1: |
| 518 | warnings.warn('model_parallel_size > 1 but torch.distributed is not initialized via SAT.' |
| 519 | 'Please carefully make sure the correctness on your own.') |
| 520 | mpu.initialize_model_parallel(args.model_parallel_size) |
| 521 | return True |
| 522 | # the automatic assignment of devices has been moved to arguments.py |
| 523 | if args.device == 'cpu': |
| 524 | pass |
| 525 | else: |
| 526 | torch.cuda.set_device(args.device) |
| 527 | # Call the init process |
| 528 | init_method = 'tcp://' |
| 529 | args.master_ip = os.getenv('MASTER_ADDR', 'localhost') |
| 530 | |
| 531 | if args.world_size == 1: |
| 532 | from sat.helpers import get_free_port |
| 533 | default_master_port = str(get_free_port()) |
| 534 | else: |
| 535 | default_master_port = '6000' |
| 536 | args.master_port = os.getenv('MASTER_PORT', default_master_port) |
| 537 | init_method += args.master_ip + ':' + args.master_port |
| 538 | torch.distributed.init_process_group( |
| 539 | backend=args.distributed_backend, |
| 540 | world_size=args.world_size, rank=args.rank, |
| 541 | init_method=init_method) |
| 542 | |
| 543 | # Set the model-parallel / data-parallel communicators. |
| 544 | mpu.initialize_model_parallel(args.model_parallel_size) |
| 545 | # Optional DeepSpeed Activation Checkpointing Features |
| 546 | if args.deepspeed: |
| 547 | import deepspeed |
| 548 | deepspeed.init_distributed( |
| 549 | dist_backend=args.distributed_backend, |
| 550 | world_size=args.world_size, rank=args.rank, init_method=init_method) |
| 551 | # It seems that it has no negative influence to configure it even without using checkpointing. |
| 552 | deepspeed.checkpointing.configure(mpu, deepspeed_config=args.deepspeed_config, num_checkpoints=args.num_layers) |
| 553 | else: |
| 554 | # in model-only mode, we don't want to init deepspeed, but we still need to init the rng tracker for model_parallel, just because we save the seed by default when dropout. |
| 555 | try: |
| 556 | import deepspeed |
| 557 | from deepspeed.runtime.activation_checkpointing.checkpointing import _CUDA_RNG_STATE_TRACKER, _MODEL_PARALLEL_RNG_TRACKER_NAME |
| 558 | _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, 1) # default seed 1 |
| 559 | except Exception as e: |
| 560 | from sat.helpers import print_rank0 |
| 561 | print_rank0(str(e), level="DEBUG") |
| 562 | |
| 563 | |
| 564 | return True |
| 565 |
no test coverage detected