| from typing import List |
| from torch import distributed |
|
|
|
|
| def barrier(): |
| if distributed.is_initialized(): |
| distributed.barrier() |
| else: |
| pass |
|
|
|
|
| def broadcast(data, src): |
| if distributed.is_initialized(): |
| distributed.broadcast(data, src) |
| else: |
| pass |
|
|
|
|
| def all_gather(data: List, src): |
| if distributed.is_initialized(): |
| distributed.all_gather(data, src) |
| else: |
| data[0] = src |
|
|
|
|
| def get_rank(): |
| if distributed.is_initialized(): |
| return distributed.get_rank() |
| else: |
| return 0 |
|
|
|
|
| def get_world_size(): |
| if distributed.is_initialized(): |
| return distributed.get_world_size() |
| else: |
| return 1 |
|
|
|
|
| def chunk_size(size, rank, world_size): |
| extra = rank < size % world_size |
| return size // world_size + extra |
