PeriodicModelAverager

class tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager(period, warmup_steps=0, process_group: ProcessGroup | None = None)[source]

Averages parameters periodically after the warm-up stage.

This can be used for running post-local SGD, by running DistributedDataParallel (DDP) using the subgroups created by new_subgroups().

Parameters:

• period (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) – 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())

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

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.