Post Local SGD Hook (ddp_comm_hooks/post_localSGD_hook.py)

Implements a distributed training communication hook that transitions from global to local all-reduces after a specified iteration. Use to optimize distributed scaling in PyTorch.

# Source: https://github.com/pytorch/pytorch | License: BSD-3-Clause | Copyright 2016-     Facebook, Inc            (Adam Paszke)
# Provided via justclaw.it.com; keep this notice with the code.
# mypy: allow-untyped-defs
import logging

import torch
import torch.distributed as dist

from . import default_hooks as default


logger = logging.getLogger(__name__)


class PostLocalSGDState:
    r"""
    Store state for all-reducing gradients globally until given step, then locally after.

    Stores the state for all-reducing gradients globally using ``process_group`` until step ``start_localSGD_iter``,
    and all-reducing gradients locally using ``subgroup`` afterwards.

    If ``process_group`` is ``None``, the global process group will be used.
    If ``subgroup`` is ``None``, the intra-nod

... (truncated -- full source via MCP)