Initialize torch.distributed.
(args)
| 474 | |
| 475 | |
| 476 | def initialize_distributed(args): |
| 477 | """Initialize torch.distributed.""" |
| 478 | |
| 479 | # Manually set the device ids. |
| 480 | device = args.rank % torch.cuda.device_count() |
| 481 | if args.local_rank is not None: |
| 482 | device = args.local_rank |
| 483 | torch.cuda.set_device(device) |
| 484 | # Call the init process |
| 485 | init_method = 'tcp://' |
| 486 | args.master_ip = os.getenv('MASTER_ADDR', 'localhost') |
| 487 | args.master_port = os.getenv('MASTER_PORT', '6000') |
| 488 | init_method += args.master_ip + ':' + args.master_port |
| 489 | if hasattr(deepspeed, "init_distributed"): |
| 490 | deepspeed.init_distributed(dist_backend=args.distributed_backend) |
| 491 | else: |
| 492 | torch.distributed.init_process_group( |
| 493 | backend=args.distributed_backend, |
| 494 | world_size=args.world_size, rank=args.rank, |
| 495 | init_method=init_method) |
| 496 | |
| 497 | # Set the model-parallel / data-parallel communicators. |
| 498 | mpu.initialize_model_parallel(args.model_parallel_size) |
| 499 | |
| 500 | # Optional DeepSpeed Activation Checkpointing Features |
| 501 | # |
| 502 | if hasattr(args, "deepspeed") and args.deepspeed and args.deepspeed_activation_checkpointing: |
| 503 | set_deepspeed_activation_checkpointing(args) |
| 504 | |
| 505 | |
| 506 | def set_random_seed(seed): |
no test coverage detected