def materialize_module(
    module: nn.Module,
) -> Generator[nn.Module, None, None]:
    """Context manager for initializing modules containing DTensor parameters.

    This context manager temporarily materializes all DTensor parameters in a module
    as full tensors for initialization, then redistributes the initialized values back
    to the original DTensor parameters.

    Args:
        module: The module whose parameters should be materialized.

    Yields:
        nn.Module: The module with materialized parameters that can be initialized.
    """
    param_placement_map: dict[tuple[nn.Module, str],
                              tuple[DeviceMesh, tuple[Placement, ...]]] = {}
    with torch.no_grad():
        for submodule in module.modules():
            for name, param in submodule.named_parameters(recurse=False):
                if isinstance(param, DTensor):
                    param_placement_map[
                        (submodule,
                         name)] = (param.device_mesh, param.placements)
                    setattr(submodule, name, nn.Parameter(param.full_tensor()))

        yield module

        for submodule in module.modules():
            for name, param in submodule.named_parameters(recurse=False):
                if (submodule, name) in param_placement_map:
                    device_mesh, placements = param_placement_map[
                        (submodule, name)]
                    setattr(
                        submodule,
                        name,
                        nn.Parameter(
                            distribute_tensor(
                                param,
                                device_mesh=device_mesh,
                                placements=placements,
                            ),
                        ),
                    )
Context manager for initializing modules containing DTensor parameters.

    This context manager temporarily materializes all DTensor parameters in a module
    as full tensors for initialization, then redistributes the initialized values back
    to the original DTensor parameters.

    Args:
        module: The module whose parameters should be materialized.

    Yields:
        nn.Module: The module with materialized parameters that can be initialized.
    
llmfoundry/models/utils/param_init_fns.py