def load_state_dict_distributed(
    module: torch.nn.Module,
    state_dict: Optional[StateDictType],
    strict: bool = True,
    prefix: str = "",
) -> _IncompatibleKeys:
    """
    Load a `state_dict` to the `module` within a distributed process. Only the global
    primary process requires the `state_dict` to not be `None`. All other processes
    will have the state tensors broadcasted to them one-by-one.

    If `strict` is `True`, then the keys of `state_dict` must exactly match the keys
    returned by `module.state_dict()`.

    !!! Note
        The returned `missing_keys` and `unexpected_keys` will only be accurate
        in the primary process.

    # Returns

    A `NamedTuple` with `missing_keys` and `unexpected_keys` fields, both of which
    are lists of strings.

    # Raises

    `RuntimeError`
        If `strict` is `True` and there are missing or unexpected keys.

    """
    if not is_distributed():
        return module.load_state_dict(state_dict, strict=strict)

    if is_global_primary():
        assert state_dict is not None
    else:
        assert state_dict is None

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

    submodules = dict(module.named_children())

    def update_key_list(original, updates):
        for key in updates:
            if key not in original:
                original.append(key)

    from allennlp.nn.parallel.sharded_module_mixin import ShardedModuleMixin

    # If we've found a sharded module or there aren't any more submodules of the current module,
    # we collect the state_dict and load it now instead of recursing further.
    if isinstance(module, ShardedModuleMixin) or not submodules:
        # Collect.
        state_dict, _missing_keys, _unexpected_keys = _collect_state_dict(
            module, state_dict, prefix=prefix
        )
        assert state_dict is not None
        update_key_list(missing_keys, _missing_keys)
        update_key_list(unexpected_keys, _unexpected_keys)
        # And load.
        _missing_keys, _unexpected_keys = module.load_state_dict(state_dict, strict=False)
        update_key_list(missing_keys, _missing_keys)
        update_key_list(unexpected_keys, _unexpected_keys)
    else:
        # We'll recursively call this function on each submodule, but first we need
        # to collect any parameters that are direct members of this module.
        direct_member_state_dict, _missing_keys, _unexpected_keys = _collect_state_dict(
            module,
            state_dict,
            recurse=False,
            prefix=prefix,
        )
        # `_unexpected_keys` will contain keys corresponding to submodules (not direct members)
        # that may be legitimate, so we ignore any `_unexpected_keys` here that correspond to submodules.
        _unexpected_keys = [
            k for k in _unexpected_keys if "." not in k or k.split(".")[0] not in submodules.keys()
        ]
        update_key_list(missing_keys, _missing_keys)
        update_key_list(unexpected_keys, _unexpected_keys)

        # `_missing_keys` here will contain any keys corresponding to submodules, but
        # we'll remove those below.
        _missing_keys, _unexpected_keys = module.load_state_dict(
            direct_member_state_dict, strict=False
        )
        update_key_list(missing_keys, _missing_keys)
        update_key_list(unexpected_keys, _unexpected_keys)

        # Okay, now for the recursive part.
        for name, submodule in submodules.items():
            # Update `missing_keys` to remove keys corresponding to this submodule.
            # If they are actually missing after this step, we add them back in below.
            missing_keys = [k for k in missing_keys if not k.startswith(name + ".")]
            submodule_state_dict: Optional[StateDictType] = None
            if is_global_primary():
                assert state_dict is not None
                submodule_state_dict = {
                    key.replace(name + ".", "", 1): value
                    for key, value in state_dict.items()
                    if key.startswith(name + ".")
                }
            _missing_keys, _unexpected_keys = load_state_dict_distributed(
                submodule,
                submodule_state_dict,
                strict=False,
                prefix=prefix + name + ".",
            )
            update_key_list(missing_keys, [f"{name}.{key}" for key in _missing_keys])
            update_key_list(unexpected_keys, [f"{name}.{key}" for key in _unexpected_keys])

    _check_incompatible_keys(module, missing_keys, unexpected_keys, strict)

    return _IncompatibleKeys(missing_keys, unexpected_keys)
Load a `state_dict` to the `module` within a distributed process. Only the global
primary process requires the `state_dict` to not be `None`. All other processes
will have the state tensors broadcasted to them one-by-one.

If `strict` is `True`, then the keys of `state_dict` must exactly match the keys
returned by `module.state_dict()`.

!!! Note
The returned `missing_keys` and `unexpected_keys` will only be accurate
in the primary process.

# Returns

A `NamedTuple` with `missing_keys` and `unexpected_keys` fields, both of which
are lists of strings.

# Raises

`RuntimeError`
If `strict` is `True` and there are missing or unexpected keys.
allennlp/nn/util.py