CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
fsdp.py54 linesDownload Raw Back to integrations
1# Copyright 2024 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from __future__ import annotations15 16import os17from typing import TYPE_CHECKING18 19from ..utils import is_torch_available, strtobool20 21 22if TYPE_CHECKING:23    from torch import nn24 25 26def is_fsdp_managed_module(module: nn.Module) -> bool:27    if not is_torch_available():28        return False29 30    import torch31 32    if not torch.distributed.is_available():33        return False34 35    import torch.distributed.fsdp36 37    return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr(38        module, "_is_fsdp_managed_module", False39    )40 41 42def is_fsdp_enabled():43    if is_torch_available():44        import torch45 46        return (47            torch.distributed.is_available()48            and torch.distributed.is_initialized()49            and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 150            and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 151        )52 53    return False54 
Aluode/PerceptionLabPortable · CoolFace