| """ |
| This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research |
| template [repo](https://github.com/buoyancy99/research-template). |
| By its MIT license, you must keep the above sentence in `README.md` |
| and the `LICENSE` file to credit the author. |
| """ |
|
|
| from typing import Callable |
| import torch |
| import torch.distributed as dist |
| from lightning.pytorch.utilities.rank_zero import rank_zero_only |
| from utils.print_utils import print_once |
|
|
| is_rank_zero = rank_zero_only.rank == 0 |
|
|
| rank_zero_print = rank_zero_only(print) |
| rank_zero_print_once = rank_zero_only(print_once) |
|
|
|
|
| def broadcast_from_zero( |
| init_func: Callable[[], torch.Tensor], func: Callable[[], torch.Tensor] |
| ) -> torch.Tensor: |
| """ |
| Broadcasts the result of func() from rank 0 to all other processes. |
| If distributed is not initialized, returns the result of func() directly. |
| Args: |
| init_func: function that initializes a tensor (will be overwritten) |
| func: function that generates the tensor to broadcast from rank 0 |
| Returns: |
| torch.Tensor: the tensor generated by func() on rank 0 |
| """ |
| if not dist.is_initialized(): |
| return func() |
| tensor = init_func() |
| if is_rank_zero: |
| tensor = func() |
| dist.broadcast(tensor, src=0) |
| return tensor |
|
|
