(args)
| 82 | |
| 83 | |
| 84 | def init_communicators(args): |
| 85 | default_init(args) |
| 86 | assert args.world_size == args.data_group_size * args.pipeline_group_size |
| 87 | if args.world_size == args.data_group_size * args.pipeline_group_size: |
| 88 | # We do the following hard code alignment of communication groups: |
| 89 | # Suppose there are 8 instances (world_size), and 4 data parallel groups (data_group_size is 2), |
| 90 | # Then there would be 2 pipeline parallel groups (pipeline_group_size is 4), then the groups will look like: |
| 91 | # pipeline parallel: <group 0: [0,1,2,3]>, <group 1: [4,5,6,7]> |
| 92 | # data parallel: <group 0: [0,4]>, <group 1: [1,5]>, <group 2: [2,6]>, <group 3: [3,7]> |
| 93 | # assert args.world_size == args.data_group_size * args.pipeline_group_size |
| 94 | global _DATA_PARALLEL_COMM |
| 95 | global _PIPELINE_PARALLEL_COMM |
| 96 | global _DATA_PARALLEL_RANK |
| 97 | global _PIPELINE_PARALLEL_RANK |
| 98 | global _DATA_PARALLEL_WORLD_SIZE |
| 99 | global _PIPELINE_PARALLEL_WORLD_SIZE |
| 100 | # We use pipeline parallel by default. |
| 101 | _PIPELINE_PARALLEL_WORLD_SIZE = args.pipeline_group_size |
| 102 | _PIPELINE_PARALLEL_RANK = args.rank % args.pipeline_group_size |
| 103 | _PIPELINE_PARALLEL_COMM = NCCLCommunicator(_PIPELINE_PARALLEL_RANK, args.cuda_id, args.pipeline_group_size, |
| 104 | "pipeline_group_"+str(args.rank // args.pipeline_group_size)) |
| 105 | if args.data_group_size != 1: |
| 106 | _DATA_PARALLEL_WORLD_SIZE = args.data_group_size |
| 107 | _DATA_PARALLEL_RANK = args.rank // args.pipeline_group_size |
| 108 | |
| 109 | dp_backend = getattr(args, 'dp_backend', 'gloo') |
| 110 | if dp_backend == 'nccl': |
| 111 | |
| 112 | _DATA_PARALLEL_COMM = NCCLCommunicator(_DATA_PARALLEL_RANK, args.cuda_id, args.data_group_size, |
| 113 | "data_group_"+str(args.rank % args.pipeline_group_size)) |
| 114 | |
| 115 | elif dp_backend == 'gloo': |
| 116 | |
| 117 | for i in range(args.pipeline_group_size): |
| 118 | ranks = [rank for rank in range(i, args.world_size, args.pipeline_group_size)] |
| 119 | print(args.rank, ranks) |
| 120 | data_group = torch.distributed.new_group(ranks, backend='gloo') |
| 121 | if args.rank in ranks: |
| 122 | def to_global_rank(dp_rank): |
| 123 | rank = _PIPELINE_PARALLEL_RANK + dp_rank * args.pipeline_group_size |
| 124 | # print(f"{dp_rank} --> {rank}") |
| 125 | return rank |
| 126 | _DATA_PARALLEL_COMM = TorchCommunicator( |
| 127 | data_group, to_global_rank=to_global_rank, |
| 128 | dp_rank=_DATA_PARALLEL_RANK, |
| 129 | comm_group_size=args.data_group_size,) |
| 130 | |
| 131 | else: |
| 132 | assert False |
| 133 | |
| 134 | print('comm init done!!') |
| 135 | |
| 136 | # elif args.world_size == args.data_group_size * args.tensor_group_size: |
| 137 | # global _DATA_PARALLEL_COMM |
| 138 | # global _TENSOR_PARALLEL_COMM |
| 139 | # global _DATA_PARALLEL_RANK |
| 140 | # global _TENSOR_PARALLEL_RANK |
| 141 | # global _DATA_PARALLEL_WORLD_SIZE |
no test coverage detected