File size: 1,278 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """
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
|