HierarchicalModelAverager

class tensorplay.distributed.algorithms.model_averaging.HierarchicalModelAverager(period_group_size_dict=None, warmup_steps=0, process_group=None)[source]

Runs hierarchical model averaging (hierarchical SGD).

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 that supports post-local SGD, 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(). 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) – 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.

average_parameters(params: Iterable)[source]

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.