Introduction

Distributed Data Parallel Wrapper (DDPW) is a lightweight Python wrapper relevant to PyTorch users. It is written in Python 3.13.

DDPW enables writing compute-intensive tasks (such as training models) without deeply worrying about the underlying compute platform (CPU, Apple SoC, GPUs, or SLURM (uses Submitit)) and instead allows specifying it simply as an argument. This considerably minimises the need to change the code for each type of platform.

DDPW handles basic logistical tasks such as creating threads on GPUs/SLURM nodes, setting up inter-process communication, etc., and provides simple, default utility methods to move modules to devices and get dataset samplers, allowing the user to focus on the main aspects of the task.

Installation

DDPW is distributed on PyPI. The source code is available on GitHub and can be used to manually build it as a dependency package.

Target platforms

This wrapper is released for all architectures but is tested only on Linux arch-64 and Apple SoC.

PyPI

With uv

# to instal and add to pyroject.toml
uv add [--active] ddpw
# or to simply instal
uv pip install ddpw

With pip

pip install ddpw

GitHub

pip install git+'https://github.com/sujaltv/ddpw'

Manual build

> git clone https://github.com/sujaltv/ddpw
> cd ddpw

> uv pip install .

Usage

As a decorator

from ddpw import Platform, wrapper

platform = Platform(device="gpu", n_cpus=32, ram=64, n_gpus=4, verbose=True)

@wrapper(platform)
def run(*args, **kwargs):
    # global and local ranks, and the process group in:
    # kwargs['global_rank'], # kwargs['local_rank'], kwargs['group']
    pass

if __name__ == '__main__':
    run(*args, **kwargs)

As a callable

from ddpw import Platform, Wrapper

# some task
def run(*args, **kwargs):
    # global and local ranks, and the process group in:
    # kwargs['global_rank'], # kwargs['local_rank'], kwargs['group']
    pass

# platform (e.g., 4 GPUs)
platform = Platform(device='gpu', n_gpus=4)

# wrapper
wrapper = Wrapper(platform=platform)

# start
wrapper.start(task, *args, **kwargs)

Refer to the API for more configuration options or the example with MNIST for an illustration.