File size: 905 Bytes
3e363e2 | 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 | """Optional DDP/FSDP construction without hard-coded GPU counts."""
from __future__ import annotations
import os
import torch
from torch import nn
def wrap_distributed(
model: nn.Module, *, use_fsdp: bool = False, backend: str = "nccl"
) -> nn.Module:
world_size = int(os.environ.get("WORLD_SIZE", "1"))
if world_size <= 1:
return model
if backend == "nccl" and not torch.cuda.is_available():
backend = "gloo"
if not torch.distributed.is_initialized():
torch.distributed.init_process_group(backend=backend)
if use_fsdp:
from torch.distributed.fsdp import FullyShardedDataParallel
return FullyShardedDataParallel(model)
device_id = torch.cuda.current_device() if torch.cuda.is_available() else None
return torch.nn.parallel.DistributedDataParallel(
model, device_ids=[device_id] if device_id is not None else None
)
|