Moved is_fsdp_available to integrations.fsdp

Author: romitjain

src/transformers/cache_utils.py

@@ -7,8 +7,8 @@
 from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_6
 
 from .configuration_utils import PretrainedConfig
+from .integrations import is_fsdp_enabled
 from .utils import (
-    is_fsdp_enabled,
     is_hqq_available,
     is_quanto_greater,
     is_torch_greater_or_equal,

src/transformers/integrations/__init__.py

@@ -55,7 +55,7 @@
     "eetq": ["replace_with_eetq_linear"],
     "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"],
     "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"],
-    "fsdp": ["is_fsdp_managed_module"],
+    "fsdp": ["is_fsdp_managed_module", "is_fsdp_enabled"],
     "ggml": [
         "GGUF_CONFIG_MAPPING",
         "GGUF_TOKENIZER_MAPPING",

src/transformers/integrations/fsdp.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
-from ..utils import is_torch_available
+from ..utils import is_torch_available, strtobool
 
 
 if TYPE_CHECKING:
@@ -36,3 +37,17 @@ def is_fsdp_managed_module(module: nn.Module) -> bool:
     return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr(
         module, "_is_fsdp_managed_module", False
     )
+
+
+def is_fsdp_enabled():
+    if is_torch_available():
+        import torch
+
+        return (
+            torch.distributed.is_available()
+            and torch.distributed.is_initialized()
+            and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1
+            and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1
+        )
+
+    return False

src/transformers/modeling_utils.py

@@ -54,7 +54,7 @@
 from .distributed import DistributedConfig
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
-from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled
+from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled, is_fsdp_enabled
 from .integrations.accelerate import find_tied_parameters, init_empty_weights
 from .integrations.deepspeed import _load_state_dict_into_zero3_model
 from .integrations.eager_paged import eager_paged_attention_forward
@@ -111,7 +111,6 @@
     is_bitsandbytes_available,
     is_flash_attn_2_available,
     is_flash_attn_3_available,
-    is_fsdp_enabled,
     is_kernels_available,
     is_offline_mode,
     is_optimum_available,

src/transformers/utils/__init__.py

@@ -161,7 +161,6 @@
     is_flute_available,
     is_fp_quant_available,
     is_fsdp_available,
-    is_fsdp_enabled,
     is_ftfy_available,
     is_g2p_en_available,
     is_galore_torch_available,

src/transformers/utils/import_utils.py

@@ -35,7 +35,7 @@
 
 from packaging import version
 
-from . import logging, strtobool
+from . import logging
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -661,18 +661,6 @@ def is_torch_bf16_available():
     return is_torch_bf16_gpu_available()
 
 
-def is_fsdp_enabled():
-    if is_torch_available():
-        import torch
-
-        return (
-            torch.distributed.is_available()
-            and torch.distributed.is_initialized()
-            and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1
-            and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1
-        )
-
-
 @lru_cache
 def is_torch_fp16_available_on_device(device):
     if not is_torch_available():