As the code example below shows, only a few lines need to be added to the setup of DDP model, and the training loop code can keep untouched. As explained previously, hierarchical SGD is an extended form of local SGD, so the enablement can be quite similar to local SGD (see PyTorch docs of PostLocalSGDOptimizer):

Register a post-local SGD communication hook to run a warmup stage of fully synchronous SGD and defer hierarchical SGD.
Create a post-local SGD optimizer that wraps an existing local optimizer and a hierarchical SGD configuration.

```
import torch.distributed.algorithms.model_averaging.hierarchical_model_averager as hierarchicalSGD
from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import (
    PostLocalSGDState,
    post_localSGD_hook,
)
from torch.distributed.optim import PostLocalSGDOptimizer
ddp_model = nn.parallel.DistributedDataParallel(
    module=model,