From 8711310cdab6e729176ac78ca230ce1e6f1b45b7 Mon Sep 17 00:00:00 2001 From: Frank Lee Date: Thu, 5 Jan 2023 13:53:28 +0800 Subject: [PATCH] [setup] remove torch dependency (#2333) --- op_builder/builder.py | 8 ++++++-- setup.py | 5 +++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/op_builder/builder.py b/op_builder/builder.py index 52f1a9cf9..2e3728397 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -3,14 +3,18 @@ import re from pathlib import Path from typing import List -import torch - def get_cuda_cc_flag() -> List: """get_cuda_cc_flag cc flag for your GPU arch """ + + # only import torch when needed + # this is to avoid importing torch when building on a machine without torch pre-installed + # one case is to build wheel for pypi release + import torch + cc_flag = [] for arch in torch.cuda.get_arch_list(): res = re.search(r'sm_(\d+)', arch) diff --git a/setup.py b/setup.py index 453f6421d..dda1061d7 100644 --- a/setup.py +++ b/setup.py @@ -15,8 +15,9 @@ try: if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10): raise RuntimeError("Colossal-AI requires Pytorch 1.10 or newer.\n" "The latest stable release can be obtained from https://pytorch.org/") + TORCH_AVAILABLE = True except ImportError: - raise ModuleNotFoundError('torch is not found. You need to install PyTorch before installing Colossal-AI.') + TORCH_AVAILABLE = False # ninja build does not work unless include_dirs are abs path @@ -24,7 +25,7 @@ this_dir = os.path.dirname(os.path.abspath(__file__)) build_cuda_ext = True ext_modules = [] -if int(os.environ.get('NO_CUDA_EXT', '0')) == 1: +if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE: build_cuda_ext = False