File size: 3,080 Bytes
e69b72a | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Pure DistributedDataParallel helpers driven by torchrun environment variables.
No DeepSpeed, no ZeRO, no CPU/parameter offload: this module only wraps the
standard ``torch.distributed`` NCCL process group and the ``torchrun`` launch
contract (``RANK``/``LOCAL_RANK``/``WORLD_SIZE``). When launched as a single
process (no ``WORLD_SIZE`` or ``WORLD_SIZE==1``) it degrades to plain
single-device execution, which is what the CPU unit tests exercise.
"""
from __future__ import annotations
import os
import random
from dataclasses import dataclass
import numpy as np
import torch
@dataclass(frozen=True, slots=True)
class DistInfo:
"""Resolved process-group topology for the current process."""
rank: int
local_rank: int
world_size: int
is_distributed: bool
@property
def is_main(self) -> bool:
return self.rank == 0
def setup_distributed(*, backend: str = "nccl") -> DistInfo:
"""Initialise the process group from torchrun env vars, if any.
Returns a :class:`DistInfo`. Safe to call when not launched by torchrun: it
reports a single-process topology and initialises nothing.
"""
world_size = int(os.environ.get("WORLD_SIZE", "1"))
if world_size <= 1:
return DistInfo(rank=0, local_rank=0, world_size=1, is_distributed=False)
import torch.distributed as dist
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
if torch.cuda.is_available():
torch.cuda.set_device(local_rank)
chosen_backend = backend
init_kwargs = {"device_id": torch.device("cuda", local_rank)}
else: # pragma: no cover - CPU multi-process is not used in practice.
chosen_backend = "gloo"
init_kwargs = {}
if not dist.is_initialized():
dist.init_process_group(backend=chosen_backend, **init_kwargs)
return DistInfo(
rank=dist.get_rank(),
local_rank=local_rank,
world_size=dist.get_world_size(),
is_distributed=True,
)
def cleanup_distributed(info: DistInfo) -> None:
"""Barrier + tear down the process group if one was created."""
if not info.is_distributed:
return
import torch.distributed as dist
if dist.is_initialized():
if torch.cuda.is_available():
dist.barrier(device_ids=[info.local_rank])
else:
dist.barrier()
dist.destroy_process_group()
def is_main_process(info: DistInfo) -> bool:
return info.rank == 0
def seed_everything(seed: int, *, rank: int = 0) -> None:
"""Seed Python, NumPy, and torch RNGs.
Model construction uses the base ``seed`` on every rank so initial weights
match before DDP broadcasts them; data-order randomness is decorrelated via
the rank offset. Full cuDNN determinism is intentionally not forced (it would
slow training); the seeds are recorded for reproducibility instead.
"""
effective = seed + rank
random.seed(effective)
np.random.seed(effective % (2**32))
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
|