From b97ca44dbbf015a399bf9affa474d2dc0b7ef349 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 19 Aug 2024 17:12:32 -0700 Subject: [PATCH] Create _auto_install.py --- unsloth/_auto_install.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 unsloth/_auto_install.py diff --git a/unsloth/_auto_install.py b/unsloth/_auto_install.py new file mode 100644 index 0000000000..2f5b62d4ca --- /dev/null +++ b/unsloth/_auto_install.py @@ -0,0 +1,16 @@ +try: import torch +except: raise ImportError('Install torch via `pip install torch`') +from packaging.version import Version as V +v = V(torch.__version__) +cuda = str(torch.version.cuda) +is_ampere = torch.cuda.get_device_capability()[0] >= 8 +if cuda != "12.1" and cuda != "11.8": raise RuntimeError(f"CUDA = {cuda} not supported!") +if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!") +elif v <= V('2.1.1'): x = 'cu{}{}-torch211' +elif v <= V('2.1.2'): x = 'cu{}{}-torch212' +elif v < V('2.3.0'): x = 'cu{}{}-torch220' +elif v < V('2.4.0'): x = 'cu{}{}-torch230' +elif v < V('2.5.0'): x = 'cu{}{}-torch240' +else: raise RuntimeError(f"Torch = {v} too new!") +x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "") +print(f'pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"') \ No newline at end of file