torchfoo.distributed.parallelize#

torchfoo.distributed.parallelize(world_size=None, master_addr=None, master_port=None, backend='auto', force=False)[source]#

Decorator that parallelizes a function across multiple devices with distributed setup.

Handles process spawning, setup_distributed, and cleanup_distributed automatically.

Use get_rank() and get_world_size() inside the decorated function to query distributed state.

Parameters:

world_size (int | None) – number of devices to use. Defaults to max(torch.cuda.device_count(), 1). If greater than the number of CUDA GPUs available, CPUs are used instead.
master_addr (str | None) – address of the master node. Falls back to MASTER_ADDR env var, then “localhost”.
master_port (str | int | None) – port of the master node. Falls back to MASTER_PORT env var, then an open port.
backend (str) – distributed backend (default: “auto”). “auto” selects “nccl” if enough CUDA GPUs are available for world_size, “gloo” otherwise.
force (bool) – Force distributed even if detected world_size is 1 (default False)

Examples:

import torchfoo as tfoo

@tfoo.dist.parallelize()
def train(cfg):
    rank = tfoo.dist.get_rank()
    world_size = tfoo.dist.get_world_size()
    ...

if __name__ == "__main__":
    train(cfg)