__init__.py 1.78 KB
Newer Older
1
2
3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
Tim Dettmers's avatar
Tim Dettmers committed
4
# LICENSE file in the root directory of this source tree.
Max Ryabinin's avatar
Max Ryabinin committed
5

6

7
8
import sys

9
10
11
import torch

from . import _ops, research, utils
12
13
14
from .autograd._functions import (
    MatmulLtState,
    matmul,
Aarni Koskela's avatar
Aarni Koskela committed
15
    matmul_4bit,
16
)
17
18
from .backends.cpu import ops as cpu_ops
from .backends.default import ops as default_ops
19
from .nn import modules
20
from .optim import adam
Max Ryabinin's avatar
Max Ryabinin committed
21

22
# This is a signal for integrations with transformers/diffusers.
23
# Eventually we may remove this but it is currently required for compatibility.
24
features = {"multi_backend"}
25
26
supported_torch_devices = {
    "cpu",
27
28
29
30
31
    "cuda",  # NVIDIA/AMD GPU
    "xpu",  # Intel GPU
    "hpu",  # Gaudi
    "npu",  # Ascend NPU
    "mps",  # Apple Silicon
32
33
34
35
36
}

if torch.cuda.is_available():
    from .backends.cuda import ops as cuda_ops

37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

def _import_backends():
    """
    Discover and autoload all available backends installed as separate packages.
    Packages with an entrypoint for "bitsandbytes.backends" will be loaded.
    Inspired by PyTorch implementation: https://pytorch.org/tutorials/prototype/python_extension_autoload.html
    """
    from importlib.metadata import entry_points

    if sys.version_info < (3, 10):
        extensions = entry_points().get("bitsandbytes.backends", [])
    else:
        extensions = entry_points(group="bitsandbytes.backends")

    for ext in extensions:
        try:
            entry = ext.load()
            entry()
        except Exception as e:
            raise RuntimeError(f"bitsandbytes: failed to load backend {ext.name}: {e}") from e


_import_backends()

61
62
63
64
65
__pdoc__ = {
    "libbitsandbytes": False,
    "optim.optimizer.Optimizer8bit": False,
    "optim.optimizer.MockArgs": False,
}
66

Matthew Douglas's avatar
Matthew Douglas committed
67
__version__ = "0.46.0.dev0"