ColossalAI/applications/ChatGPT/chatgpt/trainer/utils.py
2023-02-14 22:17:25 +08:00

6 lines
121 B
Python

import torch.distributed as dist
def is_rank_0() -> bool:
return not dist.is_initialized() or dist.get_rank() == 0