File size: 378 Bytes
681f346
8652b14
681f346
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
import os

# Check standard environment variables for distributed training
# Default to True (rank 0) if not in a distributed environment
_rank = int(os.environ.get("RANK", 0))
_local_rank = int(os.environ.get("LOCAL_RANK", 0))

# We consider it rank zero if global rank is 0.
# Local rank check is usually redundant if rank is 0, but good for sanity.
is_rank_zero = _rank == 0