ayh015's picture
Upload folder using huggingface_hub
6011e08 verified
Raw
History Blame Contribute Delete
1.15 kB
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import logging
try:
_TORCH_MEMORY_SAVER_AVAILABLE = True
except ImportError:
logging.warning("torch_memory_saver is not installed, refer to : https://github.com/fzyzcjy/torch_memory_saver")
_TORCH_MEMORY_SAVER_AVAILABLE = False
try:
_FSDP_AVAILABLE = True
except ImportError as e:
logging.warning(f"FSDP backend dependencies not available: {e}")
_FSDP_AVAILABLE = False
if _FSDP_AVAILABLE:
from .actor import FSDPTrainRayActor
from .arguments import load_fsdp_args
else:
def _raise_import_error(*args, **kwargs):
raise ImportError(
"FSDP backend is not available. "
"Please ensure PyTorch with FSDP2 support is installed. "
"For installation instructions, refer to: https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html"
)
FSDPTrainRayActor = _raise_import_error
load_fsdp_args = _raise_import_error
__all__ = ["load_fsdp_args", "FSDPTrainRayActor"]
logging.getLogger().setLevel(logging.WARNING)