Initialize model data parallel groups. Arguments: model_parallel_size: number of GPUs used to parallelize model. Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we use 2 GPUs to parallelize the model. The present function will create 4 model parallel group
(model_parallel_size_)
| 32 | import os |
| 33 | |
| 34 | def initialize_model_parallel(model_parallel_size_): |
| 35 | """ |
| 36 | Initialize model data parallel groups. |
| 37 | |
| 38 | Arguments: |
| 39 | model_parallel_size: number of GPUs used to parallelize model. |
| 40 | |
| 41 | Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we |
| 42 | use 2 GPUs to parallelize the model. The present function will |
| 43 | create 4 model parallel groups and 2 data parallel groups as: |
| 44 | 4 model parallel groups: |
| 45 | [g0, g1], [g2, g3], [g4, g5], [g6, g7] |
| 46 | 2 data parallel groups: |
| 47 | [g0, g2, g4, g6], [g1, g3, g5, g7] |
| 48 | Note that for efficiency, the caller should make sure adjacent ranks |
| 49 | are on the same DGX box. For example if we are using 2 DGX-1 boxes |
| 50 | with a total of 16 GPUs, rank 0 to 7 belong to the first box and |
| 51 | ranks 8 to 15 belong to the second box. |
| 52 | """ |
| 53 | if torch.distributed.get_rank() == 0: |
| 54 | print_rank0('> initializing model parallel with size {}'.format( |
| 55 | model_parallel_size_)) |
| 56 | # Get world size and rank. Ensure some consistencies. |
| 57 | assert torch.distributed.is_initialized() |
| 58 | world_size = torch.distributed.get_world_size() |
| 59 | model_parallel_size = min(model_parallel_size_, world_size) |
| 60 | ensure_divisibility(world_size, model_parallel_size) |
| 61 | rank = torch.distributed.get_rank() |
| 62 | |
| 63 | # Build the data parallel groups. |
| 64 | global _DATA_PARALLEL_GROUP |
| 65 | assert _DATA_PARALLEL_GROUP is None, \ |
| 66 | 'data parallel group is already initialized' |
| 67 | for i in range(model_parallel_size): |
| 68 | ranks = range(i, world_size, model_parallel_size) |
| 69 | group = torch.distributed.new_group(ranks) |
| 70 | if i == (rank % model_parallel_size): |
| 71 | _DATA_PARALLEL_GROUP = group |
| 72 | |
| 73 | # Build the model parallel groups. |
| 74 | global _MODEL_PARALLEL_GROUP |
| 75 | assert _MODEL_PARALLEL_GROUP is None, \ |
| 76 | 'model parallel group is already initialized' |
| 77 | for i in range(world_size // model_parallel_size): |
| 78 | ranks = range(i * model_parallel_size, |
| 79 | (i + 1) * model_parallel_size) |
| 80 | group = torch.distributed.new_group(ranks) |
| 81 | if i == (rank // model_parallel_size): |
| 82 | _MODEL_PARALLEL_GROUP = group |
| 83 | |
| 84 | guess_local_world_size = world_size if world_size < 8 else 8 |
| 85 | local_world_size = os.environ.get('LOCAL_WORLD_SIZE', None) |
| 86 | if local_world_size is None: |
| 87 | local_world_size = guess_local_world_size |
| 88 | print_rank0(f"You didn't pass in LOCAL_WORLD_SIZE environment variable. We use the guessed LOCAL_WORLD_SIZE={guess_local_world_size}. If this is wrong, please pass the LOCAL_WORLD_SIZE manually.") |
| 89 | local_world_size = int(local_world_size) |
| 90 | # Build the node groups. |
| 91 | global _NODE_GROUP |
no test coverage detected