Method _root_copy_hook

torch/nn/parallel/distributed.py:968–1006

For DDP mixed precision, put low precision copies on separate stream and create events to wait for them. When training with DDP mixed precision, this root pre-forward hook kicks off low precision copies on a separate stream and creates respective events to wait for

(self, *args: Any, **kwargs: Any)

        self.reducer._autograd_hook(idx)  # type: ignore[attr-defined]

    def _root_copy_hook(self, *args: Any, **kwargs: Any) -> None:
        """
        For DDP mixed precision, put low precision copies on separate stream and create events to wait for them.

        When training with DDP mixed precision, this root pre-forward hook kicks
        off low precision copies on a separate stream and creates respective
        events to wait for them.
        """
        # Clear out previous iteration submodule to event. This is because we
        # may have populated some events for modules that didn't end up being
        # used.
        self._submodule_to_event = defaultdict(deque)  # type: ignore[var-annotated]
        with torch.cuda.stream(self._mp_stream):
            for submodule in self.module.modules():
                for param in submodule.parameters(recurse=False):
                    # Do not cast DDP ignored parameters.
                    if hasattr(param, "_ddp_ignored") and param._ddp_ignored:
                        continue
                    _alloc_storage(param._mp_param, param.size())
                    # copy() implicitly casts to low precision
                    with torch.no_grad():
                        param._mp_param.copy_(param.data)
                        # TODO: when zero_grad(set_to_none=False) or in grad
                        # accumulation case, accumulated grads can be in fp32
                        # which can cause errors when running DDP backwards due
                        # to mismatched incoming and accumulated gradient types.
                        # So we manually cast the accumulated grad down for now,
                        # in the future we may shift to FSDP style gradient
                        # accumulation management where the accumulated gradient
                        # is saved and .grad field is set to None, bypassing
                        # this issue.
                        if param.grad is not None:
                            param.grad.data = param.grad.to(
                                self.mixed_precision.param_dtype  # type: ignore[union-attr]
                            )
                    param.data = param._mp_param
                copy_event = torch.cuda.Event()
                copy_event.record()
                self._submodule_to_event[submodule].append(copy_event)

    def _module_wait_for_copy_hook(
        self,