Implementation:Pyro ppl Pyro HorovodOptimizer

Property	Value
Module	pyro.optim.horovod
Source	pyro/optim/horovod.py
Lines	56
Classes	HorovodOptimizer
Parent Class	pyro.optim.optim.PyroOptim
Dependencies	pyro.optim.optim.PyroOptim, horovod.torch (optional)

Overview

HorovodOptimizer provides a distributed training wrapper for Pyro optimizers using the Horovod framework. It wraps a PyroOptim object similarly to how horovod.torch.DistributedOptimizer wraps a standard PyTorch optimizer, enabling data-parallel training across multiple workers.

The key design is a custom optim_constructor that:

1. Creates the underlying PyTorch optimizer from the wrapped PyroOptim.
2. Wraps it with hvd.DistributedOptimizer which handles gradient averaging across workers.
3. Provides named parameters by looking them up in the Pyro param store.

Parameters are sorted by name before processing to ensure deterministic ordering across all distributed workers.

Code Reference

Class: HorovodOptimizer

Constructor:

• pyro_optim (PyroOptim): A Pyro optimizer instance to distribute.
• **horovod_kwargs: Extra keyword arguments passed to hvd.DistributedOptimizer.

Methods:

• __call__(params, *args, **kwargs): Sorts parameters by name (for deterministic cross-worker ordering) before delegating to the parent PyroOptim.__call__.

I/O Contract

Method	Input	Output
__init__	pyro_optim: PyroOptim, extra Horovod kwargs	HorovodOptimizer instance
__call__	params: List/ValuesView of Tensors	None (updates parameters in-place)

Usage Examples

import pyro
from pyro.optim import Adam
from pyro.optim.horovod import HorovodOptimizer

# Note: Requires horovod to be installed (pip install pyro-ppl[horovod])
# and should be run with horovodrun

# Create a base Pyro optimizer
base_optim = Adam({"lr": 0.01})

# Wrap with Horovod for distributed training
distributed_optim = HorovodOptimizer(base_optim)

# Use in SVI as usual
svi = pyro.infer.SVI(model, guide, distributed_optim, loss=pyro.infer.Trace_ELBO())

Related Pages

• Pyro_ppl_Pyro_MultiOptimizer -- Higher-order optimizer interface
• Pyro_ppl_Pyro_PyroLRScheduler -- Learning rate scheduling