opake.distributed

Distributed training utilities for differential privacy with DDP.

Overview

The opake.distributed module provides composable primitives for multi-GPU DP training:

• Core: is_distributed(), get_rank(), get_world_size()
• Collectives: all_reduce() (+ in-place variant in the collectives submodule)
• Gradient aggregation: sum_gradients() (copy-returning) and sum_gradients_() (in-place)
• State sync: sync() (type-dispatched; handles clipping + noise states and registered DP runtime objects)
• Sharding: local_shard()

DDP is the only supported parallelism strategy. See User Guide: Distributed Training for usage.

Core Utilities

is_distributed()

Return True if torch.distributed is available and initialized.

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def is_distributed() -> bool:
    """Return True if ``torch.distributed`` is available and initialized."""
    return dist.is_available() and dist.is_initialized()

get_rank()

Return the current rank (0 if not distributed).

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def get_rank() -> int:
    """Return the current rank (0 if not distributed)."""
    return dist.get_rank() if is_distributed() else 0

get_world_size()

Return the world size (1 if not distributed).

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def get_world_size() -> int:
    """Return the world size (1 if not distributed)."""
    return dist.get_world_size() if is_distributed() else 1

all_reduce(tensor, op='sum')

Return a reduced clone of tensor; input is unchanged.

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def all_reduce(tensor: torch.Tensor, op: str = "sum") -> torch.Tensor:
    """Return a reduced clone of ``tensor``; input is unchanged."""
    reduced = tensor.clone()
    all_reduce_(reduced, op=op)
    return reduced

all_reduce_(tensor, op='sum')

All-reduce tensor across ranks in place.

Raises:

Type	Description
RuntimeError	If torch.distributed is not initialized.
ValueError	If op is not a recognized reduction.

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def all_reduce_(tensor: torch.Tensor, op: str = "sum") -> None:
    """All-reduce ``tensor`` across ranks in place.

    Raises:
        RuntimeError: If ``torch.distributed`` is not initialized.
        ValueError: If ``op`` is not a recognized reduction.
    """
    reduce_op = _resolve_op(op)
    if not is_distributed():
        raise OperationError(
            *(
                "torch.distributed is not initialized. "
                "Call torch.distributed.init_process_group() first.",
            )
        )
    dist.all_reduce(tensor, op=reduce_op)

barrier()

Block until every rank reaches this call (no-op if not distributed).

Source code in packages/opake-engine/src/opake/api/engine/distributed/collectives.py

def barrier() -> None:
    """Block until every rank reaches this call (no-op if not distributed)."""
    if is_distributed():
        dist.barrier()

Gradient Aggregation

sum_gradients(gradients)

DP-specific alias for reduce_pytree(op="sum").

Source code in packages/opake-engine/src/opake/api/engine/distributed/gradients.py

def sum_gradients(gradients: Any) -> Any:
    """DP-specific alias for ``reduce_pytree(op="sum")``."""
    return reduce_pytree(gradients, op="sum")

sum_gradients_(gradients)

DP-specific alias for reduce_pytree_(op="sum").

Source code in packages/opake-engine/src/opake/api/engine/distributed/gradients.py

def sum_gradients_(gradients: Any) -> None:
    """DP-specific alias for ``reduce_pytree_(op="sum")``."""
    reduce_pytree_(gradients, op="sum")

reduce_pytree(pytree, op='sum')

Return a pytree with each tensor leaf reduced; input unchanged.

When passed ClippedPytree or NoisedPytree, preserves and updates the wrapper metadata for supported sum and mean reductions. Paired second-moment wrappers recurse into both child streams.

Source code in packages/opake-engine/src/opake/api/engine/distributed/gradients.py

def reduce_pytree(pytree: Any, op: str = "sum") -> Any:
    """Return a pytree with each tensor leaf reduced; input unchanged.

    When passed ``ClippedPytree`` or ``NoisedPytree``, preserves and updates the
    wrapper metadata for supported ``sum`` and ``mean`` reductions.  Paired
    second-moment wrappers recurse into both child streams.
    """
    if isinstance(pytree, SecondMomentClippingOutput):
        return SecondMomentClippingOutput(
            reduce_pytree(pytree.grads, op=op),
            reduce_pytree(pytree.squared_grads, op=op),
        )

    if isinstance(pytree, SecondMomentNoiseOutput):
        return SecondMomentNoiseOutput(
            reduce_pytree(pytree.noisy_grads, op=op),
            reduce_pytree(pytree.noisy_squared_grads, op=op),
        )

    if isinstance(pytree, ClippedPytree):
        _assert_wrapper_reduction_supported(pytree, op)
        _assert_public_metadata_equal(pytree.max_norm, name="ClippedPytree.max_norm")
        if _is_noised(pytree):
            _assert_public_metadata_equal(
                pytree.noise_stddev,
                name="NoisedPytree.noise_stddev",
            )
        reduced = pytree.clone()
        reduce_pytree_(reduced.pytree, op=op)
        return _reduced_metadata(reduced, op, _metadata_world_size())

    def _clone(leaf: Any) -> Any:
        if isinstance(leaf, torch.Tensor):
            return leaf.clone()
        raise InputTypeError(
            *(
                f"reduce_pytree expects tensor leaves after wrapper dispatch; "
                f"got {type(leaf).__name__}. Unwrap paired/custom containers "
                f"explicitly or register a reduction branch.",
            )
        )

    reduced = tree_map(_clone, pytree)
    reduce_pytree_(reduced, op=op)
    return reduced

reduce_pytree_(pytree, op='sum')

All-reduce tensor leaves in place.

In-place wrapper reductions are accepted only when the wrapper metadata stays unchanged. Use :func:reduce_pytree for reductions such as noised sum or clipped/noised mean that need updated metadata.

Source code in packages/opake-engine/src/opake/api/engine/distributed/gradients.py

def reduce_pytree_(pytree: Any, op: str = "sum") -> None:
    """All-reduce tensor leaves in place.

    In-place wrapper reductions are accepted only when the wrapper metadata
    stays unchanged.  Use :func:`reduce_pytree` for reductions such as noised
    ``sum`` or clipped/noised ``mean`` that need updated metadata.
    """
    if isinstance(pytree, SecondMomentClippingOutput):
        reduce_pytree_(pytree.grads, op=op)
        reduce_pytree_(pytree.squared_grads, op=op)
        return

    if isinstance(pytree, SecondMomentNoiseOutput):
        reduce_pytree_(pytree.noisy_grads, op=op)
        reduce_pytree_(pytree.noisy_squared_grads, op=op)
        return

    if isinstance(pytree, ClippedPytree):
        _assert_wrapper_reduction_supported(pytree, op)
        if _in_place_wrapper_metadata_changes(pytree, op):
            raise InputTypeError(
                *(
                    f"In-place {type(pytree).__name__} reduction would change metadata; "
                    "use reduce_pytree() instead.",
                )
            )
        _assert_public_metadata_equal(pytree.max_norm, name="ClippedPytree.max_norm")
        if _is_noised(pytree):
            _assert_public_metadata_equal(
                pytree.noise_stddev,
                name="NoisedPytree.noise_stddev",
            )
        reduce_pytree_(pytree.pytree, op=op)
        return

    if not is_distributed():
        return

    def _reduce(leaf: Any) -> Any:
        if isinstance(leaf, torch.Tensor):
            all_reduce_(leaf, op=op)
            return leaf
        raise InputTypeError(
            *(
                f"reduce_pytree_ expects tensor leaves after wrapper dispatch; "
                f"got {type(leaf).__name__}. Unwrap paired/custom containers "
                f"explicitly or register a reduction branch.",
            )
        )

    tree_map(_reduce, pytree)

State Synchronization

sync(*states)

Synchronize one or more registered state/aux objects across ranks.

Dispatches on object type — exact match first, then __mro__ — to whichever function was registered via :func:register_sync_type. An unregistered type raises TypeError; nothing is passed through unsynchronized.

Returns:

Type	Description
Any	A single synchronized object if one argument was passed, otherwise a
Any	tuple matching the input order.

Source code in packages/opake-engine/src/opake/api/engine/distributed/_state.py

def sync(*states: Any) -> Any:
    """Synchronize one or more registered state/aux objects across ranks.

    Dispatches on object type — exact match first, then ``__mro__`` — to
    whichever function was registered via :func:`register_sync_type`.  An
    unregistered type raises ``TypeError``; nothing is passed through
    unsynchronized.

    Returns:
        A single synchronized object if one argument was passed, otherwise a
        tuple matching the input order.
    """

    def _sync_one(single: Any) -> Any:
        state_type = type(single)
        sync_fn = _resolve_sync_fn(state_type)
        if sync_fn is None:
            _ensure_builtin_sync_types_loaded()
            sync_fn = _resolve_sync_fn(state_type)
        if sync_fn is not None:
            return sync_fn(single)
        raise InputTypeError(
            *(
                f"No sync function registered for {state_type.__name__}. "
                f"Registered types: {[t.__name__ for t in _SYNC_REGISTRY]}",
            )
        )

    if len(states) == 1:
        return _sync_one(states[0])
    if len(states) > 1:
        return tuple(_sync_one(s) for s in states)
    return ()

The sync() machinery is type-dispatched: clipping and noise states register themselves with opake.distributed._state.register_sync_type and provide the right per-state aggregation rule. Lower-level scalar reductions, tensor gathers, and object syncs live in _state.py; they're internal plumbing for the registered DP runtime objects rather than headline API.

Privacy Ordering

The order of operations matters for DP guarantees:

grads, clip_state = grad_fn(params, x, y, state=clip_state)  # 1. Clip
grads = dist_utils.sum_gradients(grads)                       # 2. Aggregate (copy)
noisy_grads, noise_state = noise_fn(grads, noise_state)       # 3. Noise

sync accepts one or more objects and returns synchronized values in the same order:

clip_state, aux = dist_utils.sync(clip_state, aux)
noise_state = dist_utils.sync(noise_state)

See examples/distributed_dp_training.py for a complete working script.