Utilities for PyTorch distributed
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

import importlib
import inspect
from pathlib import Path
import sys

base_path = Path(__file__).parent.parent.resolve()
sys.path.insert(0, str(base_path))

import torch_dist_utils as du


# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

project = "torch-dist-utils"
copyright = "2023, Katherine Crowson"
author = "Katherine Crowson"
version = release = du.__version__

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.linkcode",
    "sphinx_rtd_theme",
]

templates_path = ["_templates"]
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]

# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = "sphinx_rtd_theme"
html_static_path = ["_static"]
pygments_style = "sphinx"

# -- Options for autodoc extension -------------------------------------------

autodoc_member_order = "bysource"
autodoc_typehints = "description"

# -- Options for linkcode extension ------------------------------------------

code_url = "https://github.com/crowsonkb/dist-utils/blob/main"


def linkcode_resolve(domain, info):
    if domain != "py":
        return None

    mod = importlib.import_module(info["module"])
    if "." in info["fullname"]:
        objname, attrname = info["fullname"].split(".")
        obj = getattr(mod, objname)
        try:
            # object is a method of a class
            obj = getattr(obj, attrname)
        except AttributeError:
            # object is an attribute of a class
            return None
    else:
        obj = getattr(mod, info["fullname"])

    try:
        while True:
            obj = obj.__wrapped__
    except AttributeError:
        pass

    try:
        file = inspect.getsourcefile(obj)
        lines = inspect.getsourcelines(obj)
    except TypeError:
        # e.g. object is a typing.Union
        return None
    try:
        file = Path(file).relative_to(base_path)
    except ValueError:
        return None
    start, end = lines[1], lines[1] + len(lines[0]) - 1

    return f"{code_url}/{file}#L{start}-L{end}"

docs/conf.py
crowsonkb/torch-dist-utils
#!/usr/bin/env python3

"""Trains a classifier on CIFAR-10 using torch_dist_utils and DDP.

It can be run on CPU or one GPU with:

    python train_classifier.py

or on multiple GPUs with:

    torchrun --nproc-per-node gpu train_classifier.py
"""

import argparse
from contextlib import contextmanager, nullcontext

import torch
from torch import distributed as dist, nn, optim
from torch.nn import functional as F
from torch.utils import data
import torch_dist_utils as du
from torchvision import datasets, transforms
from tqdm import tqdm

print0 = tqdm.external_write_mode()(du.print0)


@contextmanager
def train_mode(model, mode=True):
    """A context manager that places a model into training mode and restores
    the previous mode on exit."""
    modes = [module.training for module in model.modules()]
    try:
        model.train(mode)
        yield
    finally:
        for i, module in enumerate(model.modules()):
            module.training = modes[i]


def eval_mode(model):
    """A context manager that places a model into evaluation mode and restores
    the previous mode on exit."""
    return train_mode(model, False)


def sync_grads(model, enable=True):
    """A context manager that enables or disables gradient synchronization."""
    return nullcontext() if enable else model.no_sync()


@torch.no_grad()
def gradient_norm(params):
    """Computes the 2-norm of the gradients of the given parameters."""
    total = torch.tensor(0.0)
    for p in params:
        if p.grad is not None:
            total = total.to(p.grad.device)
            total += torch.norm(p.grad, dtype=total.dtype) ** 2
    return torch.sqrt(total)


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--grad-accum-steps", type=int, default=1)
    parser.add_argument("--seed", type=int, default=1234)
    args = parser.parse_args()

    du.init_distributed()

    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = du.get_device()

    print0(f"Batch size per process: {args.batch_size}")
    print0(f"World size: {world_size}")
    print0(f"Batch size: {args.batch_size * world_size}")

    if rank == 0:
        datasets.CIFAR10("_data", download=True)
    dist.barrier()

    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2470, 0.2435, 0.2616)),
        ]
    )

    train_set = datasets.CIFAR10("_data", train=True, transform=transform)
    train_sampler = data.DistributedSampler(
        train_set, num_replicas=world_size, rank=rank, seed=args.seed
    )
    train_dl = data.DataLoader(
        train_set,
        batch_size=args.batch_size,
        sampler=train_sampler,
        num_workers=4,
        persistent_workers=True,
        drop_last=True,
    )

    val_set = datasets.CIFAR10("_data", train=False, transform=transform)
    val_sampler = data.DistributedSampler(
        val_set, num_replicas=world_size, rank=rank, seed=args.seed, shuffle=False
    )
    val_dl = data.DataLoader(
        val_set,
        batch_size=args.batch_size,
        sampler=val_sampler,
        num_workers=4,
        persistent_workers=True,
    )

    torch.manual_seed(args.seed)
    model = nn.Sequential(
        nn.Conv2d(3, 32, 3, padding=1),
        nn.GELU(),
        nn.Dropout2d(0.25),
        nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1),
        nn.GELU(),
        nn.Dropout2d(0.25),
        nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 8 * 8, 64),
        nn.GELU(),
        nn.Dropout(0.25),
        nn.Linear(64, 10),
    ).to(device)
    du.broadcast_tensors(model.parameters())

    print0(f"Parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

    dist_kwargs = {}
    if device.type == "cuda":
        dist_kwargs["device_ids"] = [device.index]
        dist_kwargs["output_device"] = device.index
    model = nn.parallel.DistributedDataParallel(model, **dist_kwargs)

    opt = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-1)
    sched = optim.lr_scheduler.CosineAnnealingLR(opt, args.epochs, verbose=rank == 0)

    @torch.no_grad()
    @eval_mode(model)
    def validate():
        print0("Validating...")
        correct = torch.zeros((), dtype=torch.long, device=device)
        total = torch.zeros((), dtype=torch.long, device=device)
        for x, y in tqdm(val_dl, disable=rank != 0):
            x, y = x.to(device), y.to(device)
            logits = model(x)
            correct_batch = (logits.argmax(dim=-1) == y).sum()
            total_batch = torch.tensor(y.numel(), device=device)
            dist.all_reduce_coalesced([correct_batch, total_batch])
            correct += correct_batch
            total += total_batch
        acc = correct / total
        print0(f"Validation accuracy: {acc:.2%}")

    step = 0
    substep = 1
    loss_parts = []

    try:
        for epoch in range(args.epochs):
            train_sampler.set_epoch(epoch)
            for x, y in tqdm(train_dl, disable=rank != 0):
                sync = substep >= args.grad_accum_steps
                x, y = x.to(device), y.to(device)

                with sync_grads(model, sync):
                    logits = model(x)
                    loss = F.cross_entropy(logits, y)

                    loss_part = loss / args.grad_accum_steps
                    loss_part.backward()
                    loss_parts.append(loss_part.detach())
                if sync:
                    opt.step()
                    substep = 0

                    if step % 25 == 0:
                        loss_global = sum(loss_parts) / world_size
                        grad_norm = gradient_norm(model.parameters())
                        dist.all_reduce(loss_global)
                        print0(
                            f"epoch: {epoch}, step: {step}, loss: {loss_global.item():g}, grad norm: {grad_norm.item():g}"
                        )

                    step += 1
                    loss_parts.clear()
                    opt.zero_grad()

                substep += 1

            validate()
            sched.step()

        print0("Done!")
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()

examples/train_classifier.py
crowsonkb/torch-dist-utils
"""Utilities for PyTorch distributed."""

from .torch_dist_utils import (
    get_local_group,
    get_device,
    init_distributed,
    cleanup_distributed,
    do_in_order,
    do_in_local_order,
    on_rank_0,
    on_local_rank_0,
    print0,
    printl0,
    all_gather_object,
    broadcast_object,
    gather_object,
    scatter_objects,
    all_gather_into_new,
    broadcast_tensors,
)

del torch_dist_utils


__all__ = [
    "get_local_group",
    "get_device",
    "init_distributed",
    "cleanup_distributed",
    "do_in_order",
    "do_in_local_order",
    "on_rank_0",
    "on_local_rank_0",
    "print0",
    "printl0",
    "all_gather_object",
    "broadcast_object",
    "gather_object",
    "scatter_objects",
    "all_gather_into_new",
    "broadcast_tensors",
]

__version__ = "0.1.0"

torch_dist_utils/__init__.py
crowsonkb/torch-dist-utils
#!/usr/bin/env python3

"""Test torch_dist_utils.py."""

import argparse
from functools import partial
import os

import torch
import torch.distributed as dist

from . import torch_dist_utils as du

printf0 = partial(du.print0, flush=True)


def assert_equal(actual, expected):
    """Assert that two objects are equal."""
    assert actual == expected, f"{actual} != {expected}"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device-type", required=True, choices=["cpu", "cuda"])
    args = parser.parse_args()

    du.init_distributed()
    printf0("Testing init_distributed...")

    rank = dist.get_rank()
    local_rank = dist.get_rank(du.get_local_group())
    group_rank = int(os.environ.get("GROUP_RANK", "0"))
    world_size = dist.get_world_size()
    printf0(f"World size: {world_size}")

    if dist.is_torchelastic_launched():
        printf0("Testing local rank...")
        assert_equal(local_rank, int(os.environ["LOCAL_RANK"]))

    printf0("Testing do_in_order (please verify that these print in order)...")
    with du.do_in_order():
        print(f"Hello from rank {rank}")

    printf0("Testing do_in_local_order (please verify that these print in local order)...")
    with du.do_in_local_order():
        print(f"Hello from group rank {group_rank}, local rank {local_rank}")

    printf0("Testing printl0 (please verify only local rank 0 prints)...")
    du.printl0(f"Hello from rank {rank}, local rank {local_rank}", flush=True)

    printf0("Testing get_device...")
    device = du.get_device() if args.device_type == "cuda" else torch.device("cpu")
    if args.device_type == "cuda":
        assert_equal(device.index, local_rank)

    torch.set_default_device(device)

    printf0("Testing all_gather_object...")
    actual = du.all_gather_object(rank)
    expected = list(range(world_size))
    assert_equal(actual, expected)

    printf0("Testing broadcast_object...")
    actual = du.broadcast_object(rank)
    expected = 0
    assert_equal(actual, expected)

    printf0("Testing gather_object...")
    actual = du.gather_object(rank)
    expected = list(range(world_size)) if rank == 0 else None
    assert_equal(actual, expected)

    printf0("Testing scatter_objects...")
    actual = du.scatter_objects(list(range(world_size)))
    expected = rank
    assert_equal(actual, expected)

    printf0("Testing all_gather_into_new...")
    tensor = torch.eye(rank + 1)
    actual = du.all_gather_into_new(tensor)
    expected = [torch.eye(i) for i in range(1, world_size + 1)]
    for a, e in zip(actual, expected):
        assert torch.all(a == e), f"{a} != {e}"

    printf0("Testing broadcast_tensors...")
    torch.manual_seed(rank)
    actual = [torch.randn(4, 4), torch.randn(6, 6)]
    du.broadcast_tensors(actual)
    torch.manual_seed(0)
    expected = [torch.randn(4, 4), torch.randn(6, 6)]
    for a, e in zip(actual, expected):
        assert torch.all(a == e), f"{a} != {e}"

    printf0("Testing cleanup_distributed...")
    du.cleanup_distributed()
    assert_equal(du.get_local_group(), None)
    success = False
    try:
        dist.barrier()
    except (RuntimeError, ValueError):
        success = True
    assert success, "dist.barrier() should fail after cleanup_distributed()"


if __name__ == "__main__":
    main()

torch_dist_utils/test.py
crowsonkb/torch-dist-utils
"""Utilities for PyTorch distributed."""

import atexit
from collections import defaultdict
from contextlib import contextmanager
from functools import wraps
import os
from typing import Any, Callable, Iterable, List, Optional

import torch
import torch.distributed as dist


# Type hints

Group = Optional[dist.ProcessGroup]


# State

all_groups = []
local_group = None


def get_local_group() -> dist.ProcessGroup:
    """Get the process group containing only the local processes.

    :returns: The process group containing only the local processes.
    """
    return local_group


def get_device() -> torch.device:
    """Get the device of the current process.

    :returns: The device of the current process.
    """
    if torch.cuda.is_available():
        return torch.device("cuda", dist.get_rank(local_group))
    return torch.device("cpu")


# Initialization and cleanup


def init_distributed():
    """Initialize distributed communication. If the process is not launched with ``torchrun``,
    then assume it is the only process."""
    global local_group
    backend = None if torch.cuda.is_available() else "gloo"
    if dist.is_torchelastic_launched():
        dist.init_process_group(backend)
        all_groups.append(None)
        group_rank = int(os.environ["GROUP_RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
        rank_list = all_gather_object((group_rank, local_rank))
        groups = defaultdict(list)
        for global_rank, (group_rank, local_rank) in enumerate(rank_list):
            groups[group_rank].append((local_rank, global_rank))
        ranks_per_group = [[rank[1] for rank in sorted(group)] for group in groups.values()]
        local_group, new_groups = dist.new_subgroups_by_enumeration(ranks_per_group)
        all_groups.extend(new_groups)
        if torch.cuda.is_available():
            torch.cuda.set_device(get_device())
    else:
        dist.init_process_group(backend, world_size=1, rank=0, store=dist.HashStore())
        all_groups.append(None)
        local_group = dist.new_group([0])
        all_groups.append(local_group)
    atexit.register(cleanup_distributed)


def cleanup_distributed():
    """Clean up distributed communication."""
    global local_group
    atexit.unregister(cleanup_distributed)
    for group in reversed(all_groups):
        dist.destroy_process_group(group)
    all_groups.clear()
    local_group = None


# Context managers and decorators


@contextmanager
def do_in_order(group: Group = None):
    """A context manager that ensures that all processes execute the block in order.

    :param group: The process group. If ``None``, use the default group.
    """
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)
    dist.barrier(group)
    for i in range(world_size):
        if rank == i:
            yield
        dist.barrier(group)


@contextmanager
def do_in_local_order():
    """A context manager that ensures that all local processes execute the block in order."""
    with do_in_order(local_group):
        yield


def on_rank_0(group: Group = None) -> Callable:
    """A decorator that ensures that only process 0 executes the function.

    :param group: The process group. If ``None``, use the default group.
    :returns: A decorator.
    """

    def decorator(fn):
        @wraps(fn)
        def wrapped(*args, **kwargs):
            if dist.get_rank(group) == 0:
                return fn(*args, **kwargs)

        return wrapped

    return decorator


def on_local_rank_0() -> Callable:
    """A decorator that ensures that only the local process 0 executes the function.

    :returns: A decorator.
    """

    def decorator(fn):
        @wraps(fn)
        def wrapped(*args, **kwargs):
            if dist.get_rank(local_group) == 0:
                return fn(*args, **kwargs)

        return wrapped

    return decorator


print0 = on_rank_0()(print)
print0.__doc__ = """A version of ``print()`` that only prints on process 0."""
printl0 = on_local_rank_0()(print)
printl0.__doc__ = """A version of ``print()`` that only prints on local process 0."""

# Transferring objects


def all_gather_object(obj: Any, group: Group = None) -> List[Any]:
    """Gather an object from each process and return a list of gathered objects in all
    processes.

    :param obj: The object to gather.
    :param group: The process group. If ``None``, use the default group.
    :returns: A list of gathered objects.
    """
    object_list = [None] * dist.get_world_size(group)
    dist.all_gather_object(object_list, obj, group=group)
    return object_list


def broadcast_object(obj: Optional[Any] = None, src: int = 0, group: Group = None) -> Any:
    """Broadcast an object from the source process and return the object in all processes.

    :param obj: The object to broadcast. Ignored in processes other than ``src``.
    :param src: The source process.
    :param group: The process group. If ``None``, use the default group.
    :returns: The object broadcasted from the source process.
    """
    object_list = [obj]
    dist.broadcast_object_list(object_list, src=src, group=group)
    return object_list[0]


def gather_object(obj: Any, dst: int = 0, group: Group = None) -> Optional[List[Any]]:
    """Gather an object from each process and return a list of gathered objects in the
    destination process.

    :param obj: The object to gather.
    :param dst: The destination process.
    :param group: The process group. If ``None``, use the default group.
    :returns: A list of gathered objects in the destination process, ``None`` in all other
        processes.
    """
    rank = dist.get_rank(group)
    object_list = [None] * dist.get_world_size(group) if rank == dst else None
    dist.gather_object(obj, object_list, dst=dst, group=group)
    if rank == dst:
        return object_list


def scatter_objects(objs: Optional[List[Any]] = None, src: int = 0, group: Group = None) -> Any:
    """Scatter a list of objects from the source process and return each object in each
    process.

    :param objs: The list of objects to scatter. Ignored in processes other than ``src``.
    :param src: The source process.
    :param group: The process group. If ``None``, use the default group.
    :returns: The object scattered to the current process.
    """
    object_list = [None]
    dist.scatter_object_list(object_list, objs, src=src, group=group)
    return object_list[0]


# Transferring tensors


def all_gather_into_new(tensor: torch.Tensor, group: Group = None) -> List[torch.Tensor]:
    """Gather a tensor from each process and return a list of gathered tensors in all
    processes. Tensors can have different shapes. Tensors must be all on CPU or all on GPU.

    :param tensor: The tensor to gather.
    :param group: The process group. If ``None``, use the default group.
    :returns: A list of gathered tensors.
    """
    shapes = all_gather_object(tensor.shape, group=group)
    if tensor.device.type == "cuda":
        tensors = [torch.empty(shape, dtype=tensor.dtype, device=tensor.device) for shape in shapes]
        dist.all_gather(tensors, tensor, group=group)
    else:
        rank = dist.get_rank(group)
        tensors = []
        for i, shape in enumerate(shapes):
            if i == rank:
                tensors.append(tensor)
            else:
                tensors.append(torch.empty(shape, dtype=tensor.dtype, device=tensor.device))
        handles = [
            dist.broadcast(tensors[i], src=i, group=group, async_op=True)
            for i in range(len(tensors))
        ]
        for handle in handles:
            handle.wait()
    return tensors


def broadcast_tensors(tensors: Iterable[torch.Tensor], src: int = 0, group: Group = None):
    """Broadcast an iterable of tensors from the given source process to all other processes.

    To synchronize a model's parameters in all processes to the versions in process 0:

    .. code-block:: python

        broadcast_tensors(model.parameters())

    :param tensors: The tensors to broadcast.
    :param src: The source process.
    :param group: The process group. If ``None``, use the default group.
    """
    handles = [dist.broadcast(tensor, src=src, group=group, async_op=True) for tensor in tensors]
    for handle in handles:
        handle.wait()

torch_dist_utils/torch_dist_utils.py
crowsonkb/torch-dist-utils
#!/usr/bin/env python3

"""Launch a PyTorch distributed job on a Slurm cluster."""

import os
from subprocess import Popen, PIPE
import sys


def main():
    try:
        nnodes = os.environ["SLURM_NPROCS"]
        node_rank = os.environ["SLURM_PROCID"]
        nodelist = os.environ["SLURM_JOB_NODELIST"]
    except KeyError:
        print(
            "This script must be run within a Slurm job. The environment variables "
            "SLURM_NPROCS, SLURM_PROCID, and SLURM_JOB_NODELIST were not found.",
            file=sys.stderr,
        )
        sys.exit(1)

    proc = Popen(["scontrol", "show", "hostnames", nodelist], stdout=PIPE)
    master_addr = proc.stdout.read().decode().splitlines()[0]

    command = [
        "torchrun",
        "--nnodes",
        nnodes,
        "--nproc-per-node",
        "gpu",
        "--node_rank",
        node_rank,
        "--master_addr",
        master_addr,
        *sys.argv[1:],
    ]
    Popen(command).wait()


if __name__ == "__main__":
    main()

torch_dist_utils/torchrun_slurm.py
crowsonkb/torch-dist-utils