def _clean_path(path: str) -> str:
    """Remove FSDP related wrapper modules from a given state dict key str path."""
    return ".".join([split for split in path.split(".") if split not in {"_fsdp_wrapped_module", "_fpw_module"}])