Function get_accelerator

accelerator/real_accelerator.py:51–246

()

def get_accelerator():
    global ds_accelerator
    if ds_accelerator is not None:
        return ds_accelerator

    accelerator_name = None
    ds_set_method = None
    # 1. Detect whether there is override of DeepSpeed accelerators from environment variable.
    if "DS_ACCELERATOR" in os.environ.keys():
        accelerator_name = os.environ["DS_ACCELERATOR"]
        if accelerator_name == "xpu":
            try:
                import torch
                assert hasattr(torch, 'xpu') and torch.xpu.is_available(), \
                    "XPU_Accelerator requires PyTorch with XPU support (torch.xpu)."
            except (ImportError, AssertionError) as e:
                raise ValueError(f"XPU_Accelerator requires PyTorch with XPU support: {e}")
        elif accelerator_name == "cpu":
            pass
        elif accelerator_name == "npu":
            try:
                import torch_npu  # noqa: F401 # type: ignore
            except ImportError as e:
                raise ValueError("NPU_Accelerator requires torch_npu, which is not installed on this system.")
            pass
        elif accelerator_name == "sdaa":
            try:
                import torch_sdaa  # noqa: F401 # type: ignore
            except ImportError as e:
                raise ValueError("SDAA_Accelerator requires torch_sdaa, which is not installed on this system.")
            pass
        elif accelerator_name == "mps":
            try:
                import torch.mps

                # should use torch.mps.is_available() if it exists someday but this is used as proxy
                torch.mps.current_allocated_memory()
            except (RuntimeError, ImportError) as e:
                raise ValueError("MPS_Accelerator requires torch.mps, which is not installed on this system.")
        elif accelerator_name == "hpu":
            try:
                import habana_frameworks.torch.hpu  # noqa: F401
            except ImportError as e:
                raise ValueError(
                    "HPU_Accelerator requires habana_frameworks.torch.hpu, which is not installed on this system.")
        elif accelerator_name == "mlu":
            try:
                import torch_mlu  # noqa: F401
            except ImportError as e:
                raise ValueError("MLU_Accelerator requires torch_mlu, which is not installed on this system.")
        elif accelerator_name == "supa":
            try:
                import torch_supa  # noqa: F401 # type: ignore
            except ImportError as e:
                raise ValueError("SUPA_Accelerator requires torch_supa, which is not installed on this system.")
        elif accelerator_name not in SUPPORTED_ACCELERATOR_LIST:
            raise ValueError(f'DS_ACCELERATOR must be one of {SUPPORTED_ACCELERATOR_LIST}. '
                             f'Value "{accelerator_name}" is not supported')