(context_parallel_size)
| 30 | |
| 31 | |
| 32 | def initialize_context_parallel(context_parallel_size): |
| 33 | global _CONTEXT_PARALLEL_GROUP |
| 34 | global _CONTEXT_PARALLEL_SIZE |
| 35 | |
| 36 | assert _CONTEXT_PARALLEL_GROUP is None, "context parallel group is already initialized" |
| 37 | _CONTEXT_PARALLEL_SIZE = context_parallel_size |
| 38 | |
| 39 | rank = torch.distributed.get_rank() |
| 40 | world_size = torch.distributed.get_world_size() |
| 41 | |
| 42 | for i in range(0, world_size, context_parallel_size): |
| 43 | ranks = range(i, i + context_parallel_size) |
| 44 | group = torch.distributed.new_group(ranks) |
| 45 | if rank in ranks: |
| 46 | _CONTEXT_PARALLEL_GROUP = group |
| 47 | break |
| 48 | |
| 49 | |
| 50 | def get_context_parallel_group(): |
no outgoing calls
no test coverage detected