File size: 1,111 Bytes
b9fc164 | 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 | """Distributed process setup for data-parallel and FSDP training."""
import os
from datetime import timedelta
import torch
import torch.distributed as dist
def print_rank0(message: str) -> None:
if int(os.getenv("RANK", "0")) == 0:
print(message)
def init_distributed_environment(backend: str = "nccl") -> None:
if dist.is_initialized():
return
global_rank = int(os.getenv("RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
print_rank0(f"Initializing distributed environment with world_size={world_size}")
dist.init_process_group(
backend=backend, world_size=world_size, rank=global_rank,
timeout=timedelta(minutes=30),
)
def setup_distributed_training() -> tuple[int, int, int, torch.device]:
local_rank = int(os.getenv("LOCAL_RANK", "0"))
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
init_distributed_environment()
return local_rank, dist.get_rank(), dist.get_world_size(), device
def clean_dist_env() -> None:
if dist.is_initialized():
dist.destroy_process_group()
|