def _collect_state_dict(
    module: torch.nn.Module,
    state_dict: Optional[StateDictType],
    recurse: bool = True,
    prefix: str = "",
) -> Tuple[StateDictType, List[str], List[str]]:
    """
    Collect a module's state dict across distributed processes.

    Returns the syncronized state dictionary, which will always be a valid state dict,
    and then the missing and unexpected keys corresponding to the original `state_dict`.
    Parameters that missing from the original `state_dict` will be populated from the
    corresponding parameter in the primary processes' module's state dict.

    !!! Note

        `missing_keys` and `unexpected_keys` are only populated in the primary process.
    """
    # This is the device we'll use for the broadcast operation.
    dist_device = distributed_device()
    # We'll keep tensors on CPU in the returned state dict.
    state_dict_device = int_to_device(-1)

    missing_keys: List[str] = []
    unexpected_keys: List[str] = []

    # Gather current state dict and prepare to iterator over it.
    # We iterate over this state dict instead of `state_dict` so we can be sure
    # that the order is consistent across processes.
    # We'll also update this state dict as we go and return it at the end.

    if recurse:
        current_state_dict = module.state_dict()
    else:
        # Only collect state of direct members, including both parameters and buffers.
        current_state_dict = OrderedDict(
            chain(
                # Paramaters
                ((n, p.data) for (n, p) in module.named_parameters(recurse=False)),
                # Buffers
                module.named_buffers(recurse=False),
            )
        )

    keys = list(current_state_dict.keys())

    # Gather unexpected_keys.
    if is_global_primary():
        assert state_dict is not None
        for key in state_dict:
            if key not in keys:
                unexpected_keys.append(key)

    for key in keys:
        tensor = current_state_dict[key]
        if is_global_primary():
            assert state_dict is not None
            if key in state_dict:
                # Update `tensor` to the value in `state_dict`.
                tensor = state_dict[key]
            else:
                missing_keys.append(key)
        logger.debug("Broadcasting distributed parameter '%s'", prefix + key)
        tensor = tensor.to(dist_device).contiguous()
        dist.broadcast(tensor, 0)
        current_state_dict[key] = tensor.to(state_dict_device)

    return current_state_dict, missing_keys, unexpected_keys
Collect a module's state dict across distributed processes.

Returns the syncronized state dictionary, which will always be a valid state dict,
and then the missing and unexpected keys corresponding to the original `state_dict`.
Parameters that missing from the original `state_dict` will be populated from the
corresponding parameter in the primary processes' module's state dict.

!!! Note

`missing_keys` and `unexpected_keys` are only populated in the primary process.
allennlp/nn/util.py