# tensorplay.distributed.algorithms.model_averaging API

Source: https://www.tensorplay.cn/docs/api/tensorplay.distributed.algorithms.model_averaging.html

## Classes 3

[#](#api-tensorplay.distributed.algorithms.model_averaging.HierarchicalModelAverager)

### HierarchicalModelAverager

class[Full reference ↗](/docs/generated/tensorplay.distributed.algorithms.model_averaging.HierarchicalModelAverager.html)

```python
class tensorplay.distributed.algorithms.model_averaging.HierarchicalModelAverager(period_group_size_dict=None, warmup_steps=0, process_group=None)
```

Runs hierarchical model averaging ([hierarchical SGD](https://arxiv.org/pdf/2010.12998.pdf)).

Process groups of different sizes are organized in a hierarchy, and they average parameters
by using different periods concurrently after the warm-up stage.
This is an extension of [PeriodicModelAverager](/docs/generated/tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager.html#tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager)
that supports [post-local SGD](https://arxiv.org/abs/1808.07217), which essentially only supports
a two-level hierarchy: the intra-machine level and the global level, where the intra-machine
level is usually embedded in [post_localSGD_hook()](/docs/distributed.html#module-tensorplay.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook).
Similarly, the process groups within this class do not have such an intra-machine process
subgroup, which should be embedded by the post-local SGD communication hook instead.

Parameters:

- period_group_size_dict – An ordered dict mapping keys of model averaging period to process group size, used for initializing process groups of different sizes in a hierarchy to average parameters concurrently.

- warmup_steps ([int](https://docs.python.org/3/builtins/functions.html#int)) – The number of warm-up steps. During this stage, model averaging is skipped.

- process_group (ProcessGroup, optional) – The overall process group containing all the processes that runs model averaging.

> **Warning**
>
> The last group size in the dict must be the size of the provided process_group, which indicates model averaging at the highest level of the hierarchy. If process_group is not provided, then the last group size should be equal to the world size.

> **Warning**
>
> HierarchicalModelAverager is experimental and subject to change.

```python
average_parameters(params: Iterable)
```

Averages parameters or parameter groups of an optimizer.

Averaging only occurs if step is no less than warmup_steps
and it can be divided by a period in the keys of period_process_group_dict,
where step is increased by 1 at each iteration in the training loop.
If step can be divided by multiple periods in the keys of period_process_group_dict,
only the largest period is used, and the corresponding process group is used for averaging parameters.
:param params: The parameters of a model or parameter groups of an optimizer.

[#](#api-tensorplay.distributed.algorithms.model_averaging.ModelAverager)

### ModelAverager

class[Full reference ↗](/docs/generated/tensorplay.distributed.algorithms.model_averaging.ModelAverager.html)

```python
class tensorplay.distributed.algorithms.model_averaging.ModelAverager(process_group: ProcessGroup | None = None)
```

Base class for all model averagers.

Parameters:

process_group – The process group to be used for all-reduce.
If None, the default process group, which
is created by tensorplay.distributed.init_process_group(),
will be used. (default: None)

[#](#api-tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager)

### PeriodicModelAverager

class[Full reference ↗](/docs/generated/tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager.html)

```python
class tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager(period, warmup_steps=0, process_group: ProcessGroup | None = None)
```

Averages parameters periodically after the warm-up stage.

This can be used for running [post-local SGD](https://arxiv.org/abs/1808.07217),
by running DistributedDataParallel (DDP)
using the subgroups created by new_subgroups().

Parameters:

- period ([int](https://docs.python.org/3/builtins/functions.html#int)) – The number of steps per model averaging. Usually the period should be greater than 1 to reduce the communication cost. Otherwise, only DDP needs to be used.

- warmup_steps ([int](https://docs.python.org/3/builtins/functions.html#int)) – The number of warm-up steps. During this stage, model averaging is skipped.

- process_group – The process group to be used for all-reduce. If None, the default process group, which is created by tensorplay.distributed.init_process_group(), will be used. (default: None)

Example:

```
>>> # xdoctest: +SKIP("undefined variables")
>>> averager = averagers.PeriodicModelAverager(period=4, warmup_steps=100)
>>> for step in range(0, 200):
>>>    optimizer.zero_grad()
>>>    loss = loss_fn(output, labels)
>>>    loss.backward()
>>>    optimizer.step()
>>>    # Will average model parameters globally every 4 steps.
>>>    averager.average_parameters(model.parameters())
```

```python
average_parameters(params: Iterable[Parameter] | Iterable[dict[str, Parameter]])
```

Averages parameters or parameter groups of an optimizer if step is no less than warmup_steps.

Can be divided by period, where step is increased by 1
at each iteration in the training loop.
:param params: The parameters of a model or parameter groups of an optimizer.