def _load_raw_module_state(
    state_dict: dict[str, Any], module: Module, world_size: int = 1, strict: bool = True
) -> None:
    """Loads the state dict into the module by gathering all weights first and then and writing back to each shard."""
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    if _has_dtensor_modules(module):
        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict

        state_dict_options = StateDictOptions(
            broadcast_from_rank0=True,
            full_state_dict=True,
            # must be set False to allow loading each param separately below
            strict=False,
        )

        for submodule_name, submodule in module.named_modules():
            for param_name, _ in _named_parameters_and_buffers_to_load(submodule):
                full_param_name = f"{submodule_name}{'.' if submodule_name else ''}{param_name}"
                if full_param_name not in state_dict:
                    if not strict:
                        continue
                    raise KeyError(
                        f"The model contains a key '{full_param_name}' that does not exist in the loaded checkpoint."
                        " To disable strict loading, set `strict=False`."
                    )
                local_state_dict = {param_name: state_dict[full_param_name]}
                set_model_state_dict(submodule, local_state_dict, options=state_dict_options)

    elif isinstance(module, FSDP):
        with _get_full_state_dict_context(module, world_size=world_size, rank0_only=False):
            module.load_state_dict(state_dict, strict=strict)
    else:
        module.load_state_dict(state_dict, strict=strict)
Loads the state dict into the module by gathering all weights first and then and writing back to each shard.
src/lightning/fabric/strategies/model_parallel.py