From b24f35c5ef4fb616bf0933a361cc048e39fa98ed Mon Sep 17 00:00:00 2001 From: George Pearse Date: Mon, 29 Dec 2025 23:20:39 -0500 Subject: [PATCH] feat: add auto-DDP training support --- .gitignore | 6 + scripts/train_auto_ddp.py | 40 +++ tests/test_runtime/test_dist_env_fallback.py | 48 +++ visdet/engine/dist/__init__.py | 288 +++++------------ visdet/engine/dist/dist.py | 317 +++++++++++++++++++ visdet/engine/dist/dist_utils.py | 54 ++++ visdet/engine/dist/utils.py | 298 ++++++++++++++++- visdet/engine/runner/__init__.py | 2 + visdet/engine/runner/auto_train.py | 183 +++++++++++ 9 files changed, 1015 insertions(+), 221 deletions(-) create mode 100644 scripts/train_auto_ddp.py create mode 100644 tests/test_runtime/test_dist_env_fallback.py create mode 100644 visdet/engine/dist/dist.py create mode 100644 visdet/engine/dist/dist_utils.py create mode 100644 visdet/engine/runner/auto_train.py diff --git a/.gitignore b/.gitignore index d49cb08b..04c56d28 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,12 @@ __pycache__/ build/ develop-eggs/ dist/ + +# Keep visdet runtime dist package tracked +!visdet/engine/dist/ +!visdet/engine/dist/** +visdet/engine/dist/__pycache__/ + downloads/ eggs/ .eggs/ diff --git a/scripts/train_auto_ddp.py b/scripts/train_auto_ddp.py new file mode 100644 index 00000000..7e208e69 --- /dev/null +++ b/scripts/train_auto_ddp.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +"""Train with automatic single-node DDP (no torchrun). + +This is a thin wrapper around `visdet.engine.runner.auto_train.auto_train`. + +Usage: + python scripts/train_auto_ddp.py path/to/config.py + +Notes: +- The config must be compatible with `visdet.engine.runner.Runner.from_cfg()`. +- If multiple GPUs are available, one worker process is spawned per GPU. +""" + +import argparse + +from visdet.engine.config import Config +from visdet.engine.runner import auto_train + +_CONFIG_PATH: str | None = None + + +def _config_builder(_rank: int, _world_size: int) -> tuple[Config, dict]: + assert _CONFIG_PATH is not None + cfg = Config.fromfile(_CONFIG_PATH) + return cfg, {"config": _CONFIG_PATH} + + +def main() -> None: + parser = argparse.ArgumentParser(description="visdet auto-DDP training") + parser.add_argument("config", help="Path to a Runner.from_cfg config") + args = parser.parse_args() + + global _CONFIG_PATH + _CONFIG_PATH = args.config + + auto_train(_config_builder) + + +if __name__ == "__main__": + main() diff --git a/tests/test_runtime/test_dist_env_fallback.py b/tests/test_runtime/test_dist_env_fallback.py new file mode 100644 index 00000000..4b4d3e5e --- /dev/null +++ b/tests/test_runtime/test_dist_env_fallback.py @@ -0,0 +1,48 @@ +import os + +from visdet.engine import dist + + +def test_get_rank_world_size_env_fallback(monkeypatch): + # Ensure process group isn't initialized in this unit test + assert not dist.is_distributed() + + monkeypatch.setenv("RANK", "3") + monkeypatch.setenv("WORLD_SIZE", "8") + + assert dist.get_rank() == 3 + assert dist.get_world_size() == 8 + assert dist.get_dist_info() == (3, 8) + + +def test_infer_launcher_env(monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "2") + assert dist.infer_launcher() == "pytorch" + + monkeypatch.delenv("WORLD_SIZE", raising=False) + monkeypatch.setenv("SLURM_NTASKS", "2") + assert dist.infer_launcher() == "slurm" + + monkeypatch.delenv("SLURM_NTASKS", raising=False) + monkeypatch.setenv("OMPI_COMM_WORLD_LOCAL_RANK", "0") + assert dist.infer_launcher() == "mpi" + + monkeypatch.delenv("OMPI_COMM_WORLD_LOCAL_RANK", raising=False) + assert dist.infer_launcher() == "none" + + +def test_master_only_decorator(monkeypatch): + monkeypatch.setenv("RANK", "1") + + called = {"value": False} + + @dist.master_only + def _fn(): + called["value"] = True + + _fn() + assert called["value"] is False + + monkeypatch.setenv("RANK", "0") + _fn() + assert called["value"] is True diff --git a/visdet/engine/dist/__init__.py b/visdet/engine/dist/__init__.py index d8695cb4..5655f398 100644 --- a/visdet/engine/dist/__init__.py +++ b/visdet/engine/dist/__init__.py @@ -1,221 +1,77 @@ +# ruff: noqa +# type: ignore # Copyright (c) OpenMMLab. All rights reserved. -"""Distributed utilities for visdet.""" -import functools -import os -import pickle -import warnings -from typing import Any, List, Optional - -import torch -import torch.distributed as dist_lib - - -def _is_dist_available_and_initialized(): - """Check if distributed training is available and initialized.""" - return dist_lib.is_available() and dist_lib.is_initialized() - - -def get_dist_info(): - """Get distributed training info. - - Returns: - tuple: rank, world_size - """ - if _is_dist_available_and_initialized(): - rank = dist_lib.get_rank() - world_size = dist_lib.get_world_size() - else: - rank = 0 - world_size = 1 - return rank, world_size - - -def get_rank(): - """Get rank of current process.""" - if _is_dist_available_and_initialized(): - return dist_lib.get_rank() - return 0 - - -def get_world_size(): - """Get world size.""" - if _is_dist_available_and_initialized(): - return dist_lib.get_world_size() - return 1 - - -def is_distributed(): - """Check if distributed training is initialized.""" - return _is_dist_available_and_initialized() - - -def is_main_process(): - """Check if current process is main process (rank 0).""" - return get_rank() == 0 - - -def master_only(func): - """Decorator to make a function only execute on master process.""" - @functools.wraps(func) - def wrapper(*args, **kwargs): - if is_main_process(): - return func(*args, **kwargs) - return wrapper - - -def barrier(): - """Synchronize all processes.""" - if _is_dist_available_and_initialized(): - dist_lib.barrier() - - -def broadcast(data: Any, src: int = 0, group: Any | None = None) -> Any: - """Broadcast data from src rank to all ranks.""" - if not _is_dist_available_and_initialized(): - return data - - if isinstance(data, torch.Tensor): - dist_lib.broadcast(data, src, group=group) - return data - else: - # For non-tensor data, convert to tensor, broadcast, then convert back - if get_rank() == src: - data_tensor = torch.tensor(data, device='cuda' if torch.cuda.is_available() else 'cpu') - else: - data_tensor = torch.zeros_like(torch.tensor(data, device='cuda' if torch.cuda.is_available() else 'cpu')) - dist_lib.broadcast(data_tensor, src, group=group) - return data_tensor.item() if data_tensor.dim() == 0 else data_tensor - - -def broadcast_object_list(obj_list, src=0, group=None): - """Broadcast a list of objects from src rank to all ranks.""" - if not _is_dist_available_and_initialized(): - return obj_list - - dist_lib.broadcast_object_list(obj_list, src, group=group) - return obj_list - - -def all_reduce_params(model): - """All reduce model parameters for synchronization.""" - if not _is_dist_available_and_initialized(): - return - - world_size = get_world_size() - for param in model.parameters(): - if param.requires_grad and param.grad is not None: - dist_lib.all_reduce(param.grad.data) - param.grad.data.div_(world_size) - - -def init_dist(launcher: str = "pytorch", backend: str = "nccl", **kwargs: Any) -> tuple[int, int]: - """Initialize distributed environment.""" - if _is_dist_available_and_initialized(): - return get_dist_info() - - if launcher == 'pytorch': - dist_lib.init_process_group(backend=backend, **kwargs) - else: - raise NotImplementedError(f'Launcher {launcher} is not supported') - - return get_dist_info() - - -def collect_results(result_part, size, tmpdir=None): - """Collect results from all processes and merge them.""" - rank, world_size = get_dist_info() - - # Non-distributed mode: just return the results directly - if world_size == 1: - return result_part - - if tmpdir is None: - tmpdir = '.' - - # Create result file - result_file = os.path.join(tmpdir, f'result_rank_{rank}.pkl') - with open(result_file, 'wb') as f: - pickle.dump(result_part, f) - - dist_lib.barrier() - - if rank == 0: - results = [] - for i in range(world_size): - result_file = os.path.join(tmpdir, f'result_rank_{i}.pkl') - with open(result_file, 'rb') as f: - results.append(pickle.load(f)) - - # Clean up - for i in range(world_size): - result_file = os.path.join(tmpdir, f'result_rank_{i}.pkl') - if os.path.exists(result_file): - os.remove(result_file) - - # Merge results (flatten list of lists) - merged_results = [] - for result in results: - if isinstance(result, list): - merged_results.extend(result) - else: - merged_results.append(result) - return merged_results - - return None - - -def sync_random_seed(seed=None, device='cuda'): - """Make sure different ranks share the same seed. - - All workers must call this function, otherwise it will deadlock. - """ - import numpy as np - - if seed is None: - seed = np.random.randint(2**31) - - rank, world_size = get_dist_info() - - if world_size == 1: - return seed - - if not _is_dist_available_and_initialized(): - return seed - - if rank == 0: - random_num = torch.tensor(seed, dtype=torch.int32, device=device) - else: - random_num = torch.tensor(0, dtype=torch.int32, device=device) - - dist_lib.broadcast(random_num, src=0) - return random_num.item() - - -def infer_launcher(): - """Infer launcher type from environment variables.""" - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: - return 'pytorch' - else: - return None - - -# Backward compatibility -utils = type('utils', (), {'master_only': master_only})() +from visdet.engine.dist.dist import ( + all_gather, + all_gather_object, + all_reduce, + all_reduce_dict, + all_reduce_params, + broadcast, + broadcast_object_list, + collect_results, + collect_results_cpu, + collect_results_gpu, + gather, + gather_object, + sync_random_seed, +) +from visdet.engine.dist.dist_utils import broadcast_from_rank_0, rank_0_only, rank_0_only_method +from visdet.engine.dist.utils import ( + barrier, + cast_data_device, + get_backend, + get_comm_device, + get_data_device, + get_default_group, + get_dist_info, + get_local_group, + get_local_rank, + get_local_size, + get_rank, + get_world_size, + infer_launcher, + init_dist, + init_local_group, + is_distributed, + is_main_process, + master_only, +) __all__ = [ - 'get_dist_info', - 'get_rank', - 'get_world_size', - 'is_distributed', - 'is_main_process', - 'master_only', - 'barrier', - 'broadcast', - 'broadcast_object_list', - 'all_reduce_params', - 'init_dist', - 'collect_results', - 'sync_random_seed', - 'infer_launcher', + "all_gather", + "all_gather_object", + "all_reduce", + "all_reduce_dict", + "all_reduce_params", + "barrier", + "broadcast", + "broadcast_from_rank_0", + "broadcast_object_list", + "cast_data_device", + "collect_results", + "collect_results_cpu", + "collect_results_gpu", + "gather", + "gather_object", + "get_backend", + "get_comm_device", + "get_data_device", + "get_default_group", + "get_dist_info", + "get_local_group", + "get_local_rank", + "get_local_size", + "get_rank", + "get_world_size", + "infer_launcher", + "init_dist", + "init_local_group", + "is_distributed", + "is_main_process", + "master_only", + "rank_0_only", + "rank_0_only_method", + "sync_random_seed", ] diff --git a/visdet/engine/dist/dist.py b/visdet/engine/dist/dist.py new file mode 100644 index 00000000..96749aa8 --- /dev/null +++ b/visdet/engine/dist/dist.py @@ -0,0 +1,317 @@ +# ruff: noqa +# type: ignore +# Copyright (c) OpenMMLab. All rights reserved. + +import os.path as osp +import pickle +import shutil +import tempfile +from collections import OrderedDict +from collections.abc import Generator +from itertools import chain, zip_longest +from typing import Any + +import numpy as np +import torch +from torch import Tensor +from torch import distributed as torch_dist +from torch._utils import _flatten_dense_tensors, _take_tensors, _unflatten_dense_tensors +from torch.distributed import ProcessGroup + +from visdet.engine.utils import mkdir_or_exist + +from visdet.engine.dist.utils import ( + barrier, + cast_data_device, + get_backend, + get_comm_device, + get_data_device, + get_default_group, + get_dist_info, + get_rank, + get_world_size, +) + + +def _get_reduce_op(name: str) -> torch_dist.ReduceOp: + op_mappings = { + "sum": torch_dist.ReduceOp.SUM, + "product": torch_dist.ReduceOp.PRODUCT, + "min": torch_dist.ReduceOp.MIN, + "max": torch_dist.ReduceOp.MAX, + "band": torch_dist.ReduceOp.BAND, + "bor": torch_dist.ReduceOp.BOR, + "bxor": torch_dist.ReduceOp.BXOR, + } + if name.lower() not in op_mappings: + raise ValueError(f"reduce op should be one of {op_mappings.keys()}, but got {name}") + return op_mappings[name.lower()] + + +def all_reduce(data: Tensor, op: str = "sum", group: ProcessGroup | None = None) -> None: + world_size = get_world_size(group) + if world_size > 1: + if group is None: + group = get_default_group() + + input_device = get_data_device(data) + backend_device = get_comm_device(group) + data_on_device = cast_data_device(data, backend_device) + + if op.lower() == "mean": + torch_dist.all_reduce(data_on_device, _get_reduce_op("sum"), group) + data_on_device = torch.true_divide(data_on_device, world_size) + else: + torch_dist.all_reduce(data_on_device, _get_reduce_op(op), group) + + cast_data_device(data_on_device, input_device, out=data) + + +def all_gather(data: Tensor, group: ProcessGroup | None = None) -> list[Tensor]: + world_size = get_world_size(group) + if world_size == 1: + return [data] + + if group is None: + group = get_default_group() + + input_device = get_data_device(data) + backend_device = get_comm_device(group) + data_on_device = cast_data_device(data, backend_device) + + gather_list = [torch.empty_like(data, device=backend_device) for _ in range(world_size)] + torch_dist.all_gather(gather_list, data_on_device, group) + + return cast_data_device(gather_list, input_device) # type: ignore + + +def gather(data: Tensor, dst: int = 0, group: ProcessGroup | None = None) -> list[Tensor | None]: + world_size = get_world_size(group) + if world_size == 1: + return [data] + + if group is None: + group = get_default_group() + + input_device = get_data_device(data) + backend_device = get_comm_device(group) + + if get_rank(group) == dst: + gather_list = [torch.empty_like(data, device=backend_device) for _ in range(world_size)] + else: + gather_list = [] + + torch_dist.gather(data, gather_list, dst, group) + + if get_rank(group) == dst: + return cast_data_device(gather_list, input_device) # type: ignore + return gather_list + + +def broadcast(data: Tensor, src: int = 0, group: ProcessGroup | None = None) -> None: + if get_world_size(group) > 1: + if group is None: + group = get_default_group() + + input_device = get_data_device(data) + backend_device = get_comm_device(group) + data_on_device = cast_data_device(data, backend_device) + data_on_device = data_on_device.contiguous() # type: ignore + torch_dist.broadcast(data_on_device, src, group) + + if get_rank(group) != src: + cast_data_device(data_on_device, input_device, data) + + +def sync_random_seed(group: ProcessGroup | None = None) -> int: + seed = np.random.randint(2**31) + if get_world_size(group) == 1: + return seed + + if group is None: + group = get_default_group() + + backend_device = get_comm_device(group) + + if get_rank(group) == 0: + random_num = torch.tensor(seed, dtype=torch.int32).to(backend_device) + else: + random_num = torch.tensor(0, dtype=torch.int32).to(backend_device) + + torch_dist.broadcast(random_num, src=0, group=group) + return int(random_num.item()) + + +def broadcast_object_list(data: list[Any], src: int = 0, group: ProcessGroup | None = None) -> None: + assert isinstance(data, list) + + if get_world_size(group) > 1: + if group is None: + group = get_default_group() + + torch_dist.broadcast_object_list(data, src, group) + + +def all_reduce_dict(data: dict[str, Tensor], op: str = "sum", group: ProcessGroup | None = None) -> None: + assert isinstance(data, dict) + + world_size = get_world_size(group) + if world_size > 1: + if group is None: + group = get_default_group() + + keys = sorted(data.keys()) + tensor_shapes = [data[k].shape for k in keys] + tensor_sizes = [data[k].numel() for k in keys] + + flatten_tensor = torch.cat([data[k].flatten() for k in keys]) + all_reduce(flatten_tensor, op=op, group=group) + + split_tensors = [ + x.reshape(shape) for x, shape in zip(torch.split(flatten_tensor, tensor_sizes), tensor_shapes, strict=False) + ] + + for k, v in zip(keys, split_tensors, strict=False): + data[k] = v + + +def all_gather_object(data: Any, group: ProcessGroup | None = None) -> list[Any]: + world_size = get_world_size(group) + if world_size == 1: + return [data] + + if group is None: + group = get_default_group() + + object_list: list[Any] = [None] * world_size + torch_dist.all_gather_object(object_list, data, group) + return object_list + + +def gather_object(data: Any, dst: int = 0, group: ProcessGroup | None = None) -> list[Any] | None: + world_size = get_world_size(group) + if world_size == 1: + return [data] + + if group is None: + group = get_default_group() + + gather_list = [None] * world_size if get_rank(group) == dst else None + torch_dist.gather_object(data, gather_list, dst, group) + return gather_list + + +def collect_results(results: list, size: int, device: str = "cpu", tmpdir: str | None = None) -> list | None: + if device not in ["gpu", "cpu"]: + raise NotImplementedError(f"device must be 'cpu' or 'gpu', but got {device}") + + if device == "gpu": + assert tmpdir is None, "tmpdir should be None when device is 'gpu'" + return _collect_results_device(results, size) + + return collect_results_cpu(results, size, tmpdir) + + +def collect_results_cpu(result_part: list, size: int, tmpdir: str | None = None) -> list | None: + rank, world_size = get_dist_info() + if world_size == 1: + return result_part[:size] + + if tmpdir is None: + max_len = 512 + dir_tensor = torch.full((max_len,), 32, dtype=torch.uint8) + if rank == 0: + mkdir_or_exist(".dist_test") + tmpdir_path = tempfile.mkdtemp(dir=".dist_test") + tmpdir_tensor = torch.tensor(bytearray(tmpdir_path.encode()), dtype=torch.uint8) + dir_tensor[: len(tmpdir_tensor)] = tmpdir_tensor + broadcast(dir_tensor, 0) + tmpdir = dir_tensor.numpy().tobytes().decode().rstrip() + else: + mkdir_or_exist(tmpdir) + + with open(osp.join(tmpdir, f"part_{rank}.pkl"), "wb") as f: + pickle.dump(result_part, f, protocol=2) + + barrier() + + if rank != 0: + return None + + part_list = [] + for i in range(world_size): + path = osp.join(tmpdir, f"part_{i}.pkl") + if not osp.exists(path): + raise FileNotFoundError( + f"{tmpdir} is not a shared directory for rank {i}. Ensure it is shared for all ranks." + ) + with open(path, "rb") as f: + part_list.append(pickle.load(f)) + + zipped_results = zip_longest(*part_list) + ordered_results = [i for i in chain.from_iterable(zipped_results) if i is not None] + ordered_results = ordered_results[:size] + + shutil.rmtree(tmpdir) + return ordered_results + + +def _collect_results_device(result_part: list, size: int) -> list | None: + rank, world_size = get_dist_info() + if world_size == 1: + return result_part[:size] + + part_list = all_gather_object(result_part) + + if rank == 0: + zipped_results = zip_longest(*part_list) + ordered_results = [i for i in chain.from_iterable(zipped_results) if i is not None] + return ordered_results[:size] + + return None + + +def collect_results_gpu(result_part: list, size: int) -> list | None: + return _collect_results_device(result_part, size) + + +def _all_reduce_coalesced( + tensors: list[torch.Tensor], + bucket_size_mb: int = -1, + op: str = "sum", + group: ProcessGroup | None = None, +) -> None: + if bucket_size_mb > 0: + bucket_size_bytes = bucket_size_mb * 1024 * 1024 + buckets = _take_tensors(tensors, bucket_size_bytes) + else: + buckets = OrderedDict() + for tensor in tensors: + tp = tensor.type() + buckets.setdefault(tp, []).append(tensor) + buckets = buckets.values() + + for bucket in buckets: + flat_tensors = _flatten_dense_tensors(bucket) + all_reduce(flat_tensors, op=op, group=group) + for tensor, synced in zip(bucket, _unflatten_dense_tensors(flat_tensors, bucket), strict=False): + tensor.copy_(synced) + + +def all_reduce_params( + params: list | Generator[torch.Tensor, None, None], + coalesce: bool = True, + bucket_size_mb: int = -1, + op: str = "sum", + group: ProcessGroup | None = None, +) -> None: + world_size = get_world_size(group) + if world_size == 1: + return + + params_data = [param.data for param in params] + if coalesce: + _all_reduce_coalesced(params_data, bucket_size_mb, op=op, group=group) + else: + for tensor in params_data: + all_reduce(tensor, op=op, group=group) diff --git a/visdet/engine/dist/dist_utils.py b/visdet/engine/dist/dist_utils.py new file mode 100644 index 00000000..0e2808f7 --- /dev/null +++ b/visdet/engine/dist/dist_utils.py @@ -0,0 +1,54 @@ +"""Distributed training utilities and decorators. + +This module provides convenient decorators and helpers for distributed training, +helping centralize rank-specific operations and reduce boilerplate. +""" + +import functools +from collections.abc import Callable +from typing import Any, TypeVar + +F = TypeVar("F", bound=Callable[..., Any]) + + +def rank_0_only(func: F) -> F: + """Decorator that ensures a function only executes on rank 0.""" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + from visdet.engine.dist import is_main_process + + if not is_main_process(): + return None + + return func(*args, **kwargs) + + return wrapper # type: ignore + + +def rank_0_only_method(func: F) -> F: + """Decorator for class methods that should only execute on rank 0.""" + + @functools.wraps(func) + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + from visdet.engine.dist import is_main_process + + if not is_main_process(): + return None + + return func(self, *args, **kwargs) + + return wrapper # type: ignore + + +def broadcast_from_rank_0(obj: Any) -> Any: + """Broadcast a picklable object from rank 0 to all ranks.""" + + from visdet.engine.dist import broadcast_object_list, get_world_size + + if get_world_size() == 1: + return obj + + obj_list = [obj] + broadcast_object_list(obj_list, src=0) + return obj_list[0] diff --git a/visdet/engine/dist/utils.py b/visdet/engine/dist/utils.py index 6fe1191c..4e4914db 100644 --- a/visdet/engine/dist/utils.py +++ b/visdet/engine/dist/utils.py @@ -1,18 +1,306 @@ +# ruff: noqa +# type: ignore # Copyright (c) OpenMMLab. All rights reserved. -"""Utility functions for distributed training.""" +import datetime import functools +import os +import subprocess +from collections.abc import Callable, Iterable, Mapping +import numpy as np +import torch +import torch.multiprocessing as mp +from torch import Tensor +from torch import distributed as torch_dist +from torch.distributed import ProcessGroup -def master_only(func): - """Decorator to make a function only execute on master process.""" - from visdet.engine.dist import is_main_process +_LOCAL_PROCESS_GROUP: ProcessGroup | None = None + +def is_distributed() -> bool: + """Return True if distributed environment has been initialized.""" + + return torch_dist.is_available() and torch_dist.is_initialized() + + +def get_local_group() -> ProcessGroup | None: + """Return local process group.""" + + if not is_distributed(): + return None + + if _LOCAL_PROCESS_GROUP is None: + raise RuntimeError( + "Local process group is not created, please use `init_local_group` to setup local process group." + ) + + return _LOCAL_PROCESS_GROUP + + +def get_default_group() -> ProcessGroup | None: + """Return default process group.""" + + return torch_dist.distributed_c10d._get_default_group() + + +def infer_launcher() -> str: + if "WORLD_SIZE" in os.environ: + return "pytorch" + if "SLURM_NTASKS" in os.environ: + return "slurm" + if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: + return "mpi" + return "none" + + +def init_dist(launcher: str, backend: str = "nccl", init_backend: str = "torch", **kwargs) -> None: + """Initialize distributed environment. + + Args: + launcher: Way to launch multi-process. Supported launchers are + 'pytorch', 'mpi', 'slurm'. + backend: torch.distributed backend. Typically 'nccl' or 'gloo'. + init_backend: Initialization backend. Keep as 'torch' for visdet. + **kwargs: Passed to torch.distributed.init_process_group. + """ + + timeout = kwargs.get("timeout", None) + if timeout is not None: + kwargs["timeout"] = datetime.timedelta(seconds=timeout) + + if mp.get_start_method(allow_none=True) is None: + mp.set_start_method("spawn") + + if launcher == "pytorch": + _init_dist_pytorch(backend, init_backend=init_backend, **kwargs) + elif launcher == "mpi": + _init_dist_mpi(backend, **kwargs) + elif launcher == "slurm": + _init_dist_slurm(backend, init_backend=init_backend, **kwargs) + else: + raise ValueError(f"Invalid launcher type: {launcher}") + + +def _init_dist_pytorch(backend: str, init_backend: str = "torch", **kwargs) -> None: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + + if init_backend != "torch": + raise ValueError(f'Only init_backend="torch" is supported, got {init_backend!r}') + + torch_dist.init_process_group(backend=backend, rank=rank, world_size=int(os.environ["WORLD_SIZE"]), **kwargs) + + +def _init_dist_mpi(backend: str, **kwargs) -> None: + local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = "29500" + if "MASTER_ADDR" not in os.environ: + raise KeyError("The environment variable MASTER_ADDR is not set") + + os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"] + os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] + os.environ["LOCAL_RANK"] = str(local_rank) + + torch_dist.init_process_group(backend=backend, **kwargs) + + +def _init_dist_slurm(backend: str, port: int | None = None, init_backend: str = "torch", **kwargs) -> None: + proc_id = int(os.environ["SLURM_PROCID"]) + ntasks = int(os.environ["SLURM_NTASKS"]) + node_list = os.environ["SLURM_NODELIST"] + + local_rank_env = os.environ.get("SLURM_LOCALID", None) + if local_rank_env is not None: + local_rank = int(local_rank_env) + else: + local_rank = proc_id % max(torch.cuda.device_count(), 1) + + addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1") + + if port is not None: + os.environ["MASTER_PORT"] = str(port) + elif "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = "29500" + + if "MASTER_ADDR" not in os.environ: + os.environ["MASTER_ADDR"] = addr + + os.environ["WORLD_SIZE"] = str(ntasks) + os.environ["LOCAL_RANK"] = str(local_rank) + os.environ["RANK"] = str(proc_id) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + + if init_backend != "torch": + raise ValueError(f'Only init_backend="torch" is supported, got {init_backend!r}') + + torch_dist.init_process_group(backend=backend, **kwargs) + + +def init_local_group(node_rank: int, num_gpus_per_node: int) -> None: + global _LOCAL_PROCESS_GROUP + assert _LOCAL_PROCESS_GROUP is None + + ranks = list(range(node_rank * num_gpus_per_node, (node_rank + 1) * num_gpus_per_node)) + _LOCAL_PROCESS_GROUP = torch_dist.new_group(ranks) + + +def get_backend(group: ProcessGroup | None = None) -> str | None: + if is_distributed(): + if group is None: + group = get_default_group() + return torch_dist.get_backend(group) + return None + + +def get_world_size(group: ProcessGroup | None = None) -> int: + if is_distributed(): + if group is None: + group = get_default_group() + return torch_dist.get_world_size(group) + + return int(os.environ.get("WORLD_SIZE", 1)) + + +def get_rank(group: ProcessGroup | None = None) -> int: + if is_distributed(): + if group is None: + group = get_default_group() + return torch_dist.get_rank(group) + + return int(os.environ.get("RANK", 0)) + + +def get_local_size() -> int: + if not is_distributed(): + return 1 + if _LOCAL_PROCESS_GROUP is None: + raise RuntimeError( + "Local process group is not created, please use `init_local_group` to setup local process group." + ) + return torch_dist.get_world_size(_LOCAL_PROCESS_GROUP) + + +def get_local_rank() -> int: + if not is_distributed(): + return 0 + if _LOCAL_PROCESS_GROUP is None: + raise RuntimeError( + "Local process group is not created, please use `init_local_group` to setup local process group." + ) + return torch_dist.get_rank(_LOCAL_PROCESS_GROUP) + + +def get_dist_info(group: ProcessGroup | None = None) -> tuple[int, int]: + world_size = get_world_size(group) + rank = get_rank(group) + return rank, world_size + + +def is_main_process(group: ProcessGroup | None = None) -> bool: + return get_rank(group) == 0 + + +def master_only(func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args, **kwargs): if is_main_process(): return func(*args, **kwargs) + return wrapper -__all__ = ['master_only'] +def barrier(group: ProcessGroup | None = None) -> None: + if is_distributed(): + if group is None: + group = get_default_group() + torch_dist.barrier(group) + + +def get_data_device(data: Tensor | Mapping | Iterable) -> torch.device: + if isinstance(data, Tensor): + return data.device + if isinstance(data, Mapping): + pre = None + for v in data.values(): + cur = get_data_device(v) + if pre is None: + pre = cur + elif cur != pre: + raise ValueError(f"device type in data should be consistent, but got {cur} and {pre}") + if pre is None: + raise ValueError("data should not be empty") + return pre + if isinstance(data, Iterable) and not isinstance(data, str) and not isinstance(data, np.ndarray): + pre = None + for item in data: + cur = get_data_device(item) + if pre is None: + pre = cur + elif cur != pre: + raise ValueError(f"device type in data should be consistent, but got {cur} and {pre}") + if pre is None: + raise ValueError("data should not be empty") + return pre + + raise TypeError(f"data should be a Tensor, sequence of tensor or dict, but got {type(data)}") + + +def get_comm_device(group: ProcessGroup | None = None) -> torch.device: + backend = get_backend(group) + if backend == torch_dist.Backend.NCCL: + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +def cast_data_device( + data: Tensor | Mapping | Iterable, + device: torch.device, + out: Tensor | Mapping | Iterable | None = None, +) -> Tensor | Mapping | Iterable: + if out is not None and type(data) is not type(out): + raise TypeError(f"out should be same type as data, got {type(data)} vs {type(out)}") + + if isinstance(data, Tensor): + data_on_device = data if get_data_device(data) == device else data.to(device) + if out is not None: + out.copy_(data_on_device) # type: ignore + return data_on_device + + if isinstance(data, Mapping): + data_on_device: dict = {} + if out is not None: + if len(data) != len(out): # type: ignore + raise ValueError("length of data and out should be same") + for k, v in data.items(): + data_on_device[k] = cast_data_device(v, device, out[k]) # type: ignore + else: + for k, v in data.items(): + data_on_device[k] = cast_data_device(v, device) + if not data_on_device: + raise ValueError("data should not be empty") + return type(data)(data_on_device) # type: ignore + + if isinstance(data, Iterable) and not isinstance(data, str) and not isinstance(data, np.ndarray): + data_on_device_list = [] + if out is not None: + for v1, v2 in zip(data, out, strict=False): + data_on_device_list.append(cast_data_device(v1, device, v2)) + else: + for v in data: + data_on_device_list.append(cast_data_device(v, device)) + if not data_on_device_list: + raise ValueError("data should not be empty") + return type(data)(data_on_device_list) # type: ignore + + raise TypeError(f"data should be a Tensor, list of tensor or dict, but got {type(data)}") diff --git a/visdet/engine/runner/__init__.py b/visdet/engine/runner/__init__.py index cae21846..d347af7b 100644 --- a/visdet/engine/runner/__init__.py +++ b/visdet/engine/runner/__init__.py @@ -18,6 +18,7 @@ save_checkpoint, weights_to_cpu, ) +from visdet.engine.runner.auto_train import auto_train from visdet.engine.runner.log_processor import LogProcessor from visdet.engine.runner.loops import EpochBasedTrainLoop, IterBasedTrainLoop, TestLoop, ValLoop from visdet.engine.runner.priority import Priority, get_priority @@ -35,6 +36,7 @@ "Runner", "TestLoop", "ValLoop", + "auto_train", "autocast", "find_latest_checkpoint", "get_deprecated_model_names", diff --git a/visdet/engine/runner/auto_train.py b/visdet/engine/runner/auto_train.py new file mode 100644 index 00000000..143250fb --- /dev/null +++ b/visdet/engine/runner/auto_train.py @@ -0,0 +1,183 @@ +"""Automatic distributed training with torch.multiprocessing. + +This module is ported from BinItAI/core (PR #5414) to provide an +"auto-DDP" mode that keeps the Python entrypoint unchanged. + +It detects the number of available GPUs and automatically configures +single-node DistributedDataParallel training when multiple GPUs are present. + +Key points: +- Uses `torch.multiprocessing.spawn()` (no `torchrun` required) +- Rank 0 builds the config, then broadcasts it to other ranks +- Each worker sets its CUDA device before initializing distributed state +""" + +import logging +import os +import socket +from collections.abc import Callable +from datetime import timedelta +from typing import Any + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +logger = logging.getLogger(__name__) + + +def _find_free_port() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return str(s.getsockname()[1]) + + +def _get_gpu_count_safe() -> int: + """Detect GPU count without initializing CUDA context in parent process.""" + + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if cuda_visible is not None: + if cuda_visible == "": + return 0 + devices = [d.strip() for d in cuda_visible.split(",") if d.strip()] + return len(devices) + + # Prefer nvidia-smi if available; it doesn't initialize CUDA runtime. + import subprocess + + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + lines = [line for line in result.stdout.strip().split("\n") if line.strip()] + return len(lines) + except Exception: + pass + + # Do not call torch.cuda.device_count() here. + logger.warning( + "Could not detect GPU count via nvidia-smi or CUDA_VISIBLE_DEVICES; " + "defaulting to single-GPU mode. Set CUDA_VISIBLE_DEVICES to enable multi-GPU." + ) + return 1 + + +def _worker_fn( + rank: int, + world_size: int, + master_addr: str, + master_port: str, + config_builder: Callable[[int, int], tuple[Any, dict]], + rank_0_callback: Callable[[Any, dict], None] | None, +) -> None: + try: + os.environ["MASTER_ADDR"] = master_addr + os.environ["MASTER_PORT"] = master_port + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + if torch.cuda.is_available(): + torch.cuda.set_device(rank) + + backend = os.environ.get("DIST_BACKEND", "nccl" if torch.cuda.is_available() else "gloo") + dist.init_process_group( + backend=backend, + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + timeout=timedelta(seconds=3600), + ) + + if backend == "nccl" and torch.cuda.is_available(): + dist.barrier(device_ids=[rank]) + else: + dist.barrier() + + from visdet.engine.runner import Runner + + if rank == 0: + cfg, readable = config_builder(rank, world_size) + cfg.launcher = "pytorch" + + if rank_0_callback is not None: + rank_0_callback(cfg, readable) + else: + cfg = None + readable = None + + object_list = [cfg, readable] + dist.broadcast_object_list(object_list, src=0) + cfg, readable = object_list + + if cfg is None: + raise RuntimeError(f"[Rank {rank}] Config broadcast failed") + + runner = Runner.from_cfg(cfg) + runner.readable_config = readable + runner.train() + + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _single_gpu_train( + config_builder: Callable[[int, int], tuple[Any, dict]], + rank_0_callback: Callable[[Any, dict], None] | None, +) -> None: + from visdet.engine.runner import Runner + + cfg, readable = config_builder(0, 1) + cfg.launcher = "none" + + if rank_0_callback is not None: + rank_0_callback(cfg, readable) + + runner = Runner.from_cfg(cfg) + runner.readable_config = readable + runner.train() + + +def auto_train( + config_builder: Callable[[int, int], tuple[Any, dict]], + rank_0_callback: Callable[[Any, dict], None] | None = None, +) -> None: + """Automatically run training in single or multi-GPU mode. + + If multiple GPUs are detected, this spawns one worker per GPU using + `torch.multiprocessing.spawn()` and initializes a single-node process group. + + Args: + config_builder: Function taking `(rank, world_size)` and returning + `(cfg, readable_dict)` where `cfg` is compatible with `Runner.from_cfg`. + Note: this function must be picklable (module-level) to work with spawn. + rank_0_callback: Optional callback executed only on rank 0 after config + is built but before training starts. + """ + + if mp.get_start_method(allow_none=True) != "spawn": + mp.set_start_method("spawn", force=True) + + gpu_count = _get_gpu_count_safe() + + if gpu_count <= 1: + logger.info("Detected %d GPU(s); running single-process training", gpu_count) + return _single_gpu_train(config_builder, rank_0_callback) + + logger.info("Detected %d GPUs; enabling automatic DDP", gpu_count) + + master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1") + master_port = os.environ.get("MASTER_PORT", _find_free_port()) + + mp.spawn( + _worker_fn, + args=(gpu_count, master_addr, master_port, config_builder, rank_0_callback), + nprocs=gpu_count, + join=True, + )