| 23 | |
| 24 | |
| 25 | def initialize_context_parallel(context_parallel_size): |
| 26 | global _CONTEXT_PARALLEL_GROUP |
| 27 | global _CONTEXT_PARALLEL_SIZE |
| 28 | |
| 29 | assert _CONTEXT_PARALLEL_GROUP is None, "context parallel group is already initialized" |
| 30 | _CONTEXT_PARALLEL_SIZE = context_parallel_size |
| 31 | |
| 32 | rank = torch.distributed.get_rank() |
| 33 | world_size = torch.distributed.get_world_size() |
| 34 | |
| 35 | for i in range(0, world_size, context_parallel_size): |
| 36 | ranks = range(i, i + context_parallel_size) |
| 37 | group = torch.distributed.new_group(ranks) |
| 38 | if rank in ranks: |
| 39 | _CONTEXT_PARALLEL_GROUP = group |
| 40 | break |
| 41 | |
| 42 | |
| 43 | def get_context_parallel_group(): |