| """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 | |
| ) | |