def _get_fqn_to_fsdp_param_info(model: nn.Module) -> Dict[str, FSDPParamInfo]:
    """
    Construct the mapping from a param's fqn to its corresponding ``FSDPParamInfo``
    if the param is managed by FSDP. ``FlatParameter._fqns`` only stores the first
    FQN of a shared parameter. So the keys in the mapping are guaranteed to map
    to unique parameters.
    """

    def module_fn(module, prefix, fqn_to_param_info):
        fsdp_state = _get_module_fsdp_state_if_fully_sharded_module(module)
        if fsdp_state is None:
            return
        _lazy_init(fsdp_state, module)
        handles = _module_handles(fsdp_state, module)
        if not handles:
            return
        flat_param = handles[0].flat_param
        fsdp_param_info = FSDPParamInfo(fsdp_state, flat_param, {})
        for idx, local_fqn in enumerate(flat_param._fqns):
            fqn = clean_tensor_name(prefix + local_fqn)
            if fqn in fqn_to_param_info:
                assert fqn_to_param_info[fqn].flat_param == flat_param
            fqn_to_param_info[fqn] = fsdp_param_info
            fsdp_param_info.param_indices[fqn] = idx

    def return_fn(fqn_to_param_info):
        return fqn_to_param_info

    fqn_to_param_info: Dict[str, FSDPParamInfo] = {}
    # FlatParameter._fqns stores the local fqn, starting from the root of the
    # FSDP. Using _apply_to_modules() with model (may not be the FSDP root
    # module) allows us to construct the global fqn.
    return _apply_to_modules(
        model,
        module_fn,
        return_fn,
        [fqn for fqn, _ in model.named_parameters()],
        fqn_to_param_info,
    )
Construct the mapping from a param's fqn to its corresponding ``FSDPParamInfo``
if the param is managed by FSDP. ``FlatParameter._fqns`` only stores the first
FQN of a shared parameter. So the keys in the mapping are guaranteed to map
to unique parameters.
_optim_utils.py